273 lines
8.6 KiB
Go
273 lines
8.6 KiB
Go
// Package migrations applies the canonical SQL migrations before the API starts.
|
|
// SQL in sql/ is generated from database/migrations by scripts/sync-go-migrations.mjs.
|
|
package migrations
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"embed"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"io/fs"
|
|
"log"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/postgres"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
)
|
|
|
|
//go:embed sql/*.sql
|
|
var sqlFiles embed.FS
|
|
|
|
const migrationLockID int64 = 7308731946202608
|
|
|
|
type migration struct {
|
|
version string
|
|
checksum string
|
|
sql string
|
|
}
|
|
|
|
// Run is a no-op unless DATABASE_AUTO_MIGRATE=true. Its context covers
|
|
// connection, advisory-lock waiting, every migration, and role provisioning.
|
|
func Run(ctx context.Context, getenv func(string) string) error {
|
|
if getenv == nil {
|
|
return errors.New("environment getter is required")
|
|
}
|
|
switch strings.ToLower(strings.TrimSpace(getenv("DATABASE_AUTO_MIGRATE"))) {
|
|
case "", "false":
|
|
return nil
|
|
case "true":
|
|
default:
|
|
return errors.New("DATABASE_AUTO_MIGRATE must be true or false")
|
|
}
|
|
if strings.ToLower(strings.TrimSpace(getenv("ZHINIAN_DATA_BACKEND"))) != "postgres" {
|
|
return errors.New("DATABASE_AUTO_MIGRATE requires ZHINIAN_DATA_BACKEND=postgres")
|
|
}
|
|
timeout := 180 * time.Second
|
|
if raw := strings.TrimSpace(getenv("DATABASE_MIGRATION_TIMEOUT_SECONDS")); raw != "" {
|
|
seconds, err := strconv.ParseInt(raw, 10, 64)
|
|
if err != nil || seconds < 1 || seconds > 86400 {
|
|
return errors.New("DATABASE_MIGRATION_TIMEOUT_SECONDS must be an integer from 1 to 86400")
|
|
}
|
|
timeout = time.Duration(seconds) * time.Second
|
|
}
|
|
role := strings.TrimSpace(getenv("DATABASE_APP_ROLE"))
|
|
if role == "" && strings.EqualFold(strings.TrimSpace(getenv("NODE_ENV")), "production") {
|
|
return errors.New("DATABASE_APP_ROLE is required in production")
|
|
}
|
|
if role != "" && !rolePattern.MatchString(role) {
|
|
return errors.New("DATABASE_APP_ROLE must match [a-z_][a-z0-9_]{0,62}")
|
|
}
|
|
all, err := discover(sqlFiles)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
migrationURL := strings.TrimSpace(getenv("DATABASE_MIGRATION_URL"))
|
|
appURL := strings.TrimSpace(getenv("DATABASE_URL"))
|
|
if migrationURL == "" {
|
|
migrationURL = appURL
|
|
}
|
|
if migrationURL == "" {
|
|
return errors.New("DATABASE_MIGRATION_URL or DATABASE_URL is required")
|
|
}
|
|
// Reuse the application's URI validation and forced transport policy.
|
|
config, err := postgres.ParseConfig(func(key string) string {
|
|
if key == "DATABASE_URL" {
|
|
return migrationURL
|
|
}
|
|
return getenv(key)
|
|
}, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("migration database configuration: %w", err)
|
|
}
|
|
pgconfig, err := pgx.ParseConfig(config.DatabaseURL)
|
|
if err != nil {
|
|
return errors.New("invalid migration database connection configuration")
|
|
}
|
|
if appURL != "" && strings.TrimSpace(getenv("DATABASE_MIGRATION_URL")) != "" {
|
|
appConfig, err := postgres.ParseConfig(func(key string) string {
|
|
if key == "DATABASE_URL" {
|
|
return appURL
|
|
}
|
|
return getenv(key)
|
|
}, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("application database configuration: %w", err)
|
|
}
|
|
appPGConfig, err := pgx.ParseConfig(appConfig.DatabaseURL)
|
|
if err != nil {
|
|
return errors.New("invalid application database connection configuration")
|
|
}
|
|
if !strings.EqualFold(pgconfig.Host, appPGConfig.Host) || pgconfig.Port != appPGConfig.Port || pgconfig.Database != appPGConfig.Database {
|
|
return errors.New("DATABASE_MIGRATION_URL must target the same host, port, and database as DATABASE_URL")
|
|
}
|
|
}
|
|
pgconfig.RuntimeParams["application_name"] = "zhinian-migrate"
|
|
pgconfig.ConnectTimeout = config.ConnectionTimeout
|
|
ctx, cancel := context.WithTimeout(ctx, timeout)
|
|
defer cancel()
|
|
log.Printf("zhinian-api checking database migrations (%d bundled versions)", len(all))
|
|
conn, err := pgx.ConnectConfig(ctx, pgconfig)
|
|
if err != nil {
|
|
return safeDBError("connect", "", err)
|
|
}
|
|
defer conn.Close(context.Background())
|
|
log.Print("zhinian-api waiting for database migration lock")
|
|
if _, err := conn.Exec(ctx, "SELECT pg_advisory_lock($1)", migrationLockID); err != nil {
|
|
return safeDBError("acquire migration lock", "", err)
|
|
}
|
|
// Session-level lock is released when conn closes, including on timeout.
|
|
if err := apply(ctx, conn, all, role); err != nil {
|
|
return err
|
|
}
|
|
log.Printf("zhinian-api database migrations ready (%d versions)", len(all))
|
|
return nil
|
|
}
|
|
|
|
func apply(ctx context.Context, conn *pgx.Conn, all []migration, role string) error {
|
|
if _, err := conn.Exec(ctx, `CREATE TABLE IF NOT EXISTS schema_migrations (
|
|
version text PRIMARY KEY,
|
|
checksum text NOT NULL,
|
|
applied_at timestamptz NOT NULL DEFAULT now()
|
|
)`); err != nil {
|
|
return safeDBError("create migration tracking table", "", err)
|
|
}
|
|
rows, err := conn.Query(ctx, "SELECT version, checksum FROM schema_migrations")
|
|
if err != nil {
|
|
return safeDBError("read migration history", "", err)
|
|
}
|
|
applied := make(map[string]string)
|
|
for rows.Next() {
|
|
var version, checksum string
|
|
if err := rows.Scan(&version, &checksum); err != nil {
|
|
rows.Close()
|
|
return safeDBError("read migration history", "", err)
|
|
}
|
|
applied[version] = checksum
|
|
}
|
|
err = rows.Err()
|
|
rows.Close()
|
|
if err != nil {
|
|
return safeDBError("read migration history", "", err)
|
|
}
|
|
known := make(map[string]string, len(all))
|
|
for _, item := range all {
|
|
known[item.version] = item.checksum
|
|
}
|
|
for version, checksum := range applied {
|
|
expected, exists := known[version]
|
|
if !exists {
|
|
return fmt.Errorf("applied migration %q is absent from bundle", version)
|
|
}
|
|
if checksum != expected {
|
|
return fmt.Errorf("applied migration %q checksum differs from bundle", version)
|
|
}
|
|
}
|
|
for _, item := range all {
|
|
if _, exists := applied[item.version]; exists {
|
|
continue
|
|
}
|
|
body, err := withoutOuterTransaction(item.sql)
|
|
if err != nil {
|
|
return fmt.Errorf("migration %s: %w", item.version, err)
|
|
}
|
|
tx, err := conn.Begin(ctx)
|
|
if err != nil {
|
|
return safeDBError("begin migration", item.version, err)
|
|
}
|
|
// PostgreSQL's simple protocol permits multi-statement migration files.
|
|
_, err = tx.Exec(ctx, body, pgx.QueryExecModeSimpleProtocol)
|
|
if err == nil {
|
|
_, err = tx.Exec(ctx, "INSERT INTO schema_migrations(version, checksum) VALUES ($1, $2)", item.version, item.checksum)
|
|
}
|
|
if err == nil {
|
|
err = tx.Commit(ctx)
|
|
}
|
|
if err != nil {
|
|
_ = tx.Rollback(context.Background())
|
|
return safeDBError("apply migration", item.version, err)
|
|
}
|
|
log.Printf("zhinian-api applied database migration %s", item.version)
|
|
}
|
|
if role != "" {
|
|
if err := provisionRole(ctx, conn, role); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func discover(fsys fs.FS) ([]migration, error) {
|
|
entries, err := fs.ReadDir(fsys, "sql")
|
|
if err != nil {
|
|
return nil, errors.New("read embedded migration bundle")
|
|
}
|
|
var all []migration
|
|
for _, entry := range entries {
|
|
name := entry.Name()
|
|
if entry.IsDir() || !strings.HasSuffix(name, ".sql") {
|
|
continue
|
|
}
|
|
data, err := fs.ReadFile(fsys, "sql/"+name)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read embedded migration %s", name)
|
|
}
|
|
checksum := sha256.Sum256(data)
|
|
all = append(all, migration{strings.TrimSuffix(name, ".sql"), hex.EncodeToString(checksum[:]), string(data)})
|
|
}
|
|
sort.Slice(all, func(i, j int) bool { return all[i].version < all[j].version })
|
|
if len(all) == 0 {
|
|
return nil, errors.New("embedded migration bundle is empty")
|
|
}
|
|
return all, nil
|
|
}
|
|
|
|
// The canonical SQL may contain an explicit outer BEGIN/COMMIT. Strip only
|
|
// when both commands are the first and last non-comment statements; otherwise
|
|
// our transaction would commit before recording the checksum.
|
|
func withoutOuterTransaction(sql string) (string, error) {
|
|
first, last := -1, -1
|
|
lines := strings.SplitAfter(sql, "\n")
|
|
for i, line := range lines {
|
|
trimmed := strings.TrimSpace(line)
|
|
if trimmed != "" && !strings.HasPrefix(trimmed, "--") {
|
|
if first < 0 {
|
|
first = i
|
|
}
|
|
last = i
|
|
}
|
|
}
|
|
if first < 0 {
|
|
return "", errors.New("empty migration SQL")
|
|
}
|
|
begin := strings.EqualFold(strings.TrimSpace(lines[first]), "begin;")
|
|
commit := strings.EqualFold(strings.TrimSpace(lines[last]), "commit;")
|
|
if begin != commit {
|
|
return "", errors.New("outer transaction must contain both BEGIN and COMMIT")
|
|
}
|
|
if !begin {
|
|
return sql, nil
|
|
}
|
|
lines[first], lines[last] = "", ""
|
|
return strings.Join(lines, ""), nil
|
|
}
|
|
|
|
func safeDBError(operation, version string, err error) error {
|
|
if version != "" {
|
|
operation += " " + version
|
|
}
|
|
var pgErr *pgconn.PgError
|
|
if errors.As(err, &pgErr) {
|
|
return fmt.Errorf("%s failed (SQLSTATE %s)", operation, pgErr.Code)
|
|
}
|
|
if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
|
|
return fmt.Errorf("%s timed out or was canceled", operation)
|
|
}
|
|
return fmt.Errorf("%s failed (%T)", operation, err)
|
|
}
|