Files
hfb_sys/backend/internal/database/testdb.go
T
yml2213andCursor 647fdc0b4a feat(distribution): 通道费改主站承担,并按冷却期自动释放结算
利润不再扣支付通道费;站点可配置结算冷却期,到期自动入账,已释放后退款从收款站钱包回冲。同时修复当日自动批次无法追加、三方回冲打错钱包,以及测试库自增重置被 MySQL 统计缓存跳过的问题。

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-10 17:11:58 +08:00

547 lines
15 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package database
import (
"context"
"crypto/sha256"
"database/sql"
"encoding/hex"
"errors"
"fmt"
"os"
"path/filepath"
"regexp"
"runtime"
"strings"
"sync"
"time"
mysqldriver "github.com/go-sql-driver/mysql"
"gorm.io/driver/mysql"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
// DefaultTestDSN 指向本地 Docker 开发环境的 MySQL 测试库。
// 可通过 HFB_TEST_DSN 覆盖,库名必须包含 test 或 e2e。
const DefaultTestDSN = "hfb:secret@tcp(127.0.0.1:13306)/hfb_sys_test?charset=utf8mb4&parseTime=True&loc=Local"
const (
schemaMetaTable = "_test_schema_meta"
templateSuffix = "__template"
)
var (
testPrepareOnce sync.Once
testConn *gorm.DB
testPrepareErr error
)
// NewTestDB 返回指向 MySQL 测试库的干净连接。
//
// 表结构来自 backend/migrations 的真实迁移,与生产一致。
// 迁移只在模板库上执行一次;各测试包克隆模板库,指纹未变则跨 `go test`
// 调用复用,避免每个包、每次运行都 DROP + 重放 68 条迁移。
// 每次调用仍会清空业务表,使测试从空库开始。
func NewTestDB() *gorm.DB {
db, err := sharedTestDB()
if err != nil {
panic(fmt.Sprintf("初始化 MySQL 测试库失败: %v", err))
}
if err := truncateAllTables(db); err != nil {
panic(fmt.Sprintf("清空 MySQL 测试库失败: %v", err))
}
return db
}
func sharedTestDB() (*gorm.DB, error) {
testPrepareOnce.Do(func() {
testConn, testPrepareErr = openTestDatabase()
})
return testConn, testPrepareErr
}
// NewDryRunDB 返回 MySQL 方言的 DryRun 连接:只生成 SQL 不执行,也无需数据库运行。
// 用于断言查询语句形态的测试。
func NewDryRunDB() *gorm.DB {
db, err := gorm.Open(mysql.New(mysql.Config{
DSN: DefaultTestDSN,
SkipInitializeWithVersion: true,
}), &gorm.Config{
DryRun: true,
DisableAutomaticPing: true,
Logger: logger.Default.LogMode(logger.Silent),
})
if err != nil {
panic(fmt.Sprintf("初始化 MySQL DryRun 连接失败: %v", err))
}
return db
}
// AutoMigrateIfMissing 只对尚不存在的表执行 AutoMigrate。
//
// 测试库的表结构由 migrations 建立,直接 AutoMigrate 会去改迁移定义的
// 索引与约束(例如 DROP 掉被外键依赖的唯一索引),因此已存在的表一律跳过,
// 仅保留给测试专用的临时表使用。
func AutoMigrateIfMissing(db *gorm.DB, models ...any) error {
for _, target := range models {
if db.Migrator().HasTable(target) {
continue
}
if err := db.AutoMigrate(target); err != nil {
return err
}
}
return nil
}
func openTestDatabase() (*gorm.DB, error) {
dsn := strings.TrimSpace(os.Getenv("HFB_TEST_DSN"))
if dsn == "" {
dsn = DefaultTestDSN
}
cfg, err := mysqldriver.ParseDSN(dsn)
if err != nil {
return nil, fmt.Errorf("解析测试库 DSN 失败: %w", err)
}
baseName := cfg.DBName
templateName := baseName + templateSuffix
packageName := baseName + "_" + packageSuffix()
if err := assertTestDatabaseName(templateName); err != nil {
return nil, err
}
if err := assertTestDatabaseName(packageName); err != nil {
return nil, err
}
dir, err := migrationDir()
if err != nil {
return nil, err
}
fingerprint, err := migrationFingerprint(dir)
if err != nil {
return nil, err
}
serverCfg := cfg.Clone()
serverCfg.DBName = ""
serverCfg.MultiStatements = false
serverDB, err := sql.Open("mysql", serverCfg.FormatDSN())
if err != nil {
return nil, fmt.Errorf("连接 MySQL 服务失败: %w", err)
}
defer func() { _ = serverDB.Close() }()
if err := serverDB.Ping(); err != nil {
return nil, fmt.Errorf("连接 MySQL 服务失败: %w", err)
}
if err := ensureTemplate(serverDB, templateName, dir, fingerprint); err != nil {
return nil, err
}
if err := ensurePackageClone(serverDB, templateName, packageName, fingerprint); err != nil {
return nil, err
}
cfg.DBName = packageName
cfg.ParseTime = true
cfg.MultiStatements = true
sqlDB, err := sql.Open("mysql", cfg.FormatDSN())
if err != nil {
return nil, fmt.Errorf("连接测试库失败: %w", err)
}
sqlDB.SetMaxOpenConns(8)
sqlDB.SetMaxIdleConns(8)
sqlDB.SetConnMaxLifetime(30 * time.Minute)
gormDB, err := gorm.Open(mysql.New(mysql.Config{Conn: sqlDB}), &gorm.Config{
Logger: logger.Default.LogMode(logger.Silent),
})
if err != nil {
_ = sqlDB.Close()
return nil, fmt.Errorf("打开 gorm 连接失败: %w", err)
}
return gormDB, nil
}
func ensureTemplate(serverDB *sql.DB, name, dir, fingerprint string) error {
return withMySQLLock(serverDB, mysqlLockName("tpl", name), func() error {
if readFingerprint(serverDB, name) == fingerprint {
return nil
}
if _, err := serverDB.Exec("DROP DATABASE IF EXISTS `" + name + "`"); err != nil {
return fmt.Errorf("删除模板库失败: %w", err)
}
if _, err := serverDB.Exec("CREATE DATABASE `" + name + "` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"); err != nil {
return fmt.Errorf("创建模板库失败: %w", err)
}
if err := applyMigrationsTo(name, dir); err != nil {
return err
}
return writeFingerprint(serverDB, name, fingerprint)
})
}
func ensurePackageClone(serverDB *sql.DB, templateName, packageName, fingerprint string) error {
return withMySQLLock(serverDB, mysqlLockName("pkg", packageName), func() error {
if readFingerprint(serverDB, packageName) == fingerprint {
return nil
}
return cloneDatabase(serverDB, templateName, packageName)
})
}
func applyMigrationsTo(dbName, dir string) error {
cfg, err := parseTestDSN()
if err != nil {
return err
}
cfg.DBName = dbName
cfg.ParseTime = true
cfg.MultiStatements = true
db, err := sql.Open("mysql", cfg.FormatDSN())
if err != nil {
return fmt.Errorf("连接模板库失败: %w", err)
}
defer func() { _ = db.Close() }()
return applyMigrations(db, dir)
}
func parseTestDSN() (*mysqldriver.Config, error) {
rawDSN := strings.TrimSpace(os.Getenv("HFB_TEST_DSN"))
if rawDSN == "" {
rawDSN = DefaultTestDSN
}
cfg, err := mysqldriver.ParseDSN(rawDSN)
if err != nil {
return nil, fmt.Errorf("解析测试库 DSN 失败: %w", err)
}
return cfg, nil
}
func applyMigrations(db *sql.DB, dir string) error {
entries, err := os.ReadDir(dir)
if err != nil {
return fmt.Errorf("读取迁移目录失败: %w", err)
}
for _, entry := range entries {
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".sql") {
continue
}
raw, err := os.ReadFile(filepath.Join(dir, entry.Name()))
if err != nil {
return fmt.Errorf("读取迁移 %s 失败: %w", entry.Name(), err)
}
if _, err := db.Exec(gooseUpSQL(string(raw))); err != nil {
return fmt.Errorf("执行迁移 %s 失败: %w", entry.Name(), err)
}
}
return nil
}
func cloneDatabase(serverDB *sql.DB, src, dest string) error {
if _, err := serverDB.Exec("DROP DATABASE IF EXISTS `" + dest + "`"); err != nil {
return fmt.Errorf("删除测试库 %s 失败: %w", dest, err)
}
if _, err := serverDB.Exec("CREATE DATABASE `" + dest + "` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"); err != nil {
return fmt.Errorf("创建测试库 %s 失败: %w", dest, err)
}
tables, err := listBaseTables(serverDB, src)
if err != nil {
return err
}
cfg, err := parseTestDSN()
if err != nil {
return err
}
cfg.DBName = dest
cfg.MultiStatements = false
destDB, err := sql.Open("mysql", cfg.FormatDSN())
if err != nil {
return fmt.Errorf("连接测试库 %s 失败: %w", dest, err)
}
defer func() { _ = destDB.Close() }()
ctx := context.Background()
conn, err := destDB.Conn(ctx)
if err != nil {
return err
}
defer func() { _ = conn.Close() }()
if _, err := conn.ExecContext(ctx, "SET FOREIGN_KEY_CHECKS = 0"); err != nil {
return err
}
defer func() { _, _ = conn.ExecContext(ctx, "SET FOREIGN_KEY_CHECKS = 1") }()
for _, table := range tables {
var name, create string
if err := serverDB.QueryRow("SHOW CREATE TABLE `"+src+"`.`"+table+"`").Scan(&name, &create); err != nil {
return fmt.Errorf("读取表结构 %s.%s 失败: %w", src, table, err)
}
if _, err := conn.ExecContext(ctx, create); err != nil {
return fmt.Errorf("克隆表 %s 失败: %w", table, err)
}
if _, err := conn.ExecContext(ctx, "INSERT INTO `"+table+"` SELECT * FROM `"+src+"`.`"+table+"`"); err != nil {
return fmt.Errorf("克隆表数据 %s 失败: %w", table, err)
}
}
return nil
}
func listBaseTables(db *sql.DB, schema string) ([]string, error) {
rows, err := db.Query(
"SELECT table_name FROM information_schema.tables "+
"WHERE table_schema = ? AND table_type = 'BASE TABLE' ORDER BY table_name",
schema,
)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var tables []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
tables = append(tables, name)
}
return tables, rows.Err()
}
func writeFingerprint(db *sql.DB, schema, fingerprint string) error {
ddl := "CREATE TABLE IF NOT EXISTS `" + schema + "`.`" + schemaMetaTable + "` (" +
"k VARCHAR(64) NOT NULL PRIMARY KEY," +
"v VARCHAR(128) NOT NULL" +
") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci"
if _, err := db.Exec(ddl); err != nil {
return fmt.Errorf("创建模板指纹表失败: %w", err)
}
if _, err := db.Exec(
"REPLACE INTO `"+schema+"`.`"+schemaMetaTable+"` (k, v) VALUES (?, ?)",
"fingerprint", fingerprint,
); err != nil {
return fmt.Errorf("写入模板指纹失败: %w", err)
}
return nil
}
func readFingerprint(db *sql.DB, schema string) string {
var value string
err := db.QueryRow(
"SELECT v FROM `" + schema + "`.`" + schemaMetaTable + "` WHERE k = 'fingerprint'",
).Scan(&value)
if err != nil {
return ""
}
return value
}
func withMySQLLock(db *sql.DB, name string, fn func() error) error {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
conn, err := db.Conn(ctx)
if err != nil {
return err
}
defer func() { _ = conn.Close() }()
var got int
if err := conn.QueryRowContext(ctx, "SELECT GET_LOCK(?, 120)", name).Scan(&got); err != nil {
return fmt.Errorf("GET_LOCK %s 失败: %w", name, err)
}
if got != 1 {
return fmt.Errorf("GET_LOCK %s 未拿到锁", name)
}
defer func() {
_, _ = conn.ExecContext(context.Background(), "SELECT RELEASE_LOCK(?)", name)
}()
return fn()
}
func mysqlLockName(kind, dbName string) string {
raw := "hfb:" + kind + ":" + dbName
if len(raw) <= 64 {
return raw
}
sum := sha256.Sum256([]byte(raw))
return ("hfb:" + kind + ":") + hex.EncodeToString(sum[:12])
}
func migrationFingerprint(dir string) (string, error) {
entries, err := os.ReadDir(dir)
if err != nil {
return "", fmt.Errorf("读取迁移目录失败: %w", err)
}
h := sha256.New()
for _, entry := range entries {
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".sql") {
continue
}
_, _ = fmt.Fprintf(h, "%s\n", entry.Name())
raw, err := os.ReadFile(filepath.Join(dir, entry.Name()))
if err != nil {
return "", fmt.Errorf("读取迁移 %s 失败: %w", entry.Name(), err)
}
_, _ = h.Write(raw)
}
return hex.EncodeToString(h.Sum(nil)), nil
}
// assertTestDatabaseName 拒绝重置非测试库,避免误删生产数据。
func assertTestDatabaseName(name string) error {
if name == "" {
return errors.New("测试库 DSN 必须包含数据库名")
}
if !regexp.MustCompile(`^[A-Za-z0-9_]+$`).MatchString(name) {
return fmt.Errorf("测试库名只允许字母、数字、下划线: %q", name)
}
if len(name) > 64 {
return fmt.Errorf("测试库名超过 MySQL 上限 64 字符: %q", name)
}
lower := strings.ToLower(name)
if !strings.Contains(lower, "test") && !strings.Contains(lower, "e2e") {
return fmt.Errorf("拒绝重置非测试库 %q,库名需包含 test 或 e2e", name)
}
return nil
}
// packageSuffix 用当前测试进程的工作目录名(即被测包目录)作为库名后缀。
// go test 会把工作目录设为包的源目录,因此同包共享一个库、不同包互相隔离。
func packageSuffix() string {
wd, err := os.Getwd()
if err != nil {
return "unknown"
}
suffix := regexp.MustCompile(`[^A-Za-z0-9_]`).ReplaceAllString(filepath.Base(wd), "_")
if suffix == "" {
return "unknown"
}
return suffix
}
func migrationDir() (string, error) {
_, currentFile, _, ok := runtime.Caller(0)
if !ok {
return "", errors.New("无法定位 database 包路径")
}
return filepath.Join(filepath.Dir(currentFile), "..", "..", "migrations"), nil
}
// gooseUpSQL 截取 goose 迁移文件的 Up 段,Down 段与 goose 指令行原样保留,
// 后者在 MySQL 中属于注释,不影响执行。
func gooseUpSQL(raw string) string {
if i := strings.Index(raw, "-- +goose Up"); i >= 0 {
raw = raw[i+len("-- +goose Up"):]
}
if i := strings.Index(raw, "-- +goose Down"); i >= 0 {
raw = raw[:i]
}
return raw
}
// truncateAllTables 清空测试库中的全部业务表,使每个测试从空库开始。
// 指纹表保留,以便跨进程复用已克隆的库。
//
// 先用一条多语句 DELETE 清数据(测试数据量很小,通常几十毫秒),再把用过的
// AUTO_INCREMENT 表打回 1。全量 TRUNCATE 61 张表是 DDL,单次约 0.6s–1.4s,
// 而不少测试硬编码了 id=1,所以自增必须重置。
func truncateAllTables(db *gorm.DB) error {
sqlDB, err := db.DB()
if err != nil {
return err
}
tables, err := listCurrentDatabaseTables(sqlDB)
if err != nil {
return err
}
var b strings.Builder
b.WriteString("SET FOREIGN_KEY_CHECKS=0;")
count := 0
for _, table := range tables {
if table == schemaMetaTable {
continue
}
b.WriteString("DELETE FROM `")
b.WriteString(table)
b.WriteString("`;")
count++
}
if count == 0 {
return nil
}
b.WriteString("SET FOREIGN_KEY_CHECKS=1;")
if _, err := sqlDB.Exec(b.String()); err != nil {
return fmt.Errorf("清空测试表失败: %w", err)
}
aiTables, err := listAutoIncrementTables(sqlDB)
if err != nil {
return err
}
if len(aiTables) == 0 {
return nil
}
b.Reset()
for _, table := range aiTables {
b.WriteString("ALTER TABLE `")
b.WriteString(table)
b.WriteString("` AUTO_INCREMENT=1;")
}
if _, err := sqlDB.Exec(b.String()); err != nil {
return fmt.Errorf("重置自增失败: %w", err)
}
return nil
}
func listAutoIncrementTables(db *sql.DB) ([]string, error) {
ctx := context.Background()
conn, err := db.Conn(ctx)
if err != nil {
return nil, err
}
defer func() { _ = conn.Close() }()
// MySQL 8 默认缓存 information_schema 的统计信息(86400s),
// 直接读取 auto_increment 可能拿到陈旧值并跳过重置,必须在本连接关闭缓存后再查。
if _, err := conn.ExecContext(ctx, "SET SESSION information_schema_stats_expiry = 0"); err != nil {
return nil, err
}
rows, err := conn.QueryContext(ctx,
"SELECT table_name FROM information_schema.tables "+
"WHERE table_schema = DATABASE() AND table_name <> ? AND auto_increment > 1",
schemaMetaTable,
)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var tables []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
tables = append(tables, name)
}
return tables, rows.Err()
}
func listCurrentDatabaseTables(db *sql.DB) ([]string, error) {
rows, err := db.Query(
"SELECT table_name FROM information_schema.tables " +
"WHERE table_schema = DATABASE() AND table_type = 'BASE TABLE'",
)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var tables []string
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
return nil, err
}
tables = append(tables, name)
}
return tables, rows.Err()
}