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

package main

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"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"
)

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)
	}

	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)
		}
		if exists {
			logger.Info("migration skipped", "version", version)
			continue
		}

		sqlBytes, err := os.ReadFile(file)
		if err != nil {
			logger.Error("read migration file failed", "file", file, "error", err)
			os.Exit(1)
		}

		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 _, 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)
			}
			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 _, 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)
		}
		logger.Info("migration applied", "version", version)
	}

	fmt.Println("migrations complete")
}

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
}