使用 gorm.DefaultTableNameHandler 可能存在的问题

王小明_HIT

发布于 2021-12-20 21:41:09

1.3K0

发布于 2021-12-20 21:41:09

文章被收录于专栏：程序员奇点

业务背景

有这样的业务场景, 线上一个表 tablea，生产环境还有一个镜像表 tablea_mirror，现在你需要当请求中有一些 tag 标识的时候，访问 tablea_mirror 表，有时候会用到 DefaultTableNameHandler

先安装 sqlite

https://wangxiaoming.blog.csdn.net/article/details/121884736

代码

可以使用 DefaultTableNameHandler 来实现加前缀或者后缀功能。


import (
 "code.byted.org/gopkg/gorm"
 "context"
)

type dbStagingPostfixKeyType struct{}

var dbStagingPostfixKey = dbStagingPostfixKeyType{}

func WithDbStagingPostfix(ctx context.Context, postfix string) context.Context {
 return context.WithValue(ctx, dbStagingPostfixKey, postfix)
}

func ReWriteTableName() {
 gorm.DefaultTableNameHandler = func(db *gorm.DB, defaultTableName string) string {
  if v := db.Ctx.Value(dbStagingPostfixKey); v != nil {
   return defaultTableName + v.(string)
  }
  return defaultTableName
 }
}

测试代码

package mysql

import (
 "fmt"
 "testing"

 "github.com/jinzhu/gorm"
 //_ "github.com/jinzhu/gorm/dialects/sqlite"
 _ "github.com/mattn/go-sqlite3"
)

type Product struct {
 gorm.Model
 Code  string
 Price uint
}

func (Product) TableName() string {
 return "hax_products"
}
func Test(t *testing.T) {
 db, err := gorm.Open("sqlite3", "test.db")
 if err != nil {
  panic("failed to connect database")
 }
 defer db.Close()
 gorm.DefaultTableNameHandler = func(db *gorm.DB, defaultTableName string) string {
  return "hax_" + defaultTableName
 }
 db.LogMode(true)
 // Migrate the schema
 db.AutoMigrate(&Product{})
 db.Create(&Product{Code: "L1212", Price: 1000})
 var product Product
 db.First(&product, 1)
 var products []Product
 db.Find(&products)
 fmt.Printf("Total count %d", len(products))
}

执行结果：

(/Users/xxx/go/src/xxx/xxx.xx/GoProject/mysql/sqllite_test.go:33) 
[2021-12-14 21:43:33]  [1.38ms]  INSERT INTO "hax_products" ("created_at","updated_at","deleted_at","code","price") VALUES ('2021-12-14 21:43:33','2021-12-14 21:43:33',NULL,'L1212',1000)  
[1 rows affected or returned ] 

(/Users/xxx/go/src/xxx/xxx.xx/GoProject/mysql/sqllite_test.go:35) 
[2021-12-14 21:43:33]  [0.23ms]  SELECT * FROM "hax_products"  WHERE "hax_products"."deleted_at" IS NULL AND (("hax_products"."id" = 1)) ORDER BY "hax_products"."id" ASC LIMIT 1  
[1 rows affected or returned ] 

((/Users/xxx/go/src/xxx/xxx.xx/GoProject/mysql/sqllite_test.go:37) 
[2021-12-14 21:43:33]  no such table: hax_hax_products 

((/Users/xxx/go/src/xxx/xxx.xx/GoProject/mysql/sqllite_test.go:37) 
[2021-12-14 21:43:33]  [0.10ms]  SELECT * FROM "hax_hax_products"  WHERE "hax_hax_products"."deleted_at" IS NULL  
[0 rows affected or returned ] 
Total count 0--- PASS: Test (0.00s)
PASS

根据执行结果，可以看到，创建语言与查询单条记录时表名为 hax_products 但是查询多条记录时，却使用了表名hax_hax_products.

这个就是坑1

查询单个记录时使用了TableName()返回的表名，而在查询结果为Array时，表名在TableName()的基础上又添加了前缀。

Gorm 结构体一般分析如下 struct

type DB struct (gorm/main.go)代表数据库连接，每次操作数据库会创建出clone对象。方法gorm.Open()返回的值类型就是这个结构体指针。
type Scope struct (gorm/scope.go) 当前数据库操作的信息，每次添加条件时也会创建clone对象。
type Callback struct (gorm/callback.go) 数据库各种操作的回调函数， SQL生成也是靠这些回调函数。每种类型的回调函数放在单独的文件里，比如查询回调函数在gorm/callback_query.go，创建的在gorm/callback_create.go

db.First() 代码分析

First()方法位于gorm/main.go文件中， .callCallbacks(s.parent.callbacks.queries)调用了query回调函数。

// file: gorm/main.go
// First find first record that match given conditions, order by primary key
func (s *DB) First(out interface{}, where ...interface{}) *DB {
 newScope := s.NewScope(out)
 newScope.Search.Limit(1)
 return newScope.Set("gorm:order_by_primary_key", "ASC").
  inlineCondition(where...).callCallbacks(s.parent.callbacks.queries).db
}

Callback结构体中定义queries为函数指针数组，而默认值的初始化在gorm/callback_query.go的init()方法中，查询方法为queryCallback, 而queryCallback()方法又调用到scope.prepareQuerySQL(), scope中的方法真正生成SQL的地方。

// file: gorm/callback.go
type Callback struct {
 logger     logger
 creates    []*func(scope *Scope)
 updates    []*func(scope *Scope)
 deletes    []*func(scope *Scope)
 queries    []*func(scope *Scope)
 rowQueries []*func(scope *Scope)
 processors []*CallbackProcessor
}
// file: gorm/callback_query.go
// Define callbacks for querying
func init() {
 DefaultCallback.Query().Register("gorm:query", queryCallback)
 DefaultCallback.Query().Register("gorm:preload", preloadCallback)
 DefaultCallback.Query().Register("gorm:after_query", afterQueryCallback)
}
// queryCallback used to query data from database
func queryCallback(scope *Scope) {
...
    scope.prepareQuerySQL()
...
}

跟踪代码到scope.go文件，函数TableName()是获取数据库表名的地方。它按如下顺序来确定表名：

scope.Search.tableName 查询条件中设置了表名，则直接使用
scope.Value.(tabler) 值对象实现了tabler接口(方法TableName() string), 则从调用方法获取
scope.Value.(dbTabler) 值对象实现了dbTabler接口(方法TableName(*DB) string), 则从调用方法获取
若以上条件都不成立，则从scope.GetModelStruct()中获取对象的结构体信息，从结构体名生成表名

具体可见 scope.go 源码

// file: gorm/scope.go
func (scope *Scope) prepareQuerySQL() {
 if scope.Search.raw {
  scope.Raw(scope.CombinedConditionSql())
 } else {
  scope.Raw(fmt.Sprintf("SELECT %v FROM %v %v", scope.selectSQL(), scope.QuotedTableName(), scope.CombinedConditionSql()))
 }
 return
}
// QuotedTableName return quoted table name
func (scope *Scope) QuotedTableName() (name string) {
 if scope.Search != nil && len(scope.Search.tableName) > 0 {
  if strings.Contains(scope.Search.tableName, " ") {
   return scope.Search.tableName
  }
  return scope.Quote(scope.Search.tableName)
 }
 return scope.Quote(scope.TableName())
}
// TableName return table name
func (scope *Scope) TableName() string {
 if scope.Search != nil && len(scope.Search.tableName) > 0 {
  return scope.Search.tableName
 }
 if tabler, ok := scope.Value.(tabler); ok {
  return tabler.TableName()
 }
 if tabler, ok := scope.Value.(dbTabler); ok {
  return tabler.TableName(scope.db)
 }
 return scope.GetModelStruct().TableName(scope.db.Model(scope.Value))
}

对比以上条件，示例中的Product结构体定义了方法TableName() string，符合条件2，那么db.First(&product, 1)使用的表名就是hax_products。

db.Find() 代码分析

Find()代码如下，与First()同样是使用了callbacks.queries回调方法，不同点在于设置了newScope.Search.Limit(1)只返回一个结果、增加了按id排序。

// Find find records that match given conditions
func (s *DB) Find(out interface{}, where ...interface{}) *DB {
 return s.NewScope(out).inlineCondition(where...).callCallbacks(s.parent.callbacks.queries).db
}

在debug模式下跟踪代码到scope.TableName()中时，两次查询的区别显示出来了：它们的结果值类型不同。db.First(&product, 1)的值类型为结构体的指针*Product，而db.Find(&products)的值类型是数组的指针*[]Product，从而导致db.Find(&products)进入条件 scope.GetModelStruct().TableName(scope.db.Model(scope.Value)) }，需要靠分析struct结构体来生成表名。

// file: gorm/model_struct.go
// TableName returns model's table name
func (s *ModelStruct) TableName(db *DB) string {
 s.l.Lock()
 defer s.l.Unlock()
 if s.defaultTableName == "" && db != nil && s.ModelType != nil {
  // Set default table name
  if tabler, ok := reflect.New(s.ModelType).Interface().(tabler); ok {
   s.defaultTableName = tabler.TableName()
  } else {
   tableName := ToTableName(s.ModelType.Name())
   db.parent.RLock()
   if db == nil || (db.parent != nil && !db.parent.singularTable) {
    tableName = inflection.Plural(tableName)
   }
   db.parent.RUnlock()
   s.defaultTableName = tableName
  }
 }
 return DefaultTableNameHandler(db, s.defaultTableName)
}

默认表名s.defaultTableName为空值时先进行求值，reflect.New(s.ModelType).Interface().(tabler)先判断是否实现了tabler接口，有则调用其TableName()取值；否则的话从结构体的名字来生成表名。结果返回之前再调用 DefaultTableNameHandler(db, s.defaultTableName)方法。

这个ModelStruct的TableName方法与scope.TableName() 中的逻辑两个不一致的地方：

scope.TableName()会判断是否实现tabler与dbTabler两个接口，而这里只判断了tabler
scope.TableName()是将tableName结果直接返回的，而这里多调用了DefaultTableNameHandler()。

因为逻辑 scope.TableName()的存在，当重写DefaultTableNameHandler()方法时，就会出现表前缀再次被添加了表名前。

问题2

DefaultTableNameHandler()在多数据库时出现混乱

通过以上代码的分析，于是发现了另一个坑：当一个程序中使用两个不同的数据库时，重写方法DefaultTableNameHandler()会影响到两个数据库中的表名。其中一个数据库需要设置表前缀时，访问另一个数据库的表也可能会被加上前缀。因为是包级别的方法，整个代码里只能设置一次值。

// file: gorm/model_struct.go
// DefaultTableNameHandler default table name handler
var DefaultTableNameHandler = func(db *DB, defaultTableName string) string {
 return defaultTableName
}