46 lines
1.1 KiB
Go
46 lines
1.1 KiB
Go
package database
|
||
|
||
import (
|
||
"github.com/gin-gonic/gin"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
// GetTenantDBFromContext 从 gin.Context 获取租户 DB
|
||
// 已弃用，请使用 GetDB
|
||
func GetTenantDBFromContext(aboutID int64) (*gorm.DB, error) {
|
||
return GetTenantDB(aboutID)
|
||
}
|
||
|
||
// GetDB 从 gin.Context 获取当前应使用的数据库连接
|
||
// 如果上下文中有租户DB则返回租户DB，否则返回主库
|
||
// 使用方式: db := database.GetDB(c)
|
||
func GetDB(c *gin.Context) *gorm.DB {
|
||
if c == nil {
|
||
return DB
|
||
}
|
||
if tenantDB, exists := c.Get("tenant_db"); exists {
|
||
if db, ok := tenantDB.(*gorm.DB); ok {
|
||
return db
|
||
}
|
||
}
|
||
return DB
|
||
}
|
||
|
||
// GetDBWithFallback 获取数据库连接，支持可选的 DB 参数覆盖
|
||
// 使用方式: db := database.GetDBWithFallback(c, dbOverride)
|
||
func GetDBWithFallback(c *gin.Context, dbOverride *gorm.DB) *gorm.DB {
|
||
if dbOverride != nil {
|
||
return dbOverride
|
||
}
|
||
return GetDB(c)
|
||
}
|
||
|
||
// OptionalDB 可选 DB 参数处理
|
||
// 使用方式: db := database.OptionalDB(db...)
|
||
func OptionalDB(db ...*gorm.DB) *gorm.DB {
|
||
if len(db) > 0 && db[0] != nil {
|
||
return db[0]
|
||
}
|
||
return DB
|
||
}
|