terdut-server/internal/db/db.go

package db

import (
	"context"
	"database/sql"
	"embed"
	"fmt"
	"io/fs"
	"log"
	"sort"
	"strings"
	"time"

	_ "github.com/jackc/pgx/v5/stdlib"
)

//go:embed migrations
var migrationsFS embed.FS

// pingAttempts and pingRetryDelay bound the retry on the first connection.
// This pod's own IP can reach the Postgres pod's node before that node's
// NetworkPolicy enforcement (kube-router, reacting to the pod's creation
// event) has added it to the allowed-source set, which fails the ping with
// "connection refused" rather than a timeout. That race resolves within
// several seconds in practice; five attempts two seconds apart give it
// comfortable room without turning a genuinely absent database into a long
// hang.
const (
	pingAttempts   = 5
	pingRetryDelay = 2 * time.Second
)

// Open connects to Postgres. dsn is a libpq connection string or URL, e.g.
// postgres://terdut:secret@localhost:5432/terdut?sslmode=disable.
//
// The pool is modest on purpose: this server's concurrency comes from a handful
// of HTTP handlers plus two background loops, and a cloud-native-pg instance
// sized for it has a low max_connections.
func Open(dsn string) (*sql.DB, error) {
	if dsn == "" {
		return nil, fmt.Errorf("empty DSN: set TERDUT_DB_DSN")
	}
	db, err := sql.Open("pgx", dsn)
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(10)
	db.SetMaxIdleConns(5)
	db.SetConnMaxLifetime(time.Hour)

	for attempt := 1; ; attempt++ {
		err = db.Ping()
		if err == nil {
			return db, nil
		}
		if attempt == pingAttempts {
			db.Close()
			return nil, fmt.Errorf("ping: %w", err)
		}
		log.Printf("open db: ping attempt %d/%d failed, retrying in %s: %v", attempt, pingAttempts, pingRetryDelay, err)
		time.Sleep(pingRetryDelay)
	}
}

// migrationLockKey is the Postgres advisory lock Migrate holds for its whole
// run. Two replicas starting at once would otherwise race the check-then-apply
// loop below against schema_migrations: the loser could crash on a
// duplicate-key insert, or contend with the winner's uncommitted DDL. Blocking
// (pg_advisory_lock, not pg_try_advisory_lock as the archiver and notifier
// use): on boot there is no later tick to defer to, so the right behaviour is
// to wait for the other replica to finish migrating, not to skip ahead and
// start serving against an unmigrated schema.
const migrationLockKey int64 = 7265_0003

// Migrate applies every embedded migration that has not been applied yet, in
// filename order, recording each in schema_migrations.
//
// Each file runs inside a transaction, so a migration that fails half way
// leaves the schema as it was: Postgres has transactional DDL.
func Migrate(db *sql.DB) error {
	ctx := context.Background()
	conn, err := db.Conn(ctx)
	if err != nil {
		return fmt.Errorf("migrate: acquire connection: %w", err)
	}
	defer conn.Close()

	if _, err := conn.ExecContext(ctx, "SELECT pg_advisory_lock($1)", migrationLockKey); err != nil {
		return fmt.Errorf("migrate: acquire advisory lock: %w", err)
	}
	defer func() {
		if _, err := conn.ExecContext(ctx, "SELECT pg_advisory_unlock($1)", migrationLockKey); err != nil {
			log.Printf("migrate: release advisory lock: %v", err)
		}
	}()

	if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS schema_migrations (
		version    TEXT PRIMARY KEY,
		applied_at BIGINT NOT NULL DEFAULT FLOOR(EXTRACT(EPOCH FROM now()))::bigint
	)`); err != nil {
		return fmt.Errorf("create schema_migrations: %w", err)
	}

	entries, err := fs.ReadDir(migrationsFS, "migrations")
	if err != nil {
		return fmt.Errorf("read migrations dir: %w", err)
	}

	var files []string
	for _, e := range entries {
		if !e.IsDir() && strings.HasSuffix(e.Name(), ".sql") {
			files = append(files, e.Name())
		}
	}
	sort.Strings(files)

	for _, name := range files {
		var count int
		if err := db.QueryRow("SELECT COUNT(*) FROM schema_migrations WHERE version = $1", name).Scan(&count); err != nil {
			return fmt.Errorf("check migration %s: %w", name, err)
		}
		if count > 0 {
			continue
		}

		data, err := migrationsFS.ReadFile("migrations/" + name)
		if err != nil {
			return fmt.Errorf("read migration %s: %w", name, err)
		}

		if err := applyMigration(db, name, string(data)); err != nil {
			return err
		}
	}
	return nil
}

func applyMigration(db *sql.DB, name, body string) error {
	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("begin migration %s: %w", name, err)
	}
	defer tx.Rollback()

	if _, err := tx.Exec(body); err != nil {
		return fmt.Errorf("apply migration %s: %w", name, err)
	}
	if _, err := tx.Exec("INSERT INTO schema_migrations (version) VALUES ($1)", name); err != nil {
		return fmt.Errorf("record migration %s: %w", name, err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit migration %s: %w", name, err)
	}
	return nil
}