pgsql-broker/pkg/broker/install/install.go

package install

import (
	"context"
	"embed"
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"strings"

	"github.com/lib/pq"

	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter"
)

//go:embed all:sql/migrations
var migrationsFS embed.FS

//go:embed all:sql/roles
var rolesFS embed.FS

const migrationsDir = "sql/migrations"
const rolesDir = "sql/roles"

// migrationsTableSQL creates the version-tracking table itself. It is applied
// unconditionally (idempotently) before any numbered migration file, and is
// not itself a numbered migration.
const migrationsTableSQL = `
CREATE SCHEMA IF NOT EXISTS broker;
CREATE TABLE IF NOT EXISTS broker.broker_schema_migrations (
    version    BIGINT PRIMARY KEY,
    name       TEXT NOT NULL,
    applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
`

var migrationFileRe = regexp.MustCompile(`^(\d+)_(.+)\.sql$`)

// migrationFile describes one embedded migration.
type migrationFile struct {
	version int64
	name    string
	path    string
}

// Installer handles database schema installation via versioned migrations.
type Installer struct {
	db     adapter.DBAdapter
	logger adapter.Logger
}

// New creates a new installer
func New(db adapter.DBAdapter, logger adapter.Logger) *Installer {
	return &Installer{
		db:     db,
		logger: logger,
	}
}

// loadMigrations reads and sorts every embedded migration file by numeric prefix.
func loadMigrations() ([]migrationFile, error) {
	entries, err := migrationsFS.ReadDir(migrationsDir)
	if err != nil {
		return nil, fmt.Errorf("failed to read migrations directory: %w", err)
	}

	var migrations []migrationFile
	for _, e := range entries {
		if e.IsDir() {
			continue
		}
		m := migrationFileRe.FindStringSubmatch(e.Name())
		if m == nil {
			continue
		}
		version, err := strconv.ParseInt(m[1], 10, 64)
		if err != nil {
			return nil, fmt.Errorf("invalid migration filename %s: %w", e.Name(), err)
		}
		migrations = append(migrations, migrationFile{
			version: version,
			name:    m[2],
			path:    migrationsDir + "/" + e.Name(),
		})
	}

	sort.Slice(migrations, func(i, j int) bool { return migrations[i].version < migrations[j].version })
	return migrations, nil
}

// ensureMigrationsTable creates the broker schema and the migrations
// tracking table if they don't already exist. This is DDL and requires
// CREATE privilege on the database -- only ApplyMigrations (run by an
// admin-privileged connection, e.g. `pgsql-broker install`) calls it.
func (i *Installer) ensureMigrationsTable(ctx context.Context) error {
	if _, err := i.db.Exec(ctx, migrationsTableSQL); err != nil {
		return fmt.Errorf("failed to ensure migrations table: %w", err)
	}
	return nil
}

// migrationsTableExists reports whether the migrations tracking table is
// present, without creating it -- a read-only check safe to run with a
// least-privilege runtime role (e.g. broker_runtime) that has no CREATE
// privilege on the database.
func (i *Installer) migrationsTableExists(ctx context.Context) (bool, error) {
	var exists bool
	err := i.db.QueryRow(ctx, "SELECT to_regclass('broker.broker_schema_migrations') IS NOT NULL").Scan(&exists)
	if err != nil {
		return false, fmt.Errorf("failed to check migrations table: %w", err)
	}
	return exists, nil
}

// appliedVersions returns the set of migration versions already recorded.
func (i *Installer) appliedVersions(ctx context.Context) (map[int64]bool, error) {
	rows, err := i.db.Query(ctx, "SELECT version FROM broker.broker_schema_migrations")
	if err != nil {
		return nil, fmt.Errorf("failed to query applied migrations: %w", err)
	}
	defer rows.Close()

	applied := make(map[int64]bool)
	for rows.Next() {
		var v int64
		if err := rows.Scan(&v); err != nil {
			return nil, fmt.Errorf("failed to scan migration version: %w", err)
		}
		applied[v] = true
	}
	return applied, rows.Err()
}

// PendingMigrations returns the names of embedded migrations that have not
// yet been applied to the database, without applying them or creating the
// migrations table -- safe to call with a least-privilege runtime role.
func (i *Installer) PendingMigrations(ctx context.Context) ([]string, error) {
	migrations, err := loadMigrations()
	if err != nil {
		return nil, err
	}

	exists, err := i.migrationsTableExists(ctx)
	if err != nil {
		return nil, err
	}
	if !exists {
		pending := make([]string, len(migrations))
		for idx, m := range migrations {
			pending[idx] = fmt.Sprintf("%04d_%s", m.version, m.name)
		}
		return pending, nil
	}

	applied, err := i.appliedVersions(ctx)
	if err != nil {
		return nil, err
	}

	var pending []string
	for _, m := range migrations {
		if !applied[m.version] {
			pending = append(pending, fmt.Sprintf("%04d_%s", m.version, m.name))
		}
	}
	return pending, nil
}

