easyai-ai-gateway/apps/api/cmd/migrate/main.go

package main

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

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

const (
	noTransactionMigrationMarker      = "-- easyai:migration:no-transaction"
	migrationStatementSeparator       = "-- easyai:migration:statement"
	acceptanceImportReplayMarker      = "-- easyai:migration:reapply-after-acceptance-import"
	acceptanceImportReplayEnvironment = "AI_GATEWAY_MIGRATION_REAPPLY_ACCEPTANCE_IMPORT"
	acceptanceLocalClusterSettingKey  = "acceptance_local_cluster_id"
)

var acceptanceLocalClusterMarkerPattern = regexp.MustCompile(`^easyai-local-[0-9a-f]{24}$`)

func main() {
	cfg := config.Load()
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	ctx := context.Background()

	conn, err := pgx.Connect(ctx, cfg.DatabaseURL)
	if err != nil {
		logger.Error("connect postgres failed", "error", err)
		os.Exit(1)
	}
	defer conn.Close(ctx)
	if _, err := conn.Exec(ctx, "SET standard_conforming_strings = on"); err != nil {
		logger.Error("enforce standard SQL string semantics failed", "error", err)
		os.Exit(1)
	}

	if _, err := conn.Exec(ctx, `
CREATE TABLE IF NOT EXISTS schema_migrations (
  version text PRIMARY KEY,
  applied_at timestamptz NOT NULL DEFAULT now()
);`); err != nil {
		logger.Error("ensure schema_migrations failed", "error", err)
		os.Exit(1)
	}
	reapplyAcceptanceImport := strings.EqualFold(strings.TrimSpace(os.Getenv(acceptanceImportReplayEnvironment)), "true")
	if reapplyAcceptanceImport {
		var localClusterID string
		if err := conn.QueryRow(ctx, `
SELECT COALESCE(value->>'clusterId', '')
FROM system_settings
WHERE setting_key = $1`, acceptanceLocalClusterSettingKey).Scan(&localClusterID); err != nil {
			logger.Error("verify local acceptance database failed", "error", err)
			os.Exit(1)
		}
		if !acceptanceLocalClusterMarkerPattern.MatchString(localClusterID) {
			logger.Error("refusing acceptance migration replay outside a marked local cluster")
			os.Exit(1)
		}
	}

	files, err := filepath.Glob("migrations/*.sql")
	if err != nil {
		logger.Error("read migrations failed", "error", err)
		os.Exit(1)
	}
	sort.Strings(files)

	for _, file := range files {
		version := strings.TrimSuffix(filepath.Base(file), filepath.Ext(file))
		var exists bool
		if err := conn.QueryRow(ctx, "SELECT EXISTS (SELECT 1 FROM schema_migrations WHERE version=$1)", version).Scan(&exists); err != nil {
			logger.Error("check migration failed", "version", version, "error", err)
			os.Exit(1)
		}
		sqlBytes, err := os.ReadFile(file)
		if err != nil {
			logger.Error("read migration file failed", "file", file, "error", err)
			os.Exit(1)
		}
		replaying := exists && reapplyAcceptanceImport && hasMigrationMarker(string(sqlBytes), acceptanceImportReplayMarker)
		if exists && !replaying {
			logger.Info("migration skipped", "version", version)
			continue
		}

		noTransaction, statements := migrationStatements(string(sqlBytes))
		if noTransaction {
			for _, statement := range statements {
				if _, err := conn.Exec(ctx, statement); err != nil {
					logger.Error("execute non-transaction migration failed", "version", version, "error", err)
					os.Exit(1)
				}
			}
			if !exists {
				if _, err := conn.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
					logger.Error("record non-transaction migration failed", "version", version, "error", err)
					os.Exit(1)
				}
			}
			if replaying {
				logger.Info("acceptance migration replayed", "version", version, "transactional", false)
			} else {
				logger.Info("migration applied", "version", version, "transactional", false)
			}
			continue
		}

		tx, err := conn.Begin(ctx)
		if err != nil {
			logger.Error("begin migration failed", "version", version, "error", err)
			os.Exit(1)
		}
		// Pin string parsing semantics inside every migration transaction. A
		// previous migration may have changed the session GUC, while each file is
		// parsed and executed independently.
		if _, err := tx.Exec(ctx, "SET LOCAL standard_conforming_strings = on"); err != nil {
			_ = tx.Rollback(ctx)
			logger.Error("enforce migration SQL string semantics failed", "version", version, "error", err)
			os.Exit(1)
		}
		if _, err := tx.Exec(ctx, string(sqlBytes)); err != nil {
			_ = tx.Rollback(ctx)
			logger.Error("execute migration failed", "version", version, "error", err)
			os.Exit(1)
		}
		if !exists {
			if _, err := tx.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
				_ = tx.Rollback(ctx)
				logger.Error("record migration failed", "version", version, "error", err)
				os.Exit(1)
			}
		}
		if err := tx.Commit(ctx); err != nil {
			logger.Error("commit migration failed", "version", version, "error", err)
			os.Exit(1)
		}
		if replaying {
			logger.Info("acceptance migration replayed", "version", version)
		} else {
			logger.Info("migration applied", "version", version)
		}
	}

	fmt.Println("migrations complete")
}

func hasMigrationMarker(sql string, marker string) bool {
	for _, line := range strings.Split(sql, "\n") {
		if strings.TrimSpace(line) == marker {
			return true
		}
	}
	return false
}

func migrationStatements(sql string) (bool, []string) {
	trimmed := strings.TrimSpace(sql)
	if !strings.HasPrefix(trimmed, noTransactionMigrationMarker) {
		return false, []string{sql}
	}
	trimmed = strings.TrimSpace(strings.TrimPrefix(trimmed, noTransactionMigrationMarker))
	parts := strings.Split(trimmed, migrationStatementSeparator)
	statements := make([]string, 0, len(parts))
	for _, part := range parts {
		if statement := strings.TrimSpace(part); statement != "" {
			statements = append(statements, statement)
		}
	}
	return true, statements
}