easyai-ai-gateway/apps/api/cmd/acceptance-snapshot/main.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptancesnapshot"
	"github.com/jackc/pgx/v5/pgxpool"
)

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(64)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	var err error
	switch os.Args[1] {
	case "export":
		err = runExport(ctx, os.Args[2:])
	case "validate":
		err = runValidate(os.Args[2:])
	case "import":
		err = runImport(ctx, os.Args[2:])
	default:
		usage()
		os.Exit(64)
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, "acceptance snapshot:", err)
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprintln(os.Stderr, `Usage:
  easyai-ai-gateway-acceptance-snapshot export --release-sha <SHA> --output <file>
  easyai-ai-gateway-acceptance-snapshot validate --input <file>
  easyai-ai-gateway-acceptance-snapshot import --input <file> --local-cluster-id <id>

export is read-only. import refuses to write unless system_settings contains a
matching acceptance_local_cluster_id marker.`)
}

func runExport(ctx context.Context, args []string) error {
	flags := flag.NewFlagSet("export", flag.ContinueOnError)
	releaseSHA := flags.String("release-sha", "", "full source release SHA")
	output := flags.String("output", "", "secret-safe snapshot output path")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 0 {
		return errors.New("unexpected export arguments")
	}
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
	if databaseURL == "" {
		return errors.New("AI_GATEWAY_DATABASE_URL is required")
	}
	outputPath, err := safeOutputPath(*output)
	if err != nil {
		return err
	}
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		return err
	}
	defer pool.Close()
	snapshot, err := acceptancesnapshot.Export(ctx, pool, acceptancesnapshot.ExportOptions{
		ReleaseSHA: *releaseSHA,
	})
	if err != nil {
		return err
	}
	payload, err := acceptancesnapshot.Encode(snapshot)
	if err != nil {
		return err
	}
	if err := os.WriteFile(outputPath, payload, 0o600); err != nil {
		return err
	}
	fmt.Printf(
		"acceptance_snapshot_export=PASS schema=%s config_hash=%s snapshot_sha256=%s\n",
		snapshot.SchemaVersion,
		snapshot.Source.ConfigHash,
		snapshot.SnapshotSHA256,
	)
	return nil
}

func runValidate(args []string) error {
	flags := flag.NewFlagSet("validate", flag.ContinueOnError)
	input := flags.String("input", "", "snapshot input path")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 0 {
		return errors.New("unexpected validate arguments")
	}
	snapshot, err := readSnapshot(*input)
	if err != nil {
		return err
	}
	fmt.Printf(
		"acceptance_snapshot_validate=PASS schema=%s release=%s config_hash=%s snapshot_sha256=%s\n",
		snapshot.SchemaVersion,
		snapshot.Source.ReleaseSHA,
		snapshot.Source.ConfigHash,
		snapshot.SnapshotSHA256,
	)
	return nil
}

func runImport(ctx context.Context, args []string) error {
	flags := flag.NewFlagSet("import", flag.ContinueOnError)
	input := flags.String("input", "", "snapshot input path")
	clusterID := flags.String("local-cluster-id", "", "required local cluster marker")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 0 {
		return errors.New("unexpected import arguments")
	}
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
	if databaseURL == "" {
		return errors.New("AI_GATEWAY_DATABASE_URL is required")
	}
	snapshot, err := readSnapshot(*input)
	if err != nil {
		return err
	}
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		return err
	}
	defer pool.Close()
	if err := acceptancesnapshot.Import(ctx, pool, snapshot, *clusterID); err != nil {
		return err
	}
	fmt.Printf(
		"acceptance_snapshot_import=PASS cluster_id=%s config_hash=%s snapshot_sha256=%s\n",
		strings.TrimSpace(*clusterID),
		snapshot.Source.ConfigHash,
		snapshot.SnapshotSHA256,
	)
	return nil
}

func readSnapshot(path string) (acceptancesnapshot.Snapshot, error) {
	path = strings.TrimSpace(path)
	if path == "" {
		return acceptancesnapshot.Snapshot{}, errors.New("snapshot path is required")
	}
	info, err := os.Lstat(path)
	if err != nil {
		return acceptancesnapshot.Snapshot{}, err
	}
	if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
		return acceptancesnapshot.Snapshot{}, errors.New("snapshot must be a regular non-symlink file")
	}
	payload, err := os.ReadFile(path)
	if err != nil {
		return acceptancesnapshot.Snapshot{}, err
	}
	return acceptancesnapshot.Decode(payload)
}

func safeOutputPath(path string) (string, error) {
	path = strings.TrimSpace(path)
	if path == "" {
		return "", errors.New("output path is required")
	}
	absolute, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	parent := filepath.Dir(absolute)
	if err := os.MkdirAll(parent, 0o700); err != nil {
		return "", err
	}
	if info, err := os.Lstat(absolute); err == nil {
		if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
			return "", errors.New("output must be a regular non-symlink file")
		}
	} else if !os.IsNotExist(err) {
		return "", err
	}
	return absolute, nil
}