vibedns/internal/database/migrate.go

package database

import (
	"context"
	"crypto/sha256"
	"database/sql"
	"embed"
	"encoding/hex"
	"fmt"
	"io/fs"
	"sort"
	"strconv"
	"strings"
)

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

// Migration is one versioned schema change, embedded in the binary.
type Migration struct {
	Version int
	Name    string
	SQL     string
}

// MigrationStatus reports whether a migration has been applied.
type MigrationStatus struct {
	Version   int    `json:"version"`
	Name      string `json:"name"`
	Applied   bool   `json:"applied"`
	AppliedAt int64  `json:"applied_at,omitempty"`
	Checksum  string `json:"checksum"`
	Drifted   bool   `json:"drifted"`
}

// loadMigrations reads and orders the embedded migration files. File names must
// look like "0001_description.sql".
func loadMigrations() ([]Migration, error) {
	entries, err := fs.ReadDir(migrationFS, "migrations")
	if err != nil {
		return nil, fmt.Errorf("read embedded migrations: %w", err)
	}
	var out []Migration
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".sql") {
			continue
		}
		base := strings.TrimSuffix(e.Name(), ".sql")
		parts := strings.SplitN(base, "_", 2)
		if len(parts) != 2 {
			return nil, fmt.Errorf("migration %q: expected NNNN_name.sql", e.Name())
		}
		v, err := strconv.Atoi(parts[0])
		if err != nil {
			return nil, fmt.Errorf("migration %q: bad version prefix: %w", e.Name(), err)
		}
		body, err := migrationFS.ReadFile("migrations/" + e.Name())
		if err != nil {
			return nil, fmt.Errorf("read migration %q: %w", e.Name(), err)
		}
		out = append(out, Migration{Version: v, Name: parts[1], SQL: string(body)})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Version < out[j].Version })

	for i := 1; i < len(out); i++ {
		if out[i].Version == out[i-1].Version {
			return nil, fmt.Errorf("duplicate migration version %d", out[i].Version)
		}
	}
	return out, nil
}

func checksum(s string) string {
	sum := sha256.Sum256([]byte(s))
	return hex.EncodeToString(sum[:])
}

// ensureMigrationTable creates the migration bookkeeping table.
func (db *DB) ensureMigrationTable(ctx context.Context) error {
	_, err := db.ExecContext(ctx, `
		CREATE TABLE IF NOT EXISTS schema_migrations (
			version    INTEGER PRIMARY KEY,
			name       TEXT    NOT NULL,
			checksum   TEXT    NOT NULL,
			applied_at INTEGER NOT NULL DEFAULT (unixepoch())
		)`)
	if err != nil {
		return fmt.Errorf("create schema_migrations: %w", err)
	}
	return nil
}

type appliedMigration struct {
	name      string
	checksum  string
	appliedAt int64
}

func (db *DB) appliedMigrations(ctx context.Context) (map[int]appliedMigration, error) {
	rows, err := db.QueryContext(ctx, `SELECT version, name, checksum, applied_at FROM schema_migrations`)
	if err != nil {
		return nil, fmt.Errorf("read schema_migrations: %w", err)
	}
	defer rows.Close()

	out := map[int]appliedMigration{}
	for rows.Next() {
		var v int
		var a appliedMigration
		if err := rows.Scan(&v, &a.name, &a.checksum, &a.appliedAt); err != nil {
			return nil, err
		}
		out[v] = a
	}
	return out, rows.Err()
}

// Migrate applies every pending migration in version order. It returns the
// number of migrations that were applied.
func (db *DB) Migrate(ctx context.Context) (int, error) {
	if err := db.ensureMigrationTable(ctx); err != nil {
		return 0, err
	}
	migrations, err := loadMigrations()
	if err != nil {
		return 0, err
	}
	applied, err := db.appliedMigrations(ctx)
	if err != nil {
		return 0, err
	}

	count := 0
	for _, m := range migrations {
		sum := checksum(m.SQL)
		if prev, ok := applied[m.Version]; ok {
			if prev.checksum != sum {
				return count, fmt.Errorf(
					"migration %04d_%s was modified after being applied (expected checksum %s, found %s); "+
						"roll the change into a new migration instead of editing history",
					m.Version, m.Name, prev.checksum, sum)
			}
			continue
		}

		// Each migration is one transaction: a failure leaves no partial schema.
		err := db.InTx(ctx, func(tx *sql.Tx) error {
			if _, err := tx.ExecContext(ctx, m.SQL); err != nil {
				return fmt.Errorf("apply migration %04d_%s: %w", m.Version, m.Name, err)
			}
			_, err := tx.ExecContext(ctx,
				`INSERT INTO schema_migrations (version, name, checksum) VALUES (?, ?, ?)`,
				m.Version, m.Name, sum)
			return err
		})
		if err != nil {
			return count, err
		}
		count++
	}
	return count, nil
}

// SchemaVersion returns the highest applied migration version, or 0.
func (db *DB) SchemaVersion(ctx context.Context) (int, error) {
	if err := db.ensureMigrationTable(ctx); err != nil {
		return 0, err
	}
	var v int
	err := db.QueryRowContext(ctx, `SELECT COALESCE(MAX(version), 0) FROM schema_migrations`).Scan(&v)
	if err != nil {
		return 0, fmt.Errorf("read schema version: %w", err)
	}
	return v, nil
}

// MigrationStatuses lists every known migration and whether it is applied.
func (db *DB) MigrationStatuses(ctx context.Context) ([]MigrationStatus, error) {
	if err := db.ensureMigrationTable(ctx); err != nil {
		return nil, err
	}
	migrations, err := loadMigrations()
	if err != nil {
		return nil, err
	}
	applied, err := db.appliedMigrations(ctx)
	if err != nil {
		return nil, err
	}
	out := make([]MigrationStatus, 0, len(migrations))
	for _, m := range migrations {
		sum := checksum(m.SQL)
		st := MigrationStatus{Version: m.Version, Name: m.Name, Checksum: sum[:12]}
		if a, ok := applied[m.Version]; ok {
			st.Applied = true
			st.AppliedAt = a.appliedAt
			st.Drifted = a.checksum != sum
		}
		out = append(out, st)
	}
	return out, nil
}