package db import ( "context" "crypto/sha256" "database/sql" "encoding/hex" "fmt" "net/url" "os" "path/filepath" "runtime" "strings" "sync" "testing" "github.com/jackc/pgx/v5/pgxpool" ) var ( migrateOnceMu sync.Mutex migratedDSNs = map[string]error{} ) // RunMigrationsOnceForTest 同一 DSN 只迁一次,避免并行包抢同一隔离库。 func RunMigrationsOnceForTest(ctx context.Context, pool *pgxpool.Pool, dsn string) error { key := strings.TrimSpace(dsn) migrateOnceMu.Lock() defer migrateOnceMu.Unlock() if err, ok := migratedDSNs[key]; ok { return err } err := RunMigrations(ctx, pool, ResolveMigrationsDir()) migratedDSNs[key] = err return err } const testDatabaseEnv = "TEST_DATABASE_URL" // RequireTestDatabaseURL 集成测试必须显式提供隔离测试库,禁止默认连开发库。 // 返回按调用方包派生的独立数据库(_):包内共享、包间隔离, // 消除 go test 并行包对同一库的迁移/种子互相干扰(2026-09-29 issue #2)。 func RequireTestDatabaseURL(t *testing.T) string { t.Helper() dsn := strings.TrimSpace(os.Getenv(testDatabaseEnv)) if dsn == "" { t.Fatalf("%s must be set for PostgreSQL integration tests (refusing default dev database)", testDatabaseEnv) } return derivedPackageDSN(t, dsn) } // SkipIfNoTestDatabaseURL 无隔离测试库时跳过(单元测试包内可选 PG 用例)。 // 同样返回按包派生库。 func SkipIfNoTestDatabaseURL(t *testing.T) string { t.Helper() dsn := strings.TrimSpace(os.Getenv(testDatabaseEnv)) if dsn == "" { t.Skipf("%s not set; skipping PostgreSQL integration test", testDatabaseEnv) } return derivedPackageDSN(t, dsn) } var ( pkgDBMu sync.Mutex pkgDBReady = map[string]string{} // 派生 DSN 缓存(进程内) ) func derivedPackageDSN(t *testing.T, baseDSN string) string { t.Helper() _, file, _, ok := runtime.Caller(2) // 跳过 helper 自身,取调用测试文件 if !ok { return baseDSN } pkgDir := filepath.ToSlash(filepath.Dir(file)) // 仅取模块内相对路径段,避免绝对路径差异影响哈希 if i := strings.Index(pkgDir, "internal/"); i >= 0 { pkgDir = pkgDir[i:] } else if i := strings.Index(pkgDir, "cmd/"); i >= 0 { pkgDir = pkgDir[i:] } sum := sha256.Sum256([]byte(pkgDir)) suffix := hex.EncodeToString(sum[:4]) u, err := url.Parse(baseDSN) if err != nil || u.Path == "" { return baseDSN // 非标准 DSN:退回原值,行为同旧版 } baseDB := strings.TrimPrefix(u.Path, "/") if baseDB == "" { return baseDSN } sanitized := strings.Map(func(r rune) rune { switch { case r >= 'a' && r <= 'z', r >= '0' && r <= '9': return r default: return '_' } }, strings.ToLower(baseDB)) name := fmt.Sprintf("%s_%s", sanitized, suffix) if len(name) > 63 { name = name[:63] } u2 := *u u2.Path = "/" + name pkgDBMu.Lock() cached, done := pkgDBReady[name] pkgDBMu.Unlock() if done { return cached } if err := ensureTestDatabase(baseDSN, name); err != nil { t.Fatalf("prepare per-package test database %s: %v", name, err) } dsn := u2.String() pkgDBMu.Lock() pkgDBReady[name] = dsn pkgDBMu.Unlock() return dsn } // ensureTestDatabase 在基库实例上幂等创建派生库(advisory lock 防并行包竞态)。 func ensureTestDatabase(baseDSN, name string) error { admindb, err := sql.Open("pgx", baseDSN) if err != nil { return err } defer admindb.Close() var exists bool if err := admindb.QueryRow(`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname=$1)`, name).Scan(&exists); err != nil { return err } if exists { return nil } if _, err := admindb.Exec(`SELECT pg_advisory_lock(hashtext($1))`, name); err != nil { return err } defer admindb.Exec(`SELECT pg_advisory_unlock(hashtext($1))`, name) if err := admindb.QueryRow(`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname=$1)`, name).Scan(&exists); err != nil { return err } if exists { return nil } _, err = admindb.Exec(fmt.Sprintf(`CREATE DATABASE %q`, name)) return err } // RefuseSharedEnvDSN fails if the DSN targets a shared env database. func RefuseSharedEnvDSN(t *testing.T, dsn string) { t.Helper() if strings.Contains(dsn, "5433") || strings.Contains(dsn, "agent_system_dev") || (strings.Contains(dsn, "5434") && strings.Contains(dsn, "agent_system_test")) || strings.Contains(dsn, "5435") { t.Fatal("TEST_DATABASE_URL points at a shared env port; refuse agent_system_dev/test/main") } }