package db import ( "context" "crypto/sha256" "encoding/hex" "fmt" "io/fs" "log/slog" "sort" "strings" "github.com/dtoro/oikos/migrations" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "gopkg.in/yaml.v3" ) // Pool wraps a pgx connection pool. type Pool struct { *pgxpool.Pool } // New creates a new connection pool. func New(ctx context.Context, databaseURL string) (*Pool, error) { cfg, err := pgxpool.ParseConfig(databaseURL) if err != nil { return nil, fmt.Errorf("parse database url: %w", err) } cfg.MaxConns = 15 pool, err := pgxpool.NewWithConfig(ctx, cfg) if err != nil { return nil, fmt.Errorf("create pool: %w", err) } if err := pool.Ping(ctx); err != nil { return nil, fmt.Errorf("ping db: %w", err) } return &Pool{pool}, nil } // migrationLockKey is the advisory-lock key serializing migration runs — // two concurrent `oikos migrate` invocations must not interleave DDL. const migrationLockKey = 0x01c05e5 // Migrate runs all embedded forward migrations in order. // Uses a schema_migrations table to track applied versions. The whole run // happens on one connection holding a session advisory lock. func (p *Pool) Migrate(ctx context.Context) error { conn, err := p.Acquire(ctx) if err != nil { return fmt.Errorf("acquire migration conn: %w", err) } defer conn.Release() if _, err := conn.Exec(ctx, "SELECT pg_advisory_lock($1)", migrationLockKey); err != nil { return fmt.Errorf("acquire migration lock: %w", err) } defer conn.Exec(context.WithoutCancel(ctx), "SELECT pg_advisory_unlock($1)", migrationLockKey) // Create tracking table if not exists _, err = conn.Exec(ctx, ` CREATE TABLE IF NOT EXISTS schema_migrations ( version INT PRIMARY KEY, applied_at TIMESTAMPTZ NOT NULL DEFAULT now() ) `) if err != nil { return fmt.Errorf("create schema_migrations: %w", err) } // List migration files entries, err := fs.ReadDir(migrations.FS, ".") if err != nil { return fmt.Errorf("read migration fs: %w", err) } var files []string for _, e := range entries { if !e.IsDir() && hasSuffix(e.Name(), ".up.sql") { files = append(files, e.Name()) } } sort.Strings(files) for _, fname := range files { // Extract version number (001, 002, etc.) var version int if _, err := fmt.Sscanf(fname, "%03d", &version); err != nil { continue } // Check if already applied var applied int err := conn.QueryRow(ctx, "SELECT COUNT(*) FROM schema_migrations WHERE version = $1", version).Scan(&applied) if err != nil { return fmt.Errorf("check migration %d: %w", version, err) } if applied > 0 { continue } // Read and execute migration — split into individual statements // because TimescaleDB CAGGs and some DDL can't run inside a transaction, // and pgx's multi-statement Exec wraps them implicitly. content, err := migrations.FS.ReadFile(fname) if err != nil { return fmt.Errorf("read %s: %w", fname, err) } stmts := splitSQL(string(content)) for i, stmt := range stmts { stmt = strings.TrimSpace(stmt) if stmt == "" { continue } _, err := conn.Exec(ctx, stmt) if err != nil { return fmt.Errorf("exec migration %s stmt %d: %w", fname, i+1, err) } } _, err = conn.Exec(ctx, "INSERT INTO schema_migrations (version) VALUES ($1)", version) if err != nil { return fmt.Errorf("record migration %d: %w", version, err) } slog.Info("migration applied", "file", fname, "version", version, "statements", len(stmts)) } return nil } // SeedIngest ingests a YAML seed file into the database. // Idempotent: if the file's content hash matches seed_versions, it's a no-op (A4). func (p *Pool) SeedIngest(ctx context.Context, filename string, content []byte, ingestFn func(ctx context.Context, tx pgx.Tx, data map[string]any) error) error { hash := contentHash(content) // Check if already applied with same hash var existing string err := p.QueryRow(ctx, "SELECT content_hash FROM seed_versions WHERE file = $1", filename).Scan(&existing) if err == nil && existing == hash { return nil // no-op, same content } // Parse YAML var data map[string]any if err := yaml.Unmarshal(content, &data); err != nil { return fmt.Errorf("parse %s: %w", filename, err) } // Apply in a single transaction tx, err := p.Begin(ctx) if err != nil { return fmt.Errorf("begin tx: %w", err) } defer tx.Rollback(ctx) if err := ingestFn(ctx, tx, data); err != nil { return fmt.Errorf("ingest %s: %w", filename, err) } // Record the seed version _, err = tx.Exec(ctx, `INSERT INTO seed_versions (file, content_hash) VALUES ($1, $2) ON CONFLICT (file) DO UPDATE SET content_hash = $2, applied_at = now()`, filename, hash) if err != nil { return fmt.Errorf("record seed version: %w", err) } if err := tx.Commit(ctx); err != nil { return fmt.Errorf("commit seed: %w", err) } return nil } // contentHash returns a SHA-256 hex digest of the content. func contentHash(content []byte) string { h := sha256.Sum256(content) return hex.EncodeToString(h[:]) } // hasSuffix reports whether the string ends with the given suffix. func hasSuffix(s, suffix string) bool { return strings.HasSuffix(s, suffix) } // splitSQL splits a SQL string into individual statements. // Handles $$ ... $$ dollar-quoted blocks, $tag$ ... $tag$ tagged quotes, // -- line comments, /* ... */ block comments, and '...' string literals // so that semicolons inside any of these constructs are not treated as // statement boundaries. func splitSQL(sql string) []string { var statements []string var current strings.Builder inDollarQuote := false dollarTag := "" i := 0 for i < len(sql) { // Handle line comments (-- to end of line) if !inDollarQuote && i+1 < len(sql) && sql[i] == '-' && sql[i+1] == '-' { for i < len(sql) && sql[i] != '\n' { current.WriteByte(sql[i]) i++ } continue } // Handle block comments (/* ... */) if !inDollarQuote && i+1 < len(sql) && sql[i] == '/' && sql[i+1] == '*' { end := strings.Index(sql[i+2:], "*/") if end >= 0 { current.WriteString(sql[i : i+end+4]) i += end + 4 continue } } // Handle single-quoted string literals ('...') if !inDollarQuote && sql[i] == '\'' { j := i + 1 for j < len(sql) { if sql[j] == '\'' { if j+1 < len(sql) && sql[j+1] == '\'' { j += 2 // skip doubled quote '' continue } break } j++ } current.WriteString(sql[i : j+1]) i = j + 1 continue } // Check for dollar-quote start/end if !inDollarQuote && sql[i] == '$' { j := i + 1 for j < len(sql) && (sql[j] == '_' || (sql[j] >= 'a' && sql[j] <= 'z') || (sql[j] >= 'A' && sql[j] <= 'Z') || (sql[j] >= '0' && sql[j] <= '9')) { j++ } if j < len(sql) && sql[j] == '$' { dollarTag = sql[i : j+1] current.WriteString(dollarTag) inDollarQuote = true i = j + 1 continue } } if inDollarQuote && strings.HasPrefix(sql[i:], dollarTag) { current.WriteString(dollarTag) i += len(dollarTag) inDollarQuote = false dollarTag = "" continue } if !inDollarQuote && sql[i] == ';' { statements = append(statements, current.String()) current.Reset() i++ continue } current.WriteByte(sql[i]) i++ } if strings.TrimSpace(current.String()) != "" { statements = append(statements, current.String()) } return statements }