Vellum/internal/core/migrate.go

package core

import (
	"context"
	"database/sql"
	"fmt"
	"io/fs"
	"log/slog"
	"sort"
	"strconv"
	"strings"
)

// Migrate applies every embedded migration whose number is greater than the
// database's current PRAGMA user_version, in ascending order, each inside its
// own transaction.
//
// A migration is recorded by its numeric prefix. This is why applied files are
// never edited: the runner has no checksum, only a version number, so changing
// an applied file would silently diverge a running database from the schema.
func Migrate(ctx context.Context, db *sql.DB, fsys fs.FS, log *slog.Logger) error {
	names, err := fs.Glob(fsys, "*.sql")
	if err != nil {
		return fmt.Errorf("list migrations: %w", err)
	}
	sort.Strings(names)

	current, err := userVersion(ctx, db)
	if err != nil {
		return err
	}
	log.Debug("schema version", "current", current, "available", len(names))

	highest := 0
	for _, name := range names {
		n, err := migrationNumber(name)
		if err != nil {
			return err
		}
		if n > highest {
			highest = n
		}
	}
	// A version above every embedded migration means this binary is older than
	// the database. Without this, the loop below skips everything and starts up
	// against a schema it cannot read.
	if current > highest {
		return fmt.Errorf("database schema is version %d but this binary only knows %d: this binary is older than the database — upgrade the binary or restore a backup", current, highest)
	}

	for _, name := range names {
		n, err := migrationNumber(name)
		if err != nil {
			return err
		}
		if n <= current {
			continue
		}
		if err := applyMigration(ctx, db, fsys, name, n); err != nil {
			return fmt.Errorf("migration %s: %w", name, err)
		}
		log.Info("applied migration", "file", name, "version", n)
	}
	return nil
}

func applyMigration(ctx context.Context, db *sql.DB, fsys fs.FS, name string, n int) error {
	body, err := fs.ReadFile(fsys, name)
	if err != nil {
		return err
	}

	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback() //nolint:errcheck // no-op once Commit succeeds

	for _, stmt := range splitStatements(string(body)) {
		if _, err := tx.ExecContext(ctx, stmt); err != nil {
			return fmt.Errorf("exec: %w\nstatement:\n%s", err, stmt)
		}
	}

	// PRAGMA does not accept a bound parameter. n is parsed from the filename
	// as an integer above, so it is not attacker-controlled.
	if _, err := tx.ExecContext(ctx, fmt.Sprintf("PRAGMA user_version = %d", n)); err != nil {
		return fmt.Errorf("set user_version: %w", err)
	}
	return tx.Commit()
}

func userVersion(ctx context.Context, db *sql.DB) (int, error) {
	var v int
	if err := db.QueryRowContext(ctx, "PRAGMA user_version").Scan(&v); err != nil {
		return 0, fmt.Errorf("read user_version: %w", err)
	}
	return v, nil
}

// migrationNumber parses the numeric prefix of "0001_core.sql" -> 1.
func migrationNumber(name string) (int, error) {
	base := name
	if i := strings.IndexByte(base, '_'); i >= 0 {
		base = base[:i]
	} else if i := strings.IndexByte(base, '.'); i >= 0 {
		base = base[:i]
	}
	n, err := strconv.Atoi(base)
	if err != nil {
		return 0, fmt.Errorf("migration %q: filename must start with a number, e.g. 0001_core.sql", name)
	}
	return n, nil
}

// splitStatements splits a SQL file on semicolons that terminate a statement.
//
// Semicolons inside BEGIN...END blocks (trigger bodies) do not split, because
// SQLite's CREATE TRIGGER syntax requires them and the whole trigger must be
// executed as one statement. It is not a general SQL parser: it does not track
// string literals beyond quotes, and it treats a bare END as closing a block,
// so a CASE...END expression inside a trigger would be mis-split. The schema
// files this project writes contain neither.
func splitStatements(sqlText string) []string {
	clean := stripSQLComments(sqlText)
	var out []string
	var b strings.Builder
	depth := 0

	for i := 0; i < len(clean); {
		c := clean[i]

		// Copy string literals verbatim so a ';' or keyword inside one is inert.
		if c == '\'' || c == '"' {
			b.WriteByte(c)
			quote := c
			i++
			for i < len(clean) {
				b.WriteByte(clean[i])
				if clean[i] == quote {
					i++
					break
				}
				i++
			}
			continue
		}

		if isWordByte(c) {
			start := i
			for i < len(clean) && isWordByte(clean[i]) {
				i++
			}
			switch strings.ToUpper(clean[start:i]) {
			case "BEGIN":
				depth++
			case "END":
				if depth > 0 {
					depth--
				}
			}
			b.WriteString(clean[start:i])
			continue
		}

		if c == ';' && depth == 0 {
			if s := strings.TrimSpace(b.String()); s != "" {
				out = append(out, s)
			}
			b.Reset()
			i++
			continue
		}

		b.WriteByte(c)
		i++
	}

	if s := strings.TrimSpace(b.String()); s != "" {
		out = append(out, s)
	}
	return out
}

func isWordByte(c byte) bool {
	return c == '_' ||
		(c >= 'a' && c <= 'z') ||
		(c >= 'A' && c <= 'Z') ||
		(c >= '0' && c <= '9')
}

// stripSQLComments removes -- line comments and /* */ block comments.
func stripSQLComments(s string) string {
	var b strings.Builder
	b.Grow(len(s))

	for i := 0; i < len(s); {
		switch {
		case strings.HasPrefix(s[i:], "--"):
			for i < len(s) && s[i] != '\n' {
				i++
			}
		case strings.HasPrefix(s[i:], "/*"):
			i += 2
			for i < len(s) && !strings.HasPrefix(s[i:], "*/") {
				i++
			}
			i += 2
		default:
			b.WriteByte(s[i])
			i++
		}
	}
	return b.String()
}