EvoBGP/internal/db/apply.go

package db

import (
	"context"
	"errors"
	"fmt"
	"io/fs"
	"path"
	"sort"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

const (
	postgresMigrationLockKey int64 = 0x45766f4247504d47 // "EvoBGPMG", scoped to the database.
	migrationLockTimeout           = 30 * time.Second
	migrationCleanupTimeout        = 5 * time.Second
)

// PostgreSQL migrations are transaction-only. Non-transactional statements
// need an explicit, separately reviewed execution path.

// ApplyPostgresMigrations runs ordered *.up.sql from an embedded subtree.
// Each migration and its version row commit atomically under a session advisory lock.
func ApplyPostgresMigrations(ctx context.Context, pool *pgxpool.Pool, fsys fs.FS, subdir string) (retErr error) {
	entries, err := fs.ReadDir(fsys, subdir)
	if err != nil {
		return fmt.Errorf("db: read migrations: %w", err)
	}
	var ups []string
	for _, e := range entries {
		if !e.IsDir() && strings.HasSuffix(e.Name(), ".up.sql") {
			ups = append(ups, e.Name())
		}
	}
	sort.Strings(ups)
	// Include pool acquisition in the lock timeout. The lock precedes even
	// schema_migrations creation, because another process may start from an empty DB.
	lockCtx, cancelLock := context.WithTimeout(ctx, migrationLockTimeout)
	conn, err := pool.Acquire(lockCtx)
	if err != nil {
		cancelLock()
		return fmt.Errorf("db: acquire migration connection: %w", err)
	}
	defer conn.Release()
	if _, err := conn.Exec(lockCtx, `SELECT pg_advisory_lock($1)`, postgresMigrationLockKey); err != nil {
		cancelLock()
		// Cancellation may race with a successful server-side lock acquisition.
		closeCtx, closeCancel := context.WithTimeout(context.Background(), migrationCleanupTimeout)
		defer closeCancel()
		if closeErr := conn.Conn().Close(closeCtx); closeErr != nil {
			return errors.Join(fmt.Errorf("db: acquire migration lock: %w", err), fmt.Errorf("db: close migration connection: %w", closeErr))
		}
		return fmt.Errorf("db: acquire migration lock: %w", err)
	}
	cancelLock()
	defer func() {
		cleanupCtx, cancel := context.WithTimeout(context.Background(), migrationCleanupTimeout)
		defer cancel()
		var unlocked bool
		if err := conn.QueryRow(cleanupCtx, `SELECT pg_advisory_unlock($1)`, postgresMigrationLockKey).Scan(&unlocked); err != nil || !unlocked {
			if err == nil {
				err = errors.New("lock was not held")
			}
			retErr = errors.Join(retErr, fmt.Errorf("db: release migration lock: %w", err))
			// Never return a possibly locked session to the pool.
			closeCtx, closeCancel := context.WithTimeout(context.Background(), migrationCleanupTimeout)
			defer closeCancel()
			if closeErr := conn.Conn().Close(closeCtx); closeErr != nil {
				retErr = errors.Join(retErr, fmt.Errorf("db: close migration connection: %w", closeErr))
			}
		}
	}()
	if _, err := conn.Exec(ctx, `CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT NOT NULL PRIMARY KEY)`); err != nil {
		return fmt.Errorf("db: schema_migrations: %w", err)
	}
	for _, name := range ups {
		body, err := fs.ReadFile(fsys, path.Join(subdir, name))
		if err != nil {
			return fmt.Errorf("db: read migration %s: %w", name, err)
		}
		if bad := unsupportedMigrationStatement(string(body)); bad != "" {
			return fmt.Errorf("db: migration %s uses unsupported non-transactional SQL %q", name, bad)
		}
		if err := applyPostgresMigration(ctx, conn, strings.TrimSuffix(name, ".up.sql"), body); err != nil {
			return fmt.Errorf("db: migrate %s: %w", name, err)
		}
	}
	return nil
}

// unsupportedMigrationStatement inspects statement starts after masking SQL
// comments and quoted strings. In particular, a comment after a semicolon
// must not conceal COMMIT or another command that would escape our transaction.
func unsupportedMigrationStatement(sql string) string {
	const (
		normal = iota
		singleQuote
		doubleQuote
		lineComment
		blockComment
		dollarQuote
	)
	state := normal
	blockDepth := 0
	dollarTag := ""
	escapeString := false
	var clean strings.Builder
	for i := 0; i < len(sql); {
		switch state {
		case normal:
			switch {
			case sql[i] == '\'':
				escapeString = i > 0 && (sql[i-1] == 'E' || sql[i-1] == 'e') && (i == 1 || !isSQLIdentifierByte(sql[i-2]))
				state = singleQuote
				clean.WriteByte(' ')
				i++
			case sql[i] == '"':
				state = doubleQuote
				clean.WriteByte(' ')
				i++
			case strings.HasPrefix(sql[i:], "--"):
				state = lineComment
				clean.WriteByte(' ')
				i += 2
			case strings.HasPrefix(sql[i:], "/*"):
				state = blockComment
				blockDepth = 1
				clean.WriteByte(' ')
				i += 2
			case sql[i] == '$' && (i == 0 || !isSQLIdentifierByte(sql[i-1])):
				end := i + 1
				if end < len(sql) && isDollarTagStartByte(sql[end]) {
					end++
					for end < len(sql) && isDollarTagByte(sql[end]) {
						end++
					}
				}
				if end < len(sql) && sql[end] == '$' {
					dollarTag = sql[i : end+1]
					state = dollarQuote
					clean.WriteByte(' ')
					i = end + 1
				} else {
					clean.WriteByte(sql[i])
					i++
				}
			default:
				clean.WriteByte(sql[i])
				i++
			}
		case singleQuote:
			switch {
			case escapeString && sql[i] == '\\' && i+1 < len(sql):
				i += 2
			case sql[i] == '\'' && i+1 < len(sql) && sql[i+1] == '\'':
				i += 2
			case sql[i] == '\'':
				state = normal
				i++
			default:
				i++
			}
		case doubleQuote:
			if sql[i] == '"' {
				if i+1 < len(sql) && sql[i+1] == '"' {
					i += 2
				} else {
					state = normal
					i++
				}
			} else {
				i++
			}
		case lineComment:
			if sql[i] == '\n' || sql[i] == '\r' {
				clean.WriteByte('\n')
				state = normal
			}
			i++
		case blockComment:
			switch {
			case strings.HasPrefix(sql[i:], "/*"):
				blockDepth++
				i += 2
			case strings.HasPrefix(sql[i:], "*/"):
				blockDepth--
				i += 2
				if blockDepth == 0 {
					state = normal
				}
			default:
				i++
			}
		case dollarQuote:
			if strings.HasPrefix(sql[i:], dollarTag) {
				i += len(dollarTag)
				state = normal
			} else {
				i++
			}
		}
	}
	for _, statement := range strings.Split(clean.String(), ";") {
		words := strings.Fields(strings.ToUpper(statement))
		if isUnsupportedMigrationCommand(words) {
			return strings.Join(words[:min(len(words), 4)], " ")
		}
	}
	return ""
}

func isSQLIdentifierByte(b byte) bool {
	return isDollarTagByte(b) || b == '$'
}

func isDollarTagStartByte(b byte) bool {
	return b >= 'a' && b <= 'z' || b >= 'A' && b <= 'Z' || b == '_' || b >= 0x80
}

func isDollarTagByte(b byte) bool {
	return isDollarTagStartByte(b) || b >= '0' && b <= '9'
}

func isUnsupportedMigrationCommand(words []string) bool {
	if len(words) == 0 {
		return false
	}
	second := ""
	if len(words) > 1 {
		second = words[1]
	}
	switch words[0] {
	case "BEGIN", "COMMIT", "END", "ROLLBACK", "ABORT", "SAVEPOINT", "REINDEX", "VACUUM", "CLUSTER":
		return true
	case "START", "PREPARE":
		return second == "TRANSACTION"
	case "RELEASE":
		return second == "SAVEPOINT" || second != ""
	case "ALTER":
		return second == "SYSTEM"
	case "CREATE":
		if second == "DATABASE" || second == "TABLESPACE" {
			return true
		}
		idx := 2
		if second == "UNIQUE" {
			idx = 3
		}
		return len(words) > idx && words[idx-1] == "INDEX" && words[idx] == "CONCURRENTLY"
	case "DROP":
		return second == "DATABASE" || second == "TABLESPACE" || second == "INDEX" && len(words) > 2 && words[2] == "CONCURRENTLY"
	case "REFRESH":
		return len(words) > 3 && second == "MATERIALIZED" && words[2] == "VIEW" && words[3] == "CONCURRENTLY"
	case "DISCARD":
		return second == "ALL"
	}
	return false
}

func applyPostgresMigration(ctx context.Context, conn *pgxpool.Conn, version string, body []byte) (retErr error) {
	tx, err := conn.Begin(ctx)
	if err != nil {
		return err
	}
	committed := false
	defer func() {
		if committed {
			return
		}
		cleanupCtx, cancel := context.WithTimeout(context.Background(), migrationCleanupTimeout)
		defer cancel()
		if err := tx.Rollback(cleanupCtx); err != nil && !errors.Is(err, pgx.ErrTxClosed) {
			retErr = errors.Join(retErr, fmt.Errorf("db: rollback migration %s: %w", version, err))
		}
	}()
	var present int
	err = tx.QueryRow(ctx, `SELECT 1 FROM schema_migrations WHERE version = $1`, version).Scan(&present)
	if err == nil {
		return nil
	}
	if !errors.Is(err, pgx.ErrNoRows) {
		return err
	}
	if _, err := tx.Exec(ctx, string(body)); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `INSERT INTO schema_migrations (version) VALUES ($1)`, version); err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return err
	}
	committed = true
	return nil
}