Golang快速创建PostgreSQL测试数据库的工具pgdbtemplate

pgdbtemplate 确实是一个解决PostgreSQL测试数据库创建性能问题的优秀方案。它通过PostgreSQL的模板数据库特性，将数据库创建时间从线性增长变为近乎恒定，这对大型测试套件来说是个巨大的性能提升。

下面是一个更完整的集成示例，展示了如何在并行测试中使用它：

package user_test

import (
    "context"
    "database/sql"
    "fmt"
    "os"
    "testing"

    "github.com/andrei-polukhin/pgdbtemplate"
    _ "github.com/lib/pq"
)

var templateManager *pgdbtemplate.TemplateManager

func TestMain(m *testing.M) {
    ctx := context.Background()
    
    // 配置连接字符串生成函数
    connStringFunc := func(dbName string) string {
        return fmt.Sprintf(
            "postgres://%s:%s@%s/%s?sslmode=disable",
            os.Getenv("PG_USER"),
            os.Getenv("PG_PASSWORD"),
            os.Getenv("PG_HOST"),
            dbName,
        )
    }
    
    // 创建连接提供者（支持连接池配置）
    provider := pgdbtemplate.NewStandardConnectionProvider(
        connStringFunc,
        pgdbtemplate.WithMaxOpenConns(20),
        pgdbtemplate.WithMaxIdleConns(5),
    )
    
    // 配置迁移运行器
    migrationRunner := pgdbtemplate.NewFileMigrationRunner(
        []string{"./migrations"},
        pgdbtemplate.AlphabeticalMigrationFilesSorting,
    )
    
    // 创建模板管理器
    tm, err := pgdbtemplate.NewTemplateManager(pgdbtemplate.Config{
        ConnectionProvider: provider,
        MigrationRunner:    migrationRunner,
        TemplateDBName:     "test_template_db", // 自定义模板数据库名
    })
    if err != nil {
        panic(fmt.Sprintf("Failed to create template manager: %v", err))
    }
    
    // 初始化模板数据库（一次性操作）
    if err := tm.Initialize(ctx); err != nil {
        panic(fmt.Sprintf("Failed to initialize template: %v", err))
    }
    
    templateManager = tm
    exitCode := m.Run()
    
    // 清理模板数据库
    if err := tm.Cleanup(ctx); err != nil {
        fmt.Printf("Warning: cleanup failed: %v\n", err)
    }
    
    os.Exit(exitCode)
}

func TestUserRepository_CreateUser(t *testing.T) {
    t.Parallel() // 支持并行测试
    
    ctx := context.Background()
    
    // 为每个测试创建独立的数据库
    testDB, testDBName, err := templateManager.CreateTestDatabase(ctx)
    if err != nil {
        t.Fatalf("Failed to create test database: %v", err)
    }
    defer testDB.Close()
    defer templateManager.DropTestDatabase(ctx, testDBName)
    
    // 使用新数据库进行测试
    repo := NewUserRepository(testDB)
    
    // 测试逻辑
    user, err := repo.CreateUser(ctx, "test@example.com")
    if err != nil {
        t.Errorf("CreateUser failed: %v", err)
    }
    
    if user.Email != "test@example.com" {
        t.Errorf("Expected email test@example.com, got %s", user.Email)
    }
}

func TestUserRepository_GetUser(t *testing.T) {
    t.Parallel()
    
    ctx := context.Background()
    
    testDB, testDBName, err := templateManager.CreateTestDatabase(ctx)
    if err != nil {
        t.Fatalf("Failed to create test database: %v", err)
    }
    defer testDB.Close()
    defer templateManager.DropTestDatabase(ctx, testDBName)
    
    repo := NewUserRepository(testDB)
    
    // 先创建用户
    createdUser, err := repo.CreateUser(ctx, "get@example.com")
    if err != nil {
        t.Fatalf("Setup failed: %v", err)
    }
    
    // 测试查询
    fetchedUser, err := repo.GetUser(ctx, createdUser.ID)
    if err != nil {
        t.Errorf("GetUser failed: %v", err)
    }
    
    if fetchedUser.Email != "get@example.com" {
        t.Errorf("Expected email get@example.com, got %s", fetchedUser.Email)
    }
}

// 用户仓库示例
type UserRepository struct {
    db *sql.DB
}

func NewUserRepository(db *sql.DB) *UserRepository {
    return &UserRepository{db: db}
}

func (r *UserRepository) CreateUser(ctx context.Context, email string) (*User, error) {
    var id int64
    err := r.db.QueryRowContext(ctx,
        "INSERT INTO users (email) VALUES ($1) RETURNING id",
        email,
    ).Scan(&id)
    
    if err != nil {
        return nil, err
    }
    
    return &User{
        ID:    id,
        Email: email,
    }, nil
}

func (r *UserRepository) GetUser(ctx context.Context, id int64) (*User, error) {
    var user User
    err := r.db.QueryRowContext(ctx,
        "SELECT id, email FROM users WHERE id = $1",
        id,
    ).Scan(&user.ID, &user.Email)
    
    if err != nil {
        return nil, err
    }
    
    return &user, nil
}

type User struct {
    ID    int64
    Email string
}

对于需要事务隔离的测试，可以这样使用：

func TestUserRepository_WithTransaction(t *testing.T) {
    ctx := context.Background()
    
    testDB, testDBName, err := templateManager.CreateTestDatabase(ctx)
    if err != nil {
        t.Fatal(err)
    }
    defer testDB.Close()
    defer templateManager.DropTestDatabase(ctx, testDBName)
    
    // 在事务中运行测试
    tx, err := testDB.BeginTx(ctx, nil)
    if err != nil {
        t.Fatal(err)
    }
    defer tx.Rollback()
    
    repo := NewUserRepository(testDB)
    
    // 使用事务进行测试
    user, err := repo.CreateUser(ctx, "tx@example.com")
    if err != nil {
        t.Errorf("CreateUser in transaction failed: %v", err)
    }
    
    // 验证数据
    var count int
    err = tx.QueryRowContext(ctx, "SELECT COUNT(*) FROM users WHERE email = $1", "tx@example.com").Scan(&count)
    if err != nil {
        t.Errorf("Query failed: %v", err)
    }
    
    if count != 1 {
        t.Errorf("Expected 1 user, got %d", count)
    }
}

这个库确实能显著提升测试性能，特别是在迁移复杂、测试数量多的情况下。通过模板数据库的方式，每个测试都能获得完全独立的数据库实例，同时避免了重复运行迁移的开销。