高级特性
第十一章:高级特性
11.1 上下文(Context)支持
GORM 完全支持 Go 的 context.Context,可用于超时控制、链路追踪等。
超时控制
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// 所有数据库操作都使用上下文
db.WithContext(ctx).Find(&users)
db.WithContext(ctx).Create(&user)
db.WithContext(ctx).Update("name", "张三")
传递元数据
// 在 context 中存储当前用户ID
ctx := context.WithValue(context.Background(), "userID", currentUserID)
// 在钩子中使用
func (u *User) BeforeCreate(tx *gorm.DB) error {
if userID := tx.Statement.Context.Value("userID"); userID != nil {
// 记录创建人
}
return nil
}
链路追踪
import "go.opentelemetry.io/otel"
func TraceMiddleware() func(*gorm.DB) {
return func(db *gorm.DB) {
ctx := db.Statement.Context
tracer := otel.Tracer("gorm")
ctx, span := tracer.Start(ctx, db.Statement.Table+"."+string(db.Statement.ReflectValue.Type().Name()))
defer span.End()
db.Statement.Context = ctx
}
}
// 注册
db.Callback().Create().Before("gorm:create").Register("trace", TraceMiddleware())
11.2 复合主键
type ProductCategory struct {
ProductID uint `gorm:"primaryKey"`
CategoryID uint `gorm:"primaryKey"`
SortOrder int
}
// 使用
db.First(&ProductCategory{}, "product_id = ? AND category_id = ?", 1, 2)
// 批量查询
db.Find(&productCategories, []ProductCategory{
{ProductID: 1, CategoryID: 1},
{ProductID: 1, CategoryID: 2},
})
11.3 多数据库支持
动态切换
// 根据上下文切换数据库
type DBKey string
func (d *DBResolver) Resolve(name string) *gorm.DB {
switch name {
case "order":
return d.OrderDB
case "user":
return d.UserDB
default:
return d.DefaultDB
}
}
// 使用
dbResolver.Resolve("order").Create(&order)
分库分表
// 根据用户 ID 分表
func GetUserTableName(userID uint64) string {
return fmt.Sprintf("user_%d", userID%10)
}
// 动态表名
type User struct {
ID uint64
Name string
}
func (User) TableName() string {
// 动态获取表名
return "user_0" // 实际通过上下文获取
}
// 使用 Scope
db.Scopes(func(db *gorm.DB) *gorm.DB {
return db.Table(fmt.Sprintf("user_%d", userID%10))
}).First(&user, userID)
11.4 SQL 生成器
ToSQL
// 生成 SQL 但不执行
stmt := db.Session(&gorm.Session{DryRun: true}).First(&user, 1).Statement
sql := stmt.SQL.String()
// SELECT * FROM `users` WHERE `users`.`id` = ? ORDER BY `users`.`id` LIMIT 1
执行前修改 SQL
db.Callback().Query().Before("gorm:query").Register("modify_sql", func(db *gorm.DB) {
// 修改 SQL 语句
if db.Statement.SQL.Len() > 0 {
modifiedSQL := strings.Replace(db.Statement.SQL.String(), "FROM", "FROM /* hint */", 1)
db.Statement.SQL.Reset()
db.Statement.SQL.WriteString(modifiedSQL)
}
})
11.5 自定义数据类型
// 实现 Scanner/Valuer 接口
type EncryptedString struct {
Value string
}
// 存入数据库时加密
func (e EncryptedString) Value() (driver.Value, error) {
if e.Value == "" {
return nil, nil
}
return encrypt(e.Value), nil
}
// 从数据库读取时解密
func (e *EncryptedString) Scan(value interface{}) error {
if value == nil {
return nil
}
switch v := value.(type) {
case string:
e.Value = decrypt(v)
case []byte:
e.Value = decrypt(string(v))
}
return nil
}
// 使用
type User struct {
ID uint
Password EncryptedString
}
11.6 自定义 Clauses
import "gorm.io/gorm/clause"
// 自定义 ON CONFLICT 行为
db.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "id"}},
DoUpdates: clause.AssignmentColumns([]string{"name", "updated_at"}),
}).Create(&users)
// MySQL 的 ON DUPLICATE KEY UPDATE
db.Clauses(clause.Insert{
Modifier: "IGNORE",
}).Create(&users)
11.7 会话模式
// 创建会话
session := db.Session(&gorm.Session{
PrepareStmt: true, // 启用 Prepared Statement
SkipHooks: true, // 跳过钩子
DisableNestedTransaction: true, // 禁用嵌套事务
AllowGlobalUpdate: true, // 允许全局更新
Context: ctx, // 使用上下文
Logger: logger.Default.LogMode(logger.Silent),
})
// 在会话中执行操作
session.Create(&user)
session.Find(&users)
11.8 连接池高级配置
sqlDB, _ := db.DB()
// 连接池配置
sqlDB.SetMaxIdleConns(10)
sqlDB.SetMaxOpenConns(100)
sqlDB.SetConnMaxLifetime(time.Hour)
// 健康检查
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := sqlDB.PingContext(ctx); err != nil {
log.Fatal("数据库连接失败:", err)
}
// 连接池统计
stats := sqlDB.Stats()
fmt.Printf("打开连接数: %d\n", stats.OpenConnections)
fmt.Printf("空闲连接数: %d\n", stats.Idle)
fmt.Printf("等待连接数: %d\n", stats.WaitCount)
11.9 诊断与调试
SQL 日志
// 启用所有 SQL 日志
db.Logger = logger.Default.LogMode(logger.Info)
// 只记录慢查询
newLogger := logger.New(
log.New(os.Stdout, "\r\n", log.LstdFlags),
logger.Config{
SlowThreshold: 100 * time.Millisecond,
LogLevel: logger.Warn,
Colorful: true,
},
)
db.Logger = newLogger
执行统计
type DBStats struct {
QueryCount int64
AvgQueryTime time.Duration
SlowQueries int64
}
func (s *DBStats) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) {
sql, _ := fc()
elapsed := time.Since(begin)
atomic.AddInt64(&s.QueryCount, 1)
if elapsed > 100*time.Millisecond {
atomic.AddInt64(&s.SlowQueries, 1)
log.Printf("慢查询 [%v]: %s", elapsed, sql)
}
}
11.10 实战:多租户实现
// 基于 Schema 的多租户(PostgreSQL)
type Tenant struct {
ID uint
Schema string
Name string
}
func (t *Tenant) DB(db *gorm.DB) *gorm.DB {
return db.Table(fmt.Sprintf("%s.users", t.Schema))
}
// 使用
tenant := &Tenant{Schema: "tenant_123"}
tenant.DB(db).Create(&user)
// =====
// 基于字段的多租户
type TenantModel struct {
TenantID uint `gorm:"index:idx_tenant_id;not null"`
}
type User struct {
ID uint
Name string
TenantModel
}
// 自动添加租户过滤
db.Callback().Query().Before("gorm:query").Register("tenant_filter", func(db *gorm.DB) {
if tenantID := db.Statement.Context.Value("tenantID"); tenantID != nil {
db.Where("tenant_id = ?", tenantID)
}
})
// 自动设置租户ID
db.Callback().Create().Before("gorm:create").Register("tenant_set", func(db *gorm.DB) {
if tenantID := db.Statement.Context.Value("tenantID"); tenantID != nil {
db.Statement.SetColumn("tenant_id", tenantID)
}
})
11.11 练习题
- 实现一个带超时的数据库查询中间件
- 设计一个支持动态分表的用户系统
- 编写一个自定义加密类型,自动加密敏感字段
11.12 小结
本章介绍了 GORM 的高级特性,包括上下文、复合主键、多数据库、自定义类型等。掌握这些特性可以应对更复杂的业务场景。
本文代码地址:https://github.com/LittleMoreInteresting/gorm_study
欢迎关注公众号,一起学习进步!