[management] Clone test databases from a migrated template

Profiling the management unit suite per engine showed that nearly every test in management/server spends 1.0 to 1.7 seconds creating its store on MySQL and about 0.9 seconds on Postgres, against 0.3 seconds on sqlite. The work is CREATE DATABASE followed by 26 pre-migrations, a 40-table AutoMigrate and 9 post-migrations, repeated for each of the package's 133 store creation sites. That put the package at 10 minutes of its 20 minute go test budget on a healthy runner, and over it on a slow one. Run the full migration once per process into a template database and clone it for every test database. Postgres copies it server-side with CREATE DATABASE ... TEMPLATE. MySQL has no equivalent, so the template's SHOW CREATE TABLE statements are replayed into the new database with foreign key checks off, which keeps the ON DELETE CASCADE constraints that CREATE TABLE ... LIKE would drop. The per-test store then opens with migrations skipped. Every test still gets its own database, so isolation is unchanged.
2026-09-12 17:59:06 +02:00 · 2026-09-12 12:00:37 +00:00
parent ba9b226b4b
commit e9b9a154c1
2 changed files with 226 additions and 13 deletions
@@ -3213,7 +3213,11 @@ func NewSqliteStoreFromFileStore(ctx context.Context, fileStore *FileStore, data

 // NewPostgresqlStoreFromSqlStore restores a store from SqlStore and stores Postgres DB.
 func NewPostgresqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics) (*SqlStore, error) {
-	store, err := NewPostgresqlStoreForTests(ctx, dsn, metrics, false)
+	return newPostgresqlStoreFromSqlStore(ctx, sqliteStore, dsn, metrics, false)
+}
+
+func newPostgresqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics, skipMigration bool) (*SqlStore, error) {
+	store, err := NewPostgresqlStoreForTests(ctx, dsn, metrics, skipMigration)
 	if err != nil {
 		return nil, err
 	}
@@ -3279,7 +3283,11 @@ func connectToPgDbForTests(ctx context.Context, dsn string) (*pgxpool.Pool, erro

 // NewMysqlStoreFromSqlStore restores a store from SqlStore and stores MySQL DB.
 func NewMysqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics) (*SqlStore, error) {
-	store, err := NewMysqlStore(ctx, dsn, metrics, false)
+	return newMysqlStoreFromSqlStore(ctx, sqliteStore, dsn, metrics, false)
+}
+
+func newMysqlStoreFromSqlStore(ctx context.Context, sqliteStore *SqlStore, dsn string, metrics telemetry.AppMetrics, skipMigration bool) (*SqlStore, error) {
+	store, err := NewMysqlStore(ctx, dsn, metrics, skipMigration)
 	if err != nil {
 		return nil, err
 	}
@@ -4,6 +4,7 @@ package store

 import (
 	"context"
+	"database/sql"
 	"errors"
 	"fmt"
 	"net"
@@ -15,6 +16,7 @@ import (
 	"runtime"
 	"slices"
 	"strings"
+	"sync"
 	"time"

 	"github.com/google/uuid"
@@ -817,19 +819,23 @@ func newReusedPostgresStore(ctx context.Context, store *SqlStore, kind types.Eng
 		return nil, nil, fmt.Errorf("failed to open postgres connection: %v", err)
 	}

-	dsn, cleanup, err := createRandomDB(dsn, db, kind)
-
-	sqlDB, _ := db.DB()
-	if sqlDB != nil {
-		sqlDB.Close()
+	template, err := postgresSchemaTemplate(ctx, dsn, db)
+	if err != nil {
+		closeGormDB(db)
+		return nil, nil, err
 	}

+	dsn, cleanup, err := createRandomDB(dsn, db, kind, template)
+
+	closeGormDB(db)
+
 	if err != nil {
 		return nil, nil, err
 	}

-	store, err = NewPostgresqlStoreFromSqlStore(ctx, store, dsn, nil)
+	store, err = newPostgresqlStoreFromSqlStore(ctx, store, dsn, nil, true)
 	if err != nil {
+		cleanup()
 		return nil, nil, err
 	}

@@ -862,7 +868,13 @@ func newReusedMysqlStore(ctx context.Context, store *SqlStore, kind types.Engine
 	sqlDB.SetMaxOpenConns(1)
 	sqlDB.SetMaxIdleConns(1)

-	dsn, cleanup, err := createRandomDB(dsn, db, kind)
+	tableDDL, err := mysqlSchemaTemplate(ctx, dsn, db)
+	if err != nil {
+		sqlDB.Close()
+		return nil, nil, err
+	}
+
+	dsn, cleanup, err := createRandomDB(dsn, db, kind, "")

 	sqlDB.Close()

@@ -870,14 +882,187 @@ func newReusedMysqlStore(ctx context.Context, store *SqlStore, kind types.Engine
 		return nil, nil, err
 	}

-	store, err = NewMysqlStoreFromSqlStore(ctx, store, dsn, nil)
+	if err := cloneMysqlSchema(ctx, dsn, tableDDL); err != nil {
+		cleanup()
+		return nil, nil, err
+	}
+
+	store, err = newMysqlStoreFromSqlStore(ctx, store, dsn, nil, true)
 	if err != nil {
+		cleanup()
 		return nil, nil, err
 	}

 	return store, cleanup, nil
 }

+// schemaTemplates remembers, per engine and server, a database that went
+// through the full migration once in this process. Every later test database
+// is cloned from it, so a test pays for CREATE DATABASE and a schema copy
+// instead of the 40-table AutoMigrate plus every pre and post migration, which
+// is what made each MySQL test store cost well over a second in CI.
+var (
+	schemaTemplatesMu sync.Mutex
+	schemaTemplates   = map[string]*schemaTemplate{}
+)
+
+type schemaTemplate struct {
+	dbName string
+	// tableDDL holds the CREATE TABLE statements of the template. MySQL has no
+	// server-side database template, so the schema is replayed statement by
+	// statement into each test database.
+	tableDDL []string
+}
+
+func schemaTemplateKey(engine types.Engine, dsn string) string {
+	return string(engine) + "|" + dsn
+}
+
+func newTestDBName(prefix string) string {
+	return fmt.Sprintf("%s_%s", prefix, strings.ReplaceAll(uuid.New().String(), "-", "_"))
+}
+
+// postgresSchemaTemplate returns the name of a fully migrated database that
+// CREATE DATABASE ... TEMPLATE can copy, creating it on first use.
+func postgresSchemaTemplate(ctx context.Context, baseDSN string, admin *gorm.DB) (string, error) {
+	schemaTemplatesMu.Lock()
+	defer schemaTemplatesMu.Unlock()
+
+	key := schemaTemplateKey(types.PostgresStoreEngine, baseDSN)
+	if tpl, ok := schemaTemplates[key]; ok {
+		return tpl.dbName, nil
+	}
+
+	name := newTestDBName("test_template")
+	if err := admin.Exec(fmt.Sprintf("CREATE DATABASE %s", name)).Error; err != nil {
+		return "", fmt.Errorf("create postgres template database: %w", err)
+	}
+
+	tplStore, err := NewPostgresqlStoreForTests(ctx, replaceDBName(baseDSN, name), nil, false)
+	if err != nil {
+		return "", fmt.Errorf("migrate postgres template database: %w", err)
+	}
+	// TEMPLATE refuses a source that still has sessions, so release both handles
+	// before the first clone.
+	tplStore.Close(ctx)
+	if tplStore.pool != nil {
+		tplStore.pool.Close()
+	}
+
+	schemaTemplates[key] = &schemaTemplate{dbName: name}
+	return name, nil
+}
+
+// mysqlSchemaTemplate returns the CREATE TABLE statements of a fully migrated
+// database, migrating one on first use.
+func mysqlSchemaTemplate(ctx context.Context, baseDSN string, admin *gorm.DB) ([]string, error) {
+	schemaTemplatesMu.Lock()
+	defer schemaTemplatesMu.Unlock()
+
+	key := schemaTemplateKey(types.MysqlStoreEngine, baseDSN)
+	if tpl, ok := schemaTemplates[key]; ok {
+		return tpl.tableDDL, nil
+	}
+
+	name := newTestDBName("test_template")
+	if err := admin.Exec(fmt.Sprintf("CREATE DATABASE %s", name)).Error; err != nil {
+		return nil, fmt.Errorf("create mysql template database: %w", err)
+	}
+
+	tplStore, err := NewMysqlStore(ctx, replaceDBName(baseDSN, name), nil, false)
+	if err != nil {
+		return nil, fmt.Errorf("migrate mysql template database: %w", err)
+	}
+	tableDDL, err := mysqlTableDDL(ctx, tplStore.db, name)
+	tplStore.Close(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	schemaTemplates[key] = &schemaTemplate{dbName: name, tableDDL: tableDDL}
+	return tableDDL, nil
+}
+
+func mysqlTableDDL(ctx context.Context, db *gorm.DB, dbName string) ([]string, error) {
+	sqlDB, err := db.DB()
+	if err != nil {
+		return nil, err
+	}
+
+	tables, err := mysqlTableNames(ctx, sqlDB, dbName)
+	if err != nil {
+		return nil, err
+	}
+
+	tableDDL := make([]string, 0, len(tables))
+	for _, table := range tables {
+		var name, createStmt string
+		row := sqlDB.QueryRowContext(ctx, fmt.Sprintf("SHOW CREATE TABLE %s.%s", dbName, table))
+		if err := row.Scan(&name, &createStmt); err != nil {
+			return nil, fmt.Errorf("read create statement of %s: %w", table, err)
+		}
+		tableDDL = append(tableDDL, createStmt)
+	}
+	return tableDDL, nil
+}
+
+func mysqlTableNames(ctx context.Context, sqlDB *sql.DB, dbName string) ([]string, error) {
+	rows, err := sqlDB.QueryContext(ctx, fmt.Sprintf("SHOW TABLES FROM %s", dbName))
+	if err != nil {
+		return nil, fmt.Errorf("list template tables: %w", err)
+	}
+	defer rows.Close()
+
+	var tables []string
+	for rows.Next() {
+		var table string
+		if err := rows.Scan(&table); err != nil {
+			return nil, fmt.Errorf("scan template table name: %w", err)
+		}
+		tables = append(tables, table)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("list template tables: %w", err)
+	}
+	return tables, nil
+}
+
+// cloneMysqlSchema replays the template's CREATE TABLE statements into the
+// database the DSN points at.
+func cloneMysqlSchema(ctx context.Context, dsn string, tableDDL []string) error {
+	db, err := gorm.Open(mysql.Open(dsn+"?charset=utf8&parseTime=True&loc=Local"), getGormConfig())
+	if err != nil {
+		return fmt.Errorf("connect to test database: %w", err)
+	}
+	sqlDB, err := db.DB()
+	if err != nil {
+		return err
+	}
+	defer sqlDB.Close()
+
+	// The statements come out of SHOW TABLES in name order, not dependency
+	// order, and their foreign keys reference tables of the session's default
+	// database. Pin a single connection so the session setting below covers
+	// every statement, and connect straight to the new database so unqualified
+	// references land there.
+	sqlDB.SetMaxOpenConns(1)
+	if _, err := sqlDB.ExecContext(ctx, "SET FOREIGN_KEY_CHECKS = 0"); err != nil {
+		return fmt.Errorf("disable foreign key checks: %w", err)
+	}
+	for _, stmt := range tableDDL {
+		if _, err := sqlDB.ExecContext(ctx, stmt); err != nil {
+			return fmt.Errorf("replay table definition: %w", err)
+		}
+	}
+	return nil
+}
+
+func closeGormDB(db *gorm.DB) {
+	if sqlDB, _ := db.DB(); sqlDB != nil {
+		sqlDB.Close()
+	}
+}
+
 func openDBWithRetry(dsn string, engine types.Engine, maxRetries int) (*gorm.DB, error) {
 	var db *gorm.DB
 	var err error
@@ -903,10 +1088,16 @@ func openDBWithRetry(dsn string, engine types.Engine, maxRetries int) (*gorm.DB,
 	return nil, err
 }

-func createRandomDB(dsn string, db *gorm.DB, engine types.Engine) (string, func(), error) {
-	dbName := fmt.Sprintf("test_db_%s", strings.ReplaceAll(uuid.New().String(), "-", "_"))
+// createRandomDB creates a uniquely named database for one test. On postgres a
+// non-empty template is copied server-side with CREATE DATABASE ... TEMPLATE.
+func createRandomDB(dsn string, db *gorm.DB, engine types.Engine, template string) (string, func(), error) {
+	dbName := newTestDBName("test_db")

-	if err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", dbName)).Error; err != nil {
+	createStmt := fmt.Sprintf("CREATE DATABASE %s", dbName)
+	if template != "" && engine == types.PostgresStoreEngine {
+		createStmt = fmt.Sprintf("CREATE DATABASE %s TEMPLATE %s", dbName, template)
+	}
+	if err := execWithTemplateRetry(db, createStmt); err != nil {
 		return "", nil, fmt.Errorf("failed to create database: %v", err)
 	}

@@ -972,6 +1163,20 @@ func createRandomDB(dsn string, db *gorm.DB, engine types.Engine) (string, func(
 	return replaceDBName(dsn, dbName), cleanup, nil
 }

+// execWithTemplateRetry runs a statement, retrying briefly when postgres still
+// sees the template's just-closed sessions and refuses to copy it.
+func execWithTemplateRetry(db *gorm.DB, stmt string) error {
+	var err error
+	for attempt := 0; attempt < 20; attempt++ {
+		err = db.Exec(stmt).Error
+		if err == nil || !strings.Contains(err.Error(), "is being accessed by other users") {
+			return err
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+	return err
+}
+
 func replaceDBName(dsn, newDBName string) string {
 	re := regexp.MustCompile(`(?P<pre>[:/@])(?P<dbname>[^/?]+)(?P<post>\?|$)`)
 	return re.ReplaceAllString(dsn, `${pre}`+newDBName+`${post}`)