snapshot of https://github.com/ory/keto.git (master, 2026-09-16; blobs > 131072B omitted)

oryx/popx/migration_box.go

@@ -0,0 +1,371 @@
+// Copyright © 2023 Ory Corp
+// SPDX-License-Identifier: Apache-2.0
+
+package popx
+
+import (
+	"database/sql"
+	"fmt"
+	"io/fs"
+	"path"
+	"regexp"
+	"slices"
+	"sort"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ory/pop/v6"
+
+	"github.com/ory/x/dbal"
+	"github.com/ory/x/logrusx"
+)
+
+type (
+	// MigrationBox is a embed migration box.
+	MigrationBox struct {
+		c                      *pop.Connection
+		migrationsUp           Migrations
+		migrationsDown         Migrations
+		perMigrationTimeout    time.Duration
+		l                      *logrusx.Logger
+		migrationContent       MigrationContent
+		disableGoldenDatabase  bool
+		sqliteTemplateCacheDir string
+	}
+	MigrationContent   func(mf Migration, c *pop.Connection, r []byte, usingTemplate bool) (string, error)
+	MigrationBoxOption func(*MigrationBox)
+)
+
+func WithTemplateValues(v map[string]interface{}) MigrationBoxOption {
+	return func(m *MigrationBox) {
+		m.migrationContent = ParameterizedMigrationContent(v)
+	}
+}
+
+func WithMigrationContentMiddleware(middleware func(content string, err error) (string, error)) MigrationBoxOption {
+	return func(m *MigrationBox) {
+		prev := m.migrationContent
+		m.migrationContent = func(mf Migration, c *pop.Connection, r []byte, usingTemplate bool) (string, error) {
+			return middleware(prev(mf, c, r, usingTemplate))
+		}
+	}
+}
+
+// WithGoMigrations adds migrations that have a custom migration runner.
+// TEST THEM THOROUGHLY!
+// It will be very hard to fix a buggy migration.
+func WithGoMigrations(migrations Migrations) MigrationBoxOption {
+	return func(mb *MigrationBox) {
+		for _, m := range migrations {
+			switch m.Direction {
+			case "up":
+				mb.migrationsUp = append(mb.migrationsUp, m)
+			case "down":
+				mb.migrationsDown = append(mb.migrationsDown, m)
+			default:
+				panic(fmt.Sprintf("unknown migration direction %q for %q", m.Direction, m.Version))
+			}
+		}
+	}
+}
+
+func WithPerMigrationTimeout(timeout time.Duration) MigrationBoxOption {
+	return func(m *MigrationBox) {
+		m.perMigrationTimeout = timeout
+	}
+}
+
+// WithoutGoldenDatabase disables the golden database optimisation for SQLite
+// test databases. Use this when the caller asserts on migration execution
+// side-effects (e.g. exact timing, call counts, or intermediate schema state).
+func WithoutGoldenDatabase() MigrationBoxOption {
+	return func(m *MigrationBox) {
+		m.disableGoldenDatabase = true
+	}
+}
+
+// WithSQLiteTemplateCacheDir stores pre-migrated SQLite test database
+// templates in dir. By default, templates are stored in os.TempDir.
+//
+// This option only affects on-disk SQLite databases while running tests.
+func WithSQLiteTemplateCacheDir(dir string) MigrationBoxOption {
+	return func(m *MigrationBox) {
+		m.sqliteTemplateCacheDir = dir
+	}
+}
+
+var testdataPattern = regexp.MustCompile(`^(\d+)_testdata(|\.[a-zA-Z0-9]+).sql$`)
+
+// WithTestdata adds testdata to the migration box.
+func WithTestdata(t *testing.T, testdata fs.FS) MigrationBoxOption {
+	return func(m *MigrationBox) {
+		// Testdata may not change the golden-database hash,
+		// so it should not share or create a golden database.
+		m.disableGoldenDatabase = true
+		require.NoError(t, fs.WalkDir(testdata, ".", func(path string, info fs.DirEntry, err error) error {
+			if err != nil {
+				return errors.WithStack(err)
+			}
+			if info.Name() == "." {
+				return nil
+			}
+			if !info.Type().IsRegular() {
+				t.Logf("skipping testdata entry that is not a file: %s", path)
+				return nil
+			}
+
+			match := testdataPattern.FindStringSubmatch(info.Name())
+			if len(match) != 2 && len(match) != 3 {
+				t.Logf(`WARNING! Found a test migration which does not match the test data pattern: %s`, info.Name())
+				return nil
+			}
+
+			version := match[1]
+			flavor := "all"
+			if len(match) == 3 && len(match[2]) > 0 {
+				flavor = pop.CanonicalDialect(strings.TrimPrefix(match[2], "."))
+			}
+
+			// t.Logf("Found test migration \"%s\" (%s, %+v): %s", flavor, match, err, info.Name())
+
+			m.migrationsUp = append(m.migrationsUp, Migration{
+				Version:   version + "9", // run testdata after version
+				Path:      path,
+				Name:      info.Name(),
+				DBType:    flavor,
+				Direction: "up",
+				Type:      "sql",
+				Runner: func(m Migration, c *pop.Connection) error {
+					b, err := fs.ReadFile(testdata, m.Path)
+					if err != nil {
+						return errors.WithStack(err)
+					}
+					if isMigrationEmpty(string(b)) {
+						return nil
+					}
+					_, err = c.Store.SQLDB().Exec(string(b))
+					return errors.WithStack(err)
+				},
+			})
+
+			m.migrationsDown = append(m.migrationsDown, Migration{
+				Version:   version + "9", // run testdata after version
+				Path:      path,
+				Name:      info.Name(),
+				DBType:    flavor,
+				Direction: "down",
+				Type:      "sql",
+				Runner:    func(m Migration, _ *pop.Connection) error { return nil },
+			})
+
+			return nil
+		}))
+	}
+}
+
+var emptySQLReplace = regexp.MustCompile(`(?m)^(\s*--.*|\s*)$`)
+
+func isMigrationEmpty(content string) bool {
+	return len(strings.ReplaceAll(emptySQLReplace.ReplaceAllString(content, ""), "\n", "")) == 0
+}
+
+type queryExecutor interface {
+	Exec(query string, args ...any) (sql.Result, error)
+}
+
+// NewMigrationBox creates a new migration box.
+func NewMigrationBox(dir fs.FS, c *pop.Connection, l *logrusx.Logger, opts ...MigrationBoxOption) (*MigrationBox, error) {
+	mb := &MigrationBox{
+		c:                c,
+		l:                l,
+		migrationContent: ParameterizedMigrationContent(nil),
+	}
+
+	for _, o := range opts {
+		o(mb)
+	}
+
+	txRunner := func(b []byte) func(Migration, *pop.Connection) error {
+		return func(mf Migration, c *pop.Connection) error {
+			content, err := mb.migrationContent(mf, c, b, true)
+			if err != nil {
+				return errors.Wrapf(err, "error processing %s", mf.Path)
+			}
+			if isMigrationEmpty(content) {
+				l.WithField("migration", mf.Path).Trace("This is usually ok - ignoring migration because content is empty. This is ok!")
+				return nil
+			}
+
+			var q queryExecutor = c.Store.SQLDB()
+			if c.TX != nil {
+				q = c.TX
+			}
+
+			if _, err = q.Exec(content); err != nil {
+				return errors.Wrapf(err, "error executing %s, sql: %s", mf.Path, content)
+			}
+			return nil
+		}
+	}
+
+	err := mb.findMigrations(dir, txRunner)
+	if err != nil {
+		return mb, err
+	}
+
+	if err := mb.check(); err != nil {
+		return nil, err
+	}
+	return mb, nil
+}
+
+// migrationFallbacks returns the ordered dialects migration selection falls
+// back to before the generic ("all") file (.yugabyte. → .postgres. → generic
+// for YugabyteDB, no fallbacks otherwise).
+func (mb *MigrationBox) migrationFallbacks() []string {
+	if mb.c.Dialect.Name() == dbal.DriverYugabyteDB {
+		return []string{dbal.DriverPostgreSQL}
+	}
+	return nil
+}
+
+// SelectedMigrations returns a copy of the migrations selected for the
+// connection's dialect and direction. It uses the same exact-match and fallback
+// ranking as Up, Down, and Status so callers can audit the effective migration
+// set without duplicating filename-selection logic.
+func (mb *MigrationBox) SelectedMigrations(direction string) (Migrations, error) {
+	dialect := mb.c.Dialect.Name()
+	fallbacks := mb.migrationFallbacks()
+	switch direction {
+	case "up":
+		return slices.Clone(mb.migrationsUp.sortAndFilter(dialect, fallbacks...)), nil
+	case "down":
+		return slices.Clone(mb.migrationsDown.sortAndFilter(dialect, fallbacks...)), nil
+	default:
+		return nil, errors.Errorf("unknown migration direction %q", direction)
+	}
+}
+
+// noTxDDL reports whether the dialect cannot run DDL inside a transaction, so
+// both migrations and the migration-status table setup must run in autocommit
+// mode. CockroachDB and MySQL auto-commit each DDL statement; YugabyteDB
+// restricts DDL inside transactions. Because there is no surrounding
+// transaction to roll back a partial failure, new YugabyteDB migrations must
+// contain one idempotent statement per file. YugabyteDB also inherits older
+// PostgreSQL migration files that predate this rule; service-level guard tests
+// track that replay risk, and a dedicated .yugabyte. override is required when
+// a new unsafe inherited migration appears.
+func (mb *MigrationBox) noTxDDL() bool {
+	switch mb.c.Dialect.Name() {
+	case dbal.DriverCockroachDB, dbal.DriverMySQL, dbal.DriverYugabyteDB:
+		return true
+	}
+	return false
+}
+
+func (mb *MigrationBox) findMigrations(
+	dir fs.FS,
+	runner func([]byte) func(m Migration, c *pop.Connection) error,
+) error {
+	err := fs.WalkDir(dir, ".", func(p string, info fs.DirEntry, err error) error {
+		if err != nil {
+			return errors.WithStack(err)
+		}
+
+		if !info.Type().IsRegular() {
+			mb.l.Tracef("ignoring non file: %s", info.Name())
+			return nil
+		}
+
+		if path.Ext(info.Name()) != ".sql" {
+			mb.l.Tracef("ignoring non SQL file: %s", info.Name())
+			return nil
+		}
+
+		details, err := parseMigrationFilename(info.Name())
+		if err != nil {
+			if errors.Is(err, errUnsupportedMigrationDialect) {
+				mb.l.Debugf("Ignoring migration file %s because its dialect is not supported: %s", info.Name(), err.Error())
+				return nil
+			}
+			return errors.WithStack(err)
+		}
+
+		if details == nil {
+			return errors.Errorf("Found a migration file that does not match the file pattern: filename=%s pattern=%s", info.Name(), MigrationFileRegexp)
+		}
+
+		content, err := fs.ReadFile(dir, p)
+		if err != nil {
+			return errors.WithStack(err)
+		}
+
+		mf := Migration{
+			Path:       p,
+			Version:    details.Version,
+			Name:       details.Name,
+			DBType:     details.DBType,
+			Direction:  details.Direction,
+			Type:       details.Type,
+			Content:    string(content),
+			Autocommit: details.Autocommit,
+		}
+
+		mf.Runner = runner(content)
+
+		switch details.Direction {
+		case "up":
+			mb.migrationsUp = append(mb.migrationsUp, mf)
+		case "down":
+			mb.migrationsDown = append(mb.migrationsDown, mf)
+		default:
+			return errors.Errorf("unknown migration direction %q for %q", details.Direction, info.Name())
+		}
+		return nil
+	})
+
+	// Sort descending.
+	sort.Sort(mb.migrationsDown)
+	slices.Reverse(mb.migrationsDown)
+
+	// Sort ascending.
+	sort.Sort(mb.migrationsUp)
+
+	return errors.WithStack(err)
+}
+
+// hasDownMigrationWithVersion checks if there is a migration with the given
+// version.
+func (mb *MigrationBox) hasDownMigrationWithVersion(version string) bool {
+	for _, down := range mb.migrationsDown {
+		if version == down.Version {
+			return true
+		}
+	}
+	return false
+}
+
+// check checks that every "up" migration has a corresponding "down" migration.
+func (mb *MigrationBox) check() error {
+	for _, up := range mb.migrationsUp {
+		if !mb.hasDownMigrationWithVersion(up.Version) {
+			return errors.Errorf("migration %s has no corresponding down migration", up.Version)
+		}
+	}
+
+	for _, n := range mb.migrationsUp {
+		if err := n.Valid(); err != nil {
+			return errors.WithStack(err)
+		}
+	}
+	for _, n := range mb.migrationsDown {
+		if err := n.Valid(); err != nil {
+			return errors.WithStack(err)
+		}
+	}
+	return nil
+}