// ApplyMigrations applies every embedded migration that has not yet been
// recorded in broker.broker_schema_migrations, each inside its own transaction.
func (i *Installer) ApplyMigrations(ctx context.Context) error {
	i.logger.Info("applying migrations")

	if err := i.ensureMigrationsTable(ctx); err != nil {
		return err
	}

	migrations, err := loadMigrations()
	if err != nil {
		return err
	}

	applied, err := i.appliedVersions(ctx)
	if err != nil {
		return err
	}

	appliedCount := 0
	for _, m := range migrations {
		if applied[m.version] {
			continue
		}

		content, err := migrationsFS.ReadFile(m.path)
		if err != nil {
			return fmt.Errorf("failed to read migration %s: %w", m.path, err)
		}

		i.logger.Info("applying migration", "version", m.version, "name", m.name)

		tx, err := i.db.Begin(ctx)
		if err != nil {
			return fmt.Errorf("failed to begin transaction for migration %s: %w", m.name, err)
		}

		if err := execStatements(ctx, tx, string(content)); err != nil {
			if rbErr := tx.Rollback(); rbErr != nil {
				i.logger.Error("failed to rollback migration transaction", "version", m.version, "error", rbErr)
			}
			return fmt.Errorf("failed to apply migration %s: %w", m.name, err)
		}

		if _, err := tx.Exec(ctx,
			"INSERT INTO broker.broker_schema_migrations (version, name) VALUES ($1, $2)",
			m.version, m.name,
		); err != nil {
			if rbErr := tx.Rollback(); rbErr != nil {
				i.logger.Error("failed to rollback migration transaction", "version", m.version, "error", rbErr)
			}
			return fmt.Errorf("failed to record migration %s: %w", m.name, err)
		}

		if err := tx.Commit(); err != nil {
			return fmt.Errorf("failed to commit migration %s: %w", m.name, err)
		}

		appliedCount++
	}

	if appliedCount == 0 {
		i.logger.Info("no pending migrations")
	} else {
		i.logger.Info("migrations applied successfully", "count", appliedCount)
	}
	return nil
}

// RolePasswords holds the login passwords for the reference broker_admin,
// broker_runtime, and broker_enqueue roles created by InstallRoles. All
// three are required -- there is no placeholder/default fallback, since
// these roles carry real database privileges.
type RolePasswords struct {
	AdminPassword   string
	RuntimePassword string
	EnqueuePassword string
}

// InstallRoles applies the embedded role/grant scripts (sql/roles), which
// create (or, if already present, rotate the password of) broker_admin,
// broker_runtime, and broker_enqueue, then grant them the appropriate
// schema/table/function privileges. The caller must connect as a superuser
// or a role with CREATEROLE -- this is intentionally separate from the
// migration-running connection.
func (i *Installer) InstallRoles(ctx context.Context, passwords RolePasswords) error {
	if passwords.AdminPassword == "" || passwords.RuntimePassword == "" || passwords.EnqueuePassword == "" {
		return fmt.Errorf("all three role passwords (admin, runtime, enqueue) are required")
	}

	entries, err := rolesFS.ReadDir(rolesDir)
	if err != nil {
		return fmt.Errorf("failed to read roles directory: %w", err)
	}

	var names []string
	for _, e := range entries {
		if !e.IsDir() {
			names = append(names, e.Name())
		}
	}
	sort.Strings(names)

	replacer := strings.NewReplacer(
		"__BROKER_ADMIN_PASSWORD__", pq.QuoteLiteral(passwords.AdminPassword),
		"__BROKER_RUNTIME_PASSWORD__", pq.QuoteLiteral(passwords.RuntimePassword),
		"__BROKER_ENQUEUE_PASSWORD__", pq.QuoteLiteral(passwords.EnqueuePassword),
	)

	for _, name := range names {
		content, err := rolesFS.ReadFile(rolesDir + "/" + name)
		if err != nil {
			return fmt.Errorf("failed to read roles script %s: %w", name, err)
		}

		i.logger.Info("applying roles script", "name", name)

		tx, err := i.db.Begin(ctx)
		if err != nil {
			return fmt.Errorf("failed to begin transaction for roles script %s: %w", name, err)
		}

		// pq.QuoteLiteral already produces a safely quoted SQL string
		// literal (doubling embedded quotes, or switching to E'...' escape
		// syntax if the password contains a backslash), so this is a plain
		// textual substitution, not string concatenation of untrusted input
		// into SQL syntax.
		rendered := replacer.Replace(string(content))

		if err := execStatements(ctx, tx, rendered); err != nil {
			if rbErr := tx.Rollback(); rbErr != nil {
				i.logger.Error("failed to rollback roles script transaction", "name", name, "error", rbErr)
			}
			return fmt.Errorf("failed to apply roles script %s: %w", name, err)
		}

		if err := tx.Commit(); err != nil {
			return fmt.Errorf("failed to commit roles script %s: %w", name, err)
		}
	}

	i.logger.Info("roles installed successfully")
	return nil
}

