运行时数据库切换

// Package database 提供可在运行时热切换的数据库句柄。
//
// 如果进程只在启动时创建一次 *gorm.DB，那么修改 config.yml 里的 database 段就
// 必须重启才能生效：GORM 的方言（占位符 ? 与 $1、SQL 生成、DataTypeOf）在
// gorm.Open 时就固化进了 *gorm.DB 的 Config，无法在不重建句柄的情况下换成另一种
// 驱动，也无法通过替换 ConnPool 实现（占位符在 SQL 组装阶段就已经写死）。
//
// Handle 把「当前生效的句柄」放进一个原子指针，所有 repository 只持有 Handle，
// 在发起语句时才解析出当时的 *gorm.DB。于是切换数据库退化成一次原子写：
// 已经在途的语句继续走旧句柄（旧连接池会等它们结束后再关闭），新语句自动落到
// 新库，全程不需要重启进程。
package database

import (
	"context"
	"sync"
	"sync/atomic"

	"gorm.io/gorm"
)

// Handle 持有当前生效的 *gorm.DB，并允许在运行时原子替换。
//
// Handle 内部含 atomic.Pointer 与 sync.Mutex，不能被复制（go vet 会告警），
// 因此所有使用方都必须通过 *Handle 传递。
type Handle struct {
	current atomic.Pointer[gorm.DB]
	swapMu  sync.Mutex
	onSwap  []func(old, next *gorm.DB)
}

// NewHandle 用初始句柄创建 Handle。
// 也可以用事务句柄创建临时 Handle，例如 repository 内部用 tx 派生子 repository。
func NewHandle(db *gorm.DB) *Handle {
	h := &Handle{}
	h.current.Store(db)
	return h
}

// Current 返回当前生效的句柄。
//
// 一次操作里需要多次使用同一个句柄时（例如先查再改），应当先取出来复用，
// 避免中途发生切换导致两次操作落在不同的库上。
func (h *Handle) Current() *gorm.DB {
	return h.current.Load()
}

// WithContext 是 repository 访问数据库的统一入口。
//
// 每条语句在发起时锁定一个句柄，之后的链式调用与事务都在同一个句柄上进行，
// 因此不会出现「SQL 按旧方言生成、却在新库执行」的错配。
func (h *Handle) WithContext(ctx context.Context) *gorm.DB {
	return h.current.Load().WithContext(ctx)
}

// AutoMigrate 用当前句柄建表。
func (h *Handle) AutoMigrate(dst ...any) error {
	return h.current.Load().AutoMigrate(dst...)
}

// OnSwap 注册切换回调。
//
// 给缓存了数据库状态的组件（例如把路由表读进内存的网关、启动时播种的内置
// Profile）一个在新库上重建内存态的机会，否则切库后这些数据会「凭空消失」。
func (h *Handle) OnSwap(fn func(old, next *gorm.DB)) {
	if fn == nil {
		return
	}
	h.swapMu.Lock()
	defer h.swapMu.Unlock()
	h.onSwap = append(h.onSwap, fn)
}

// Swap 原子切换到新句柄，并在后台关闭旧连接池。
//
// sql.DB.Close 会阻止新查询开始，并等待已经在服务端执行的查询结束，因此放在
// goroutine 里关闭既不会阻塞调用方，也不会打断在途语句。切换期间新旧两套连接池
// 短暂共存，回收后不再占用连接。
func (h *Handle) Swap(next *gorm.DB) {
	if next == nil {
		return
	}

	h.swapMu.Lock()
	old := h.current.Swap(next)
	callbacks := make([]func(old, next *gorm.DB), len(h.onSwap))
	copy(callbacks, h.onSwap)
	h.swapMu.Unlock()

	for _, cb := range callbacks {
		cb(old, next)
	}

	if old != nil && old != next {
		go func() { _ = Close(old) }()
	}
}

// Close 关闭句柄底层的连接池。
func Close(db *gorm.DB) error {
	if db == nil {
		return nil
	}
	sqlDB, err := db.DB()
	if err != nil {
		return err
	}
	return sqlDB.Close()
}