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 { tx.Rollback() 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 { tx.Rollback() 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 { tx.Rollback() 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 }