// 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()
}