// execStatements runs every statement in sql within tx.
func execStatements(ctx context.Context, tx adapter.DBTransaction, sqlText string) error {
	statements := splitSQLStatements(sqlText)

	for _, stmt := range statements {
		stmt = strings.TrimSpace(stmt)
		if stmt == "" || strings.HasPrefix(stmt, "\\") {
			continue
		}
		if _, err := tx.Exec(ctx, stmt); err != nil {
			return fmt.Errorf("failed to execute statement: %w\nStatement: %s", err, stmt)
		}
	}
	return nil
}

// splitSQLStatements splits SQL into individual statements, keeping
// $$-quoted function bodies intact.
// splitSQLStatements splits a SQL script into individual statements on
// top-level semicolons, ignoring semicolons that appear inside single-quoted
// strings ('...', with ” as an escaped quote), double-quoted identifiers,
// line comments (--), and dollar-quoted bodies ($$...$$ or $tag$...$tag$).
func splitSQLStatements(sqlText string) []string {
	var result []string
	var buffer strings.Builder

	runes := []rune(sqlText)
	n := len(runes)
	i := 0

	for i < n {
		c := runes[i]

		switch {
		case c == '-' && i+1 < n && runes[i+1] == '-':
			// Line comment: copy through end of line.
			for i < n && runes[i] != '\n' {
				buffer.WriteRune(runes[i])
				i++
			}
			continue

		case c == '\'':
			buffer.WriteRune(c)
			i++
			for i < n {
				buffer.WriteRune(runes[i])
				if runes[i] == '\'' {
					if i+1 < n && runes[i+1] == '\'' {
						buffer.WriteRune(runes[i+1])
						i += 2
						continue
					}
					i++
					break
				}
				i++
			}
			continue

		case c == '"':
			buffer.WriteRune(c)
			i++
			for i < n {
				buffer.WriteRune(runes[i])
				if runes[i] == '"' {
					i++
					break
				}
				i++
			}
			continue

		case c == '$':
			if tag, ok := matchDollarTag(runes, i); ok {
				closer := tag
				buffer.WriteString(closer)
				i += len(closer)
				end := indexOfRunes(runes, i, closer)
				if end == -1 {
					buffer.WriteString(string(runes[i:]))
					i = n
				} else {
					buffer.WriteString(string(runes[i:end]))
					buffer.WriteString(closer)
					i = end + len(closer)
				}
				continue
			}
			buffer.WriteRune(c)
			i++

		case c == ';':
			stmt := strings.TrimSpace(buffer.String())
			if stmt != "" {
				result = append(result, stmt+";")
			}
			buffer.Reset()
			i++

		default:
			buffer.WriteRune(c)
			i++
		}
	}

	if stmt := strings.TrimSpace(buffer.String()); stmt != "" {
		result = append(result, stmt)
	}

	return result
}

// matchDollarTag checks whether runes[pos:] begins a dollar-quote tag
// ($$ or $tag$) and returns that tag if so.
func matchDollarTag(runes []rune, pos int) (string, bool) {
	if runes[pos] != '$' {
		return "", false
	}
	j := pos + 1
	for j < len(runes) && (runes[j] == '_' || isAlnum(runes[j])) {
		j++
	}
	if j < len(runes) && runes[j] == '$' {
		return string(runes[pos : j+1]), true
	}
	return "", false
}

func isAlnum(r rune) bool {
	return (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9')
}

// indexOfRunes returns the index of the first occurrence of sub in
// runes[from:], or -1 if not found.
func indexOfRunes(runes []rune, from int, sub string) int {
	subRunes := []rune(sub)
	for i := from; i+len(subRunes) <= len(runes); i++ {
		match := true
		for j, r := range subRunes {
			if runes[i+j] != r {
				match = false
				break
			}
		}
		if match {
			return i
		}
	}
	return -1
}

// VerifyInstallation checks that every embedded migration has been applied.
func (i *Installer) VerifyInstallation(ctx context.Context) error {
	i.logger.Info("verifying installation")

	pending, err := i.PendingMigrations(ctx)
	if err != nil {
		return fmt.Errorf("failed to check pending migrations: %w", err)
	}

	if len(pending) > 0 {
		return fmt.Errorf("schema is behind: %d migration(s) not applied: %s", len(pending), strings.Join(pending, ", "))
	}

	i.logger.Info("installation verified successfully")
	return nil
}