GORM 中同一个 *DB 对象污染问题的分析与解决

167 阅读5分钟

现象

先来看一个示例代码(使用 PostgresDB,account 表结构与结构体 Account 相同):

type Account struct {
    UserId   string
    Username string
    Email    string
    Phone    string
}

func main() {
    // getDB实现省略
    db, err := getDB("root", "secret", "localhost", 15432, "testdb")
    if err != nil {
        log.Println(err)
        return
    }

    var cnt int64
    var items []Account

    // 假设此处只做了一个简单的 where 条件
    db = db.Table("account").Where("username='admin'")
    
    // 希望在这里保留上面 db 的 where 条件,用于统计总数
    // 但又不想让后续的操作影响到这里,所以创建了 cntDb
    cntDb := db

    // 下面这行代码对 db 做了分页限制
    // 期望 cntDb 不被影响,然而执行后会发现 cntDb 被污染了
    db = db.Limit(1).Offset(2)

    // 执行 .Count 时却带上了分页限制
    if err := cntDb.Debug().Count(&cnt).Error; err != nil {
        log.Println(err)
        return
    }

    // 查询实际数据
    if err := db.Debug().Find(&items).Error; err != nil {
        log.Println(err)
        return
    }

    log.Printf("items: %#v, cnt: %d", items, cnt)
    time.Sleep(3 * time.Second)
}

运行代码后,输出的 SQL 日志如下:

2025/01/08 00:48:26 D:/Code/.../main.go:32
[2.208ms] [rows:0] SELECT count(*) FROM "account" WHERE username='admin' LIMIT 1 OFFSET 2

2025/01/08 00:48:26 D:/Code/.../main.go:37
[0.979ms] [rows:0] SELECT * FROM "account" WHERE username='admin' LIMIT 1 OFFSET 2
2025/01/08 00:48:26 items: []main.Account{}, cnt: 0

其中可以看到,原本只想给查询数据(Find)加上 LIMIT 1 OFFSET 2 的语句,却意外影响到了统计总数(Count)的 SQL。显然,cntDb 并没有独立于 db,导致最终统计结果也带上了分页限制。

初步尝试:Session

为了解决这个问题,我希望在创建 cntDb 的时候,既能保留 db 原先的 Where 条件,又能保证后续对 db 的修改不会“污染”cntDb。初步想到的方法是使用 Session 函数:

cntDb := db.Session(&gorm.Session{})

从 GORM 的注释中可以看到:Session create new db session,这似乎意味着它会创建一个新的 db 对象。可惜实际测试中,cntDb 依然会被后续 db.Limit(...) 的操作所影响。

再试:将 NewDB 设为 true

gorm.Session 的配置结构体中还有一个 bool 类型的 NewDB 选项。如果将其置为 true,就会在新 Session 中将clone值置为1,会导致创建一个全新的 Statement,理论上可以隔离先前的条件。这样虽然可以得到一个干净的 db,但原本希望保留的 Where 条件也被清空了,仍旧无法满足需求。

并且,即使把NewDB置为true,也不会立即实现Statement的深拷贝, 也无法直接与原来的db对象的Statement分离

深入研究:为什么会“污染”?

要理解“污染”原因,需要先了解 GORM 内部对 clone 以及 Statement 的操作机制。GORM 中的大部分链式调用都会将查询条件保存在一个名为 Statement 的结构体里。根据 GORM 源码可知,当我们调用一些查询方法(如 .Where()/.Limit()/.Offset())时,都会修改同一个 Statement 的内容。

gorm.Session 内部,有一个和 clone 相关的逻辑,核心代码大致如下(简化):

func (db *DB) getInstance() *DB {
    if db.clone > 0 {
        tx := &DB{Config: db.Config, Error: db.Error}

        if db.clone == 1 {
            // 当 clone == 1 时,创建一个全新的 Statement
            tx.Statement = &Statement{...} // 干净的 Statement
        } else {
            // 当 clone == 2 (或更大) 时,就会深拷贝 db 的 Statement
            tx.Statement = db.Statement.clone()
            tx.Statement.DB = tx
        }

        return tx
    }
    return db
}
  • clone == 1 时,Statement 是全新创建的,原先的查询条件不再保留。
  • clone == 2 时,Statement 会被深拷贝到新的 tx,从而继承已设置的查询条件,却又和原对象分离,不再互相影响。

Session 函数中,如果 NewDB = true,就会把 clone 值设为 1,最终导致清空了原先的 Where 条件;如果 NewDB = false,则不会强制在此处进行一次 getInstance(),因此也就不会完成深拷贝,后面对 db 的操作仍然会影响到 cntDb

更关键的是,这个深拷贝操作往往只有在真正“调用”数据库操作(例如 Count()Find() 等)时才发生。如果只是简单调用 Session(...) 而并没有发出任何 SQL 请求,cntDb 其实还没来得及 “分身”;接着再对原 db 调用 LimitOffset,自然还是会反作用到 cntDb 上。

解决方案:在 Session 时触发 Statement 克隆

既想在 cntDb 中保留原先的查询条件,又要在后续查询里避免相互影响,关键在于——在调用 Session 时,能让 GORM 立刻进行 Statement 的克隆,保证在克隆之前没有对原来的db进行修改

观察到 GORM 在下列场景下会触发克隆逻辑:

if config.Context != nil || config.PrepareStmt || config.SkipHooks {
    tx.Statement = tx.Statement.clone()
    tx.Statement.DB = tx
}

这意味着只要我们在 Session 时给 config.Context(或者 PrepareStmtSkipHooks)赋一个值,就能够让 GORM 马上进行深拷贝。实践中,最简便的方式就是给 Context 设个值:

cntDb := db.Session(&gorm.Session{Context: context.Background()})

这样做以后,cntDb 会与 db 彻底分离,保留之前的 Where 条件又不会受后面 .Limit().Offset() 的影响。

总结

  • 问题本质:GORM 通过共享一个 Statement 对象来构建查询条件,多次链式调用会对同一个 Statement 做累积操作,从而出现“一个改了、俩都变”的情况。
  • Session 关键点:只调用 Session(&gorm.Session{}) 并不足以触发深拷贝;需要给 Context 等配置赋值,才能立刻克隆 Statement
  • 正确用法:如果你想要保留当前已有的所有查询条件,并保证后续修改不会影响到这个“快照”,可以这样使用:
cntDb := db.Session(&gorm.Session{Context: context.Background()})
// 或者
cntDb := db.Session(&gorm.Session{PrepareStmt: true})
// 或者
cntDb := db.Session(&gorm.Session{SkipHooks: true})

只要能触发克隆逻辑即可,具体选择哪种配置可根据实际需求来定。 如有更好的实践欢迎交流。