feat(core): add WithMigrationBaseline hook before goose Up

2026-09-29 14:06:36 +08:00 · 2026-08-30 11:41:00 +08:00
parent 4ce7110e23
commit a617457a3c
5 changed files with 313 additions and 52 deletions
@@ -41,6 +41,10 @@ import (
 const (
 	defaultShutdownTimeout = 15 * time.Second
 	defaultHTTPAddr        = "127.0.0.1:8000"
+
+	// migrationAdvisoryLockKey serializes baseline + plugin Up across Postgres
+	// sessions (ASCII "wave"). SQLite is single-writer and needs no extra lock.
+	migrationAdvisoryLockKey int64 = 0x77617665
 )

 // runProfileApp prepares and runs the application for a given profile.
@@ -140,17 +144,21 @@ type sharedStore struct {
 func (s *sharedStore) Tablename() string { return "w_schema_versions" }

 func (s *sharedStore) CreateVersionTable(ctx context.Context, db goosedb.DBTxConn) error {
+	_, err := db.ExecContext(ctx, schemaVersionsDDL(s.dialect))
+	return err
+}
+
+func schemaVersionsDDL(dialect string) string {
 	timeType := "TIMESTAMPTZ"
-	if s.dialect == "sqlite3" || s.dialect == "sqlite" {
+	if dialect == "sqlite3" || dialect == "sqlite" {
 		timeType = "DATETIME"
 	}
-	_, err := db.ExecContext(ctx, fmt.Sprintf(`CREATE TABLE IF NOT EXISTS w_schema_versions (
+	return fmt.Sprintf(`CREATE TABLE IF NOT EXISTS w_schema_versions (
 		plugin_id   VARCHAR(64)  NOT NULL,
 		version_id  BIGINT       NOT NULL,
 		applied_at  %s  NOT NULL DEFAULT CURRENT_TIMESTAMP,
 		PRIMARY KEY (plugin_id, version_id)
-	)`, timeType))
-	return err
+	)`, timeType)
 }

 //nolint:mnd
@@ -265,6 +273,38 @@ func (e *gooseEngine) Migrate(ctx *core.Context, entries []core.MigrationEntry)

 	dialect := gooseDialect(ctx)
 	dialectStr := string(dialect)
+	goCtx := context.Background()
+	if ctx != nil {
+		goCtx = ctx.GoContext()
+	}
+	if goCtx == nil {
+		goCtx = context.Background()
+	}
+
+	bootstrap := &sharedStore{dialect: dialectStr}
+	if err := bootstrap.CreateVersionTable(goCtx, sqlDB); err != nil {
+		return fmt.Errorf("migration: create version table: %w", err)
+	}
+
+	if dialect == goose.DialectPostgres {
+		conn, lockErr := sqlDB.Conn(goCtx)
+		if lockErr != nil {
+			return fmt.Errorf("migration: pin connection for advisory lock: %w", lockErr)
+		}
+		defer func() { _ = conn.Close() }()
+		if _, lockErr = conn.ExecContext(goCtx, "SELECT pg_advisory_lock($1)", migrationAdvisoryLockKey); lockErr != nil {
+			return fmt.Errorf("migration: advisory lock: %w", lockErr)
+		}
+		defer func() {
+			_, _ = conn.ExecContext(context.Background(), "SELECT pg_advisory_unlock($1)", migrationAdvisoryLockKey)
+		}()
+	}
+
+	if fn := ctx.MigrationBaseline(); fn != nil {
+		if err := fn(ctx); err != nil {
+			return fmt.Errorf("migration baseline: %w", err)
+		}
+	}

 	for _, entry := range entries {
 		store := &sharedStore{
@@ -0,0 +1,127 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package cmd
+
+import (
+	"Wavelet/core"
+	"Wavelet/core/contracts"
+	"context"
+	"path/filepath"
+	"testing"
+	"testing/fstest"
+
+	"github.com/glebarez/sqlite"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+type migrateTestDB struct {
+	db *gorm.DB
+}
+
+func (s migrateTestDB) GORM() *gorm.DB { return s.db }
+
+func (s migrateTestDB) DB(ctx context.Context) *gorm.DB { return s.db.WithContext(ctx) }
+
+func (s migrateTestDB) Named(string) *gorm.DB { return s.db }
+
+type migrateTestPlugin struct {
+	db *gorm.DB
+	fs fstest.MapFS
+}
+
+func (p *migrateTestPlugin) Name() string { return "t" }
+
+func (p *migrateTestPlugin) Apply(ctx *core.Context) error {
+	core.Provide[contracts.DBService](ctx, migrateTestDB{db: p.db})
+	ctx.Migrations().Register("t", p.fs)
+	return nil
+}
+
+func sqliteTableExists(t *testing.T, db *gorm.DB, name string) bool {
+	t.Helper()
+	var n int
+	err := db.Raw("SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?", name).Scan(&n).Error
+	require.NoError(t, err)
+	return n > 0
+}
+
+func testMigrationFS() fstest.MapFS {
+	return fstest.MapFS{
+		"migrations/sqlite/00001_init.sql": &fstest.MapFile{Data: []byte(`-- +goose Up
+CREATE TABLE t_up (id INTEGER PRIMARY KEY);
+
+-- +goose Down
+DROP TABLE t_up;
+`)},
+	}
+}
+
+func openMigrateTestDB(t *testing.T) *gorm.DB {
+	t.Helper()
+	dbPath := filepath.Join(t.TempDir(), "migrate.db")
+	gdb, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
+	require.NoError(t, err)
+	return gdb
+}
+
+func TestGooseEngineMigrateOrderCreateTableBaselineUp(t *testing.T) {
+	gdb := openMigrateTestDB(t)
+	var order []string
+
+	app := core.NewApp(
+		core.WithMigrationEngine(&gooseEngine{}),
+		core.WithMigrationBaseline(func(*core.Context) error {
+			require.True(t, sqliteTableExists(t, gdb, "w_schema_versions"), "version table must exist before baseline")
+			require.False(t, sqliteTableExists(t, gdb, "t_up"), "plugin Up must not run before baseline")
+			order = append(order, "create-table", "baseline")
+			return nil
+		}),
+		core.WithPlugins(&migrateTestPlugin{db: gdb, fs: testMigrationFS()}),
+	)
+
+	require.NoError(t, app.Prepare())
+	require.NoError(t, app.ApplyPlugins())
+	require.NoError(t, app.RunMigrations())
+
+	require.True(t, sqliteTableExists(t, gdb, "t_up"), "plugin Up must run after baseline")
+	order = append(order, "up")
+	assert.Equal(t, []string{"create-table", "baseline", "up"}, order)
+}
+
+func TestGooseEngineBaselineErrorSkipsUp(t *testing.T) {
+	gdb := openMigrateTestDB(t)
+
+	app := core.NewApp(
+		core.WithMigrationEngine(&gooseEngine{}),
+		core.WithMigrationBaseline(func(*core.Context) error {
+			require.True(t, sqliteTableExists(t, gdb, "w_schema_versions"), "version table must exist before baseline")
+			return assert.AnError
+		}),
+		core.WithPlugins(&migrateTestPlugin{db: gdb, fs: testMigrationFS()}),
+	)
+
+	require.NoError(t, app.Prepare())
+	require.NoError(t, app.ApplyPlugins())
+	err := app.RunMigrations()
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "migration baseline")
+	assert.False(t, sqliteTableExists(t, gdb, "t_up"), "plugin Up must not run when baseline fails")
+}
+
+func TestGooseEngineNilBaselineStillMigrates(t *testing.T) {
+	gdb := openMigrateTestDB(t)
+
+	app := core.NewApp(
+		core.WithMigrationEngine(&gooseEngine{}),
+		core.WithPlugins(&migrateTestPlugin{db: gdb, fs: testMigrationFS()}),
+	)
+
+	require.NoError(t, app.Prepare())
+	require.NoError(t, app.ApplyPlugins())
+	require.NoError(t, app.RunMigrations())
+	assert.True(t, sqliteTableExists(t, gdb, "w_schema_versions"))
+	assert.True(t, sqliteTableExists(t, gdb, "t_up"))
+}
@@ -59,6 +59,14 @@ func WithMigrationRunner(runner MigrationRunner) AppOption {
 	}
 }

+// WithMigrationBaseline registers a hook the migration engine runs after the
+// shared version table exists and before any plugin Up.
+func WithMigrationBaseline(fn func(*Context) error) AppOption {
+	return func(a *App) {
+		a.migrationBaseline = fn
+	}
+}
+
 // WithShutdownTimeout sets the fallback timeout for graceful application shutdown.
 func WithShutdownTimeout(timeout time.Duration) AppOption {
 	return func(a *App) {
@@ -99,23 +107,24 @@ func WithConfigDecl(pluginID string, bindings ...ConfigBinding) AppOption {
 // It manages plugin collection, dependency mounting, migration execution, profile-based driver startup,
 // and graceful signal-driven LIFO shutdown.
 type App struct {
-	mu               sync.RWMutex
-	ctx              *Context
-	profile          Profile
-	plugins          []Plugin
-	pluginMap        map[string]Plugin
-	fibers           []*Fiber
-	fiberMap         map[string]*Fiber
-	applied          bool
-	running          bool
-	startedDrivers   []Driver
-	migrationEngine  MigrationEngine
-	shutdownTimeout  time.Duration
-	configSource     ConfigSource
-	hostDeclOwner    string
-	hostDeclBindings []ConfigBinding
-	prepared         bool
-	applyErr         error
+	mu                sync.RWMutex
+	ctx               *Context
+	profile           Profile
+	plugins           []Plugin
+	pluginMap         map[string]Plugin
+	fibers            []*Fiber
+	fiberMap          map[string]*Fiber
+	applied           bool
+	running           bool
+	startedDrivers    []Driver
+	migrationEngine   MigrationEngine
+	migrationBaseline func(*Context) error
+	shutdownTimeout   time.Duration
+	configSource      ConfigSource
+	hostDeclOwner     string
+	hostDeclBindings  []ConfigBinding
+	prepared          bool
+	applyErr          error
 }

 // NewApp creates a new Cordis application instance with default options.
@@ -373,22 +382,20 @@ func (a *App) prepareLocked() error {
 	if a.prepared {
 		return nil
 	}
-	if a.configSource == nil {
-		a.prepared = true
-		return nil
+	if a.configSource != nil {
+		config := a.ctx.Config()
+		config.SetSource(a.configSource)
+
+		if err := config.Declare(a.hostDeclOwner, a.hostDeclBindings...); err != nil {
+			return err
+		}
+		if err := config.Resolve(); err != nil {
+			return err
+		}
 	}

-	config := a.ctx.Config()
-	config.SetSource(a.configSource)
-
-	if err := config.Declare(a.hostDeclOwner, a.hostDeclBindings...); err != nil {
-		return err
-	}
-	if err := config.Resolve(); err != nil {
-		return err
-	}
+	a.ctx.setMigrationBaseline(a.migrationBaseline)
 	a.prepared = true
-
 	return nil
 }

@@ -594,3 +594,61 @@ func TestAppSetShutdownTimeoutIgnoresNonPositive(t *testing.T) {
 	app.SetShutdownTimeout(45 * time.Second)
 	assert.Equal(t, 45*time.Second, app.ShutdownTimeout())
 }
+
+func TestWithMigrationBaselineVisibleAfterPrepare(t *testing.T) {
+	var called bool
+	fn := func(*core.Context) error {
+		called = true
+		return nil
+	}
+
+	app := core.NewApp(core.WithMigrationBaseline(fn))
+	require.Nil(t, app.Context().MigrationBaseline(), "baseline must be copied during Prepare")
+
+	require.NoError(t, app.Prepare())
+
+	got := app.Context().MigrationBaseline()
+	require.NotNil(t, got, "Prepare must copy the baseline onto the root Context")
+	require.NoError(t, got(app.Context()))
+	assert.True(t, called)
+}
+
+func TestWithMigrationBaselineRunsBeforeEngineMigrate(t *testing.T) {
+	var order []string
+	engine := core.MigrationRunner(func(ctx *core.Context, _ []extpoints.MigrationEntry) error {
+		order = append(order, "engine")
+		if ctx.MigrationBaseline() == nil {
+			t.Fatal("baseline must be visible on context inside Migrate")
+		}
+		return ctx.MigrationBaseline()(ctx)
+	})
+
+	sqlFS := fstest.MapFS{
+		"migrations/001_init.sql": &fstest.MapFile{Data: []byte("-- +goose Up\nSELECT 1;\n")},
+	}
+	app := core.NewApp(
+		core.WithMigrationEngine(engine),
+		core.WithMigrationBaseline(func(*core.Context) error {
+			order = append(order, "baseline")
+			return nil
+		}),
+		core.WithPlugins(&appMockPlugin{
+			name: "t",
+			applyFn: func(ctx *core.Context) error {
+				ctx.Migrations().Register("t", sqlFS)
+				return nil
+			},
+		}),
+	)
+
+	require.NoError(t, app.Start(context.Background()))
+	defer func() { _ = app.Stop(context.Background()) }()
+
+	assert.Equal(t, []string{"engine", "baseline"}, order)
+}
+
+func TestWithMigrationBaselineNilByDefault(t *testing.T) {
+	app := core.NewApp()
+	require.NoError(t, app.Prepare())
+	assert.Nil(t, app.Context().MigrationBaseline())
+}
@@ -29,12 +29,13 @@ type Context struct {
 	settings   extpoints.SettingExtension
 	config     extpoints.ConfigExtension

-	mu        sync.RWMutex
-	children  []*Context
-	disposers []Disposer
-	drivers   []Driver
-	values    map[any]any
-	disposed  bool
+	mu                sync.RWMutex
+	children          []*Context
+	disposers         []Disposer
+	drivers           []Driver
+	values            map[any]any
+	disposed          bool
+	migrationBaseline func(*Context) error
 }

 // NewContext creates a new root Context wrapping a standard Go context.
@@ -159,18 +160,19 @@ func (c *Context) ForkWithContext(base context.Context) *Context {
 	ctx, cancel := context.WithCancel(base)

 	child := &Context{
-		goCtx:      ctx,
-		cancel:     cancel,
-		parent:     c,
-		container:  NewContainer(c.container),
-		events:     c.events,
-		router:     c.router,
-		migrations: c.migrations,
-		tasks:      c.tasks,
-		schedules:  c.schedules,
-		settings:   c.settings,
-		config:     c.config,
-		values:     make(map[any]any),
+		goCtx:             ctx,
+		cancel:            cancel,
+		parent:            c,
+		container:         NewContainer(c.container),
+		events:            c.events,
+		router:            c.router,
+		migrations:        c.migrations,
+		tasks:             c.tasks,
+		schedules:         c.schedules,
+		settings:          c.settings,
+		config:            c.config,
+		values:            make(map[any]any),
+		migrationBaseline: c.MigrationBaseline(),
 	}

 	c.mu.Lock()
@@ -208,6 +210,33 @@ func (c *Context) Migrations() extpoints.MigrationExtension {
 	return c.migrations
 }

+// MigrationBaseline returns the hook copied onto this Context during App.Prepare.
+// Child contexts fall back to their parent so forks still see the root hook.
+func (c *Context) MigrationBaseline() func(*Context) error {
+	if c == nil {
+		return nil
+	}
+	c.mu.RLock()
+	fn := c.migrationBaseline
+	c.mu.RUnlock()
+	if fn != nil {
+		return fn
+	}
+	if c.parent != nil {
+		return c.parent.MigrationBaseline()
+	}
+	return nil
+}
+
+func (c *Context) setMigrationBaseline(fn func(*Context) error) {
+	if c == nil {
+		return
+	}
+	c.mu.Lock()
+	c.migrationBaseline = fn
+	c.mu.Unlock()
+}
+
 // Tasks returns the scoped TaskExtension registry with automatic disposer tracking.
 func (c *Context) Tasks() extpoints.TaskExtension {
 	return newScopedTaskExtension(c, c.tasks)