Use postgres template to create databases faster

Author: domdomegg

internal/database/testutil.go

@@ -4,39 +4,85 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"testing"
 	"time"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/stretchr/testify/require"
 )
 
-// NewTestDB creates an isolated PostgreSQL database for each test.
-// Each test gets its own database with a random name, eliminating all coordination issues.
+const templateDBName = "mcp_registry_test_template"
+
+// ensureTemplateDB creates a template database with migrations applied
+// Multiple processes may call this, so we handle race conditions
+func ensureTemplateDB(ctx context.Context, adminConn *pgx.Conn) error {
+	// Check if template exists
+	var exists bool
+	err := adminConn.QueryRow(ctx, "SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)", templateDBName).Scan(&exists)
+	if err != nil {
+		return fmt.Errorf("failed to check template database: %w", err)
+	}
+
+	if exists {
+		// Template already exists
+		return nil
+	}
+
+	// Create template database
+	_, err = adminConn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s", templateDBName))
+	if err != nil {
+		// Ignore duplicate error - another process created it concurrently
+		var pgErr *pgconn.PgError
+		if errors.As(err, &pgErr) && pgErr.Code == "23505" {
+			return nil
+		}
+		return fmt.Errorf("failed to create template database: %w", err)
+	}
+
+	// Connect to template and run migrations
+	templateURI := fmt.Sprintf("postgres://mcpregistry:mcpregistry@localhost:5432/%s?sslmode=disable", templateDBName)
+	templateDB, err := NewPostgreSQL(ctx, templateURI)
+	if err != nil {
+		return fmt.Errorf("failed to connect to template database: %w", err)
+	}
+	defer templateDB.Close()
+
+	// Migrations run automatically in NewPostgreSQL
+	return nil
+}
+
+// NewTestDB creates an isolated PostgreSQL database for each test by copying a template.
+// The template database has migrations pre-applied, so each test is fast.
 // Requires PostgreSQL to be running on localhost:5432 (e.g., via docker-compose).
 func NewTestDB(t *testing.T) Database {
 	t.Helper()
 
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
+	// Connect to postgres database
+	adminURI := "postgres://mcpregistry:mcpregistry@localhost:5432/postgres?sslmode=disable"
+	adminConn, err := pgx.Connect(ctx, adminURI)
+	require.NoError(t, err, "Failed to connect to PostgreSQL. Make sure PostgreSQL is running via: docker-compose up -d postgres")
+	defer adminConn.Close(ctx)
+
+	// Ensure template database exists with migrations
+	err = ensureTemplateDB(ctx, adminConn)
+	require.NoError(t, err, "Failed to initialize template database")
+
 	// Generate unique database name for this test
 	var randomBytes [8]byte
-	_, err := rand.Read(randomBytes[:])
+	_, err = rand.Read(randomBytes[:])
 	require.NoError(t, err, "Failed to generate random database id")
 	randomInt := binary.BigEndian.Uint64(randomBytes[:])
 	dbName := fmt.Sprintf("test_%d", randomInt)
 
-	// Connect to postgres database to create test database
-	adminURI := "postgres://mcpregistry:mcpregistry@localhost:5432/postgres?sslmode=disable"
-	adminConn, err := pgx.Connect(ctx, adminURI)
-	require.NoError(t, err, "Failed to connect to PostgreSQL. Make sure PostgreSQL is running via: docker-compose up -d postgres")
-	defer adminConn.Close(ctx)
-
-	// Create test database
-	_, err = adminConn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s", dbName))
-	require.NoError(t, err, "Failed to create test database")
+	// Create test database from template (fast - just copies files)
+	_, err = adminConn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s TEMPLATE %s", dbName, templateDBName))
+	require.NoError(t, err, "Failed to create test database from template")
 
 	// Register cleanup to drop database
 	t.Cleanup(func() {
@@ -53,8 +99,9 @@ func NewTestDB(t *testing.T) Database {
 		_, _ = adminConn.Exec(cleanupCtx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", dbName))
 	})
 
-	// Connect to test database
+	// Connect to test database (no migrations needed - copied from template)
 	testURI := fmt.Sprintf("postgres://mcpregistry:mcpregistry@localhost:5432/%s?sslmode=disable", dbName)
+
 	db, err := NewPostgreSQL(ctx, testURI)
 	require.NoError(t, err, "Failed to connect to test database")