高级特性


第十一章:高级特性

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 练习题

  1. 实现一个带超时的数据库查询中间件
  2. 设计一个支持动态分表的用户系统
  3. 编写一个自定义加密类型,自动加密敏感字段

11.12 小结

本章介绍了 GORM 的高级特性,包括上下文、复合主键、多数据库、自定义类型等。掌握这些特性可以应对更复杂的业务场景。


本文代码地址:https://github.com/LittleMoreInteresting/gorm_study

欢迎关注公众号,一起学习进步!

如有疑问关注公众号给我留言
wx

关注公众号

©2017-2023 鲁ICP备17023316号-1 Powered by Hugo