llmgw/internal/store/migrate.go

package store

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

//go:embed migrations/sqlite/*.sql migrations/postgres/*.sql
var migrationFS embed.FS

// migrate applies every unapplied migration in order, each in its own
// transaction, and records it.
//
// Deliberately small. A migration tool would be a dependency and a second way
// to be wrong about which migrations have run; the whole mechanism is a sorted
// directory, a version table, and a transaction each.
//
// The directory is per dialect. SQLite and Postgres diverge enough in DDL —
// SQLite's ALTER TABLE barely exists — that a single shared set would be a
// promise this could not keep. Same version numbers, same meaning, different
// files.
func migrate(ctx context.Context, db *sql.DB, dialect string) error {
	// BIGINT, not INTEGER. SQLite's INTEGER holds 64 bits; PostgreSQL's is
	// int4, and a millisecond timestamp does not fit in it — which the first
	// migration discovers by failing to record that it ran. This table is
	// created here rather than in a migration file, which is exactly how it
	// escaped the port of the others.
	stamp := "BIGINT"
	if dialect == string(SQLiteDialect) {
		stamp = "INTEGER"
	}
	if _, err := db.ExecContext(ctx,
		`CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at `+stamp+` NOT NULL)`,
	); err != nil {
		return fmt.Errorf("creating the migration table: %w", err)
	}

	applied := map[string]struct{}{}
	rows, err := db.QueryContext(ctx, `SELECT version FROM schema_migrations`)
	if err != nil {
		return err
	}
	for rows.Next() {
		var v string
		if err := rows.Scan(&v); err != nil {
			rows.Close()
			return err
		}
		applied[v] = struct{}{}
	}
	if err := rows.Err(); err != nil {
		rows.Close()
		return err
	}
	rows.Close()

	names, err := migrationNames(dialect)
	if err != nil {
		return err
	}

	for _, name := range names {
		if _, done := applied[name]; done {
			continue
		}
		body, err := migrationFS.ReadFile("migrations/" + dialect + "/" + name)
		if err != nil {
			return err
		}
		if err := applyOne(ctx, db, dialect, name, string(body)); err != nil {
			return fmt.Errorf("migration %s: %w", name, err)
		}
	}
	return nil
}

func migrationNames(dialect string) ([]string, error) {
	entries, err := fs.ReadDir(migrationFS, "migrations/"+dialect)
	if err != nil {
		return nil, fmt.Errorf("no migrations for dialect %q: %w", dialect, err)
	}
	var names []string
	for _, e := range entries {
		if !e.IsDir() && strings.HasSuffix(e.Name(), ".sql") {
			names = append(names, e.Name())
		}
	}
	// Lexical order is version order, which is why the files are zero padded.
	sort.Strings(names)
	return names, nil
}

// fkOffMarker in a SQLite migration runs it with foreign keys off.
//
// Rebuilding a table — the only way SQLite changes a column's constraints —
// drops the old one, and with foreign keys on that deletes every row other
// tables point at. SQLite's procedure for this turns them off, which only
// works outside a transaction, so this takes one connection, turns them off
// on it, runs the migration in a transaction there, and refuses to commit
// unless foreign_key_check finds nothing.
const fkOffMarker = "-- migrate: foreign_keys=off"

func applyOne(ctx context.Context, db *sql.DB, dialect, name, body string) error {
	type execer interface {
		BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
	}
	var on execer = db
	checkFKs := false
	if dialect == string(SQLiteDialect) && strings.Contains(body, fkOffMarker) {
		c, err := db.Conn(ctx)
		if err != nil {
			return err
		}
		defer c.Close()
		if _, err := c.ExecContext(ctx, `PRAGMA foreign_keys = OFF`); err != nil {
			return err
		}
		defer func() { _, _ = c.ExecContext(context.WithoutCancel(ctx), `PRAGMA foreign_keys = ON`) }()
		on, checkFKs = c, true
	}

	tx, err := on.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()

	if _, err := tx.ExecContext(ctx, body); err != nil {
		return err
	}
	if checkFKs {
		rows, err := tx.QueryContext(ctx, `PRAGMA foreign_key_check`)
		if err != nil {
			return err
		}
		broken := rows.Next()
		rows.Close()
		if broken {
			return fmt.Errorf("leaves foreign keys pointing at nothing; not applied")
		}
	}
	if _, err := tx.ExecContext(ctx,
		Dialect(dialect).rebind(`INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)`),
		name, nowMillis(),
	); err != nil {
		return err
	}
	return tx.Commit()
}