oikos/internal/db/pool.go

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
}