ai-gateway-go/cmd/gateway-legacy-import/main.go

package main

import (
	"bufio"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"

	"aigateway.local/core/internal/platform/config"
	"aigateway.local/core/internal/platform/database"
	platformid "aigateway.local/core/internal/platform/id"
	"aigateway.local/core/internal/platform/legacyid"
	"github.com/jackc/pgx/v5/pgxpool"
)

type record struct {
	SourceSystem string          `json:"source_system"`
	EntityType   string          `json:"entity_type"`
	LegacyID     string          `json:"legacy_id"`
	Data         json.RawMessage `json:"data"`
	Checksum     string          `json:"checksum"`
}
type staged struct {
	record
	NewID string
}

func main() {
	input := flag.String("input", "-", "JSONL file from export_legacy_data.py, or - for stdin")
	dryRun := flag.Bool("dry-run", false, "validate without database writes")
	flag.Parse()
	reader, closeInput, err := openInput(*input)
	if err != nil {
		fatal(err)
	}
	defer closeInput()
	records, sourceChecksum, counts, err := readRecords(reader)
	if err != nil {
		fatal(err)
	}
	summary := map[string]any{"records": len(records), "entities": counts, "source_checksum": sourceChecksum, "dry_run": *dryRun}
	if *dryRun {
		write(summary)
		return
	}
	cfg, err := config.Load()
	if err != nil {
		fatal(err)
	}
	if cfg.Database.URL == "" {
		fatal(errors.New("DATABASE_URL is required"))
	}
	ctx := context.Background()
	pool, err := database.Open(ctx, cfg.Database)
	if err != nil {
		fatal(err)
	}
	defer pool.Close()
	batchID, err := stage(ctx, pool, records, sourceChecksum, counts)
	if err != nil {
		fatal(err)
	}
	summary["batch_id"] = batchID
	summary["status"] = "staged"
	write(summary)
}

func openInput(path string) (io.Reader, func(), error) {
	if path == "-" {
		return os.Stdin, func() {}, nil
	}
	file, err := os.Open(path)
	if err != nil {
		return nil, func() {}, err
	}
	return file, func() { _ = file.Close() }, nil
}
func readRecords(reader io.Reader) ([]staged, string, map[string]int, error) {
	scanner := bufio.NewScanner(reader)
	scanner.Buffer(make([]byte, 64<<10), 32<<20)
	aggregate := sha256.New()
	items := []staged{}
	counts := map[string]int{}
	source := ""
	line := 0
	for scanner.Scan() {
		line++
		raw := bytesTrimSpace(scanner.Bytes())
		if len(raw) == 0 {
			continue
		}
		_, _ = aggregate.Write(raw)
		_, _ = aggregate.Write([]byte{'\n'})
		var item record
		if err := json.Unmarshal(raw, &item); err != nil {
			return nil, "", nil, fmt.Errorf("line %d: %w", line, err)
		}
		item.SourceSystem = strings.TrimSpace(item.SourceSystem)
		item.EntityType = strings.TrimSpace(item.EntityType)
		item.LegacyID = strings.TrimSpace(item.LegacyID)
		if source == "" {
			source = item.SourceSystem
		}
		if item.SourceSystem != source {
			return nil, "", nil, fmt.Errorf("line %d: mixed source systems", line)
		}
		canonical, _ := json.Marshal(map[string]any{"source_system": item.SourceSystem, "entity_type": item.EntityType, "legacy_id": item.LegacyID, "data": item.Data})
		digest := sha256.Sum256(canonical)
		if hex.EncodeToString(digest[:]) != item.Checksum {
			return nil, "", nil, fmt.Errorf("line %d: checksum mismatch", line)
		}
		newID, err := legacyid.UUID(item.SourceSystem, item.EntityType, item.LegacyID)
		if err != nil {
			return nil, "", nil, fmt.Errorf("line %d: %w", line, err)
		}
		items = append(items, staged{record: item, NewID: newID})
		counts[item.EntityType]++
	}
	if err := scanner.Err(); err != nil {
		return nil, "", nil, err
	}
	if len(items) == 0 {
		return nil, "", nil, errors.New("input contains no records")
	}
	return items, hex.EncodeToString(aggregate.Sum(nil)), counts, nil
}

func stage(ctx context.Context, pool *pgxpool.Pool, items []staged, sourceChecksum string, counts map[string]int) (string, error) {
	batchID, err := platformid.NewUUID()
	if err != nil {
		return "", err
	}
	tx, err := pool.Begin(ctx)
	if err != nil {
		return "", err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	countsJSON, _ := json.Marshal(counts)
	source := items[0].SourceSystem
	err = tx.QueryRow(ctx, `INSERT INTO gateway.legacy_import_batches(id,source_system,source_checksum,status,record_count,entity_counts) VALUES($1,$2,$3,'staged',$4,$5) ON CONFLICT(source_system,source_checksum) DO UPDATE SET source_checksum=excluded.source_checksum RETURNING id::text`, batchID, source, sourceChecksum, len(items), countsJSON).Scan(&batchID)
	if err != nil {
		return "", err
	}
	for _, item := range items {
		tag, insertErr := tx.Exec(ctx, `INSERT INTO gateway.legacy_import_records(first_seen_batch_id,source_system,entity_type,legacy_id,new_id,payload,payload_checksum) VALUES($1,$2,$3,$4,$5,$6,$7) ON CONFLICT(source_system,entity_type,legacy_id) DO UPDATE SET payload_checksum=gateway.legacy_import_records.payload_checksum WHERE gateway.legacy_import_records.payload_checksum=excluded.payload_checksum`, batchID, item.SourceSystem, item.EntityType, item.LegacyID, item.NewID, item.Data, item.Checksum)
		if insertErr != nil {
			return "", insertErr
		}
		if tag.RowsAffected() == 0 {
			return "", fmt.Errorf("legacy record changed since an earlier import: %s/%s", item.EntityType, item.LegacyID)
		}
		if _, insertErr = tx.Exec(ctx, `INSERT INTO gateway.legacy_import_batch_records(batch_id,source_system,entity_type,legacy_id,payload_checksum) VALUES($1,$2,$3,$4,$5) ON CONFLICT DO NOTHING`, batchID, item.SourceSystem, item.EntityType, item.LegacyID, item.Checksum); insertErr != nil {
			return "", insertErr
		}
		metadata, _ := json.Marshal(map[string]any{"batch_id": batchID, "payload_checksum": item.Checksum})
		if _, insertErr = tx.Exec(ctx, `INSERT INTO gateway.legacy_id_mappings(source_system,entity_type,legacy_id,new_id,metadata) VALUES($1,$2,$3,$4,$5) ON CONFLICT(source_system,entity_type,legacy_id) DO UPDATE SET metadata=excluded.metadata WHERE gateway.legacy_id_mappings.new_id=excluded.new_id`, item.SourceSystem, item.EntityType, item.LegacyID, item.NewID, metadata); insertErr != nil {
			return "", insertErr
		}
	}
	if err = tx.Commit(ctx); err != nil {
		return "", err
	}
	return batchID, nil
}

func bytesTrimSpace(value []byte) []byte { return []byte(strings.TrimSpace(string(value))) }
func write(value any) {
	encoded, _ := json.MarshalIndent(value, "", "  ")
	fmt.Println(string(encoded))
}
func fatal(err error) { fmt.Fprintln(os.Stderr, err); os.Exit(1) }