利用 callback 实现自动 etag 字段

为了实现乐观锁，我们经常会增加一个 etag 字段作为判断语句执行时的判断条件，类似如下的 model 定义：

type Object struct {
	gorm.Model
	Key string
	Etag string
}

对于这个 model，我们希望能在执行UPDATE/DELETE 语句的时候，自动增加一个判断条件 etag='some-uuid'，这样我们在业务代码中就不用每次都加上 etag 字段的处理了，即方便，也避免了出错。

在 GORM 中，可以使用可以使用 hook 来实现这个需求，可以通过实现 BeforeCreate/BeforeUpdate/BeforeDelete 等方法来实现 etag 字段的自动处理。但是 hook 的方式有个问题，就是每个 model 都要实现一次这个方法。

一个简单的改进是，可以自己实现一个用于嵌入的 BaseModel，类似于 gorm.Model，在这个 model 上实现一次这些方法即可，其他 embed 了这个 BaseModel 的 model 就自动实现了这些方法。但是这个方法也有个问题，如果一个 model 想实现自己的 hook，那么就需要在自己的 hook 中调用 BaseModel 的对应 hook。但是这么做的话，不仅会降低效率，也会增加出错的概率。

GORM 还有另外一个方式来实现这个需求，就是 callback 系统。通过 callback，不仅可以实现在有 etag 字段存在时就自动处理的逻辑，也不会和 hook 系统冲突。下面是一个 create 和 update 时自动处理的 etag 的实现。

package model

import (
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/schema"
)

func callbackAddEtag(db *gorm.DB) {
	etagField := db.Statement.Schema.LookUpField(FieldEtag)
	if etagField == nil {
		return
	}
	if etagField.DataType != schema.String {
		return
	}

	etag := NewEtag()
	etagField.Set(db.Statement.ReflectValue, etag)
}

func callbackUpdateCheckEtag(db *gorm.DB) {
	stmt := db.Statement
	if stmt.Schema == nil {
		// Schema will be nil if db was not called with Model() function.
		return
	}

	etagField := stmt.Schema.LookUpField(FieldEtag)
	if etagField == nil {
		return
	}
	if etagField.DataType != schema.String {
		return
	}
	etag, isZero := etagField.ValueOf(stmt.ReflectValue)
	if isZero {
		return
	}

	where := map[string]interface{}{
		ColEtag: etag,
	}
	conds := stmt.BuildCondition(where)
	newWhereClause := clause.Where{Exprs: conds}
	stmt.AddClause(newWhereClause)

	destValue := reflect.ValueOf(stmt.Dest)
	for destValue.Kind() == reflect.Ptr {
		destValue = destValue.Elem()
	}
	switch dest := destValue.Interface().(type) {
	case map[string]interface{}:
		dest[ColEtag] = NewEtag()
	default:
	}
}

// RegisterDB registers a database as basicDB which is used in future operations
func RegisterDB(arg *ConnectionArg) (err error) {
	basicDBLock.Lock()
	defer basicDBLock.Unlock()

	if basicDB != nil {
		if err = closeDB(basicDB); err != nil {
			return err
		}
		basicDB = nil
	}

	dsn := arg.DSN()
	basicDB, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
	if err != nil {
		return err
	}

	basicDB.Callback().Create().
		Before("gorm:before_create").
		Register("add_etag", callbackAddEtag)
	basicDB.Callback().Update().Before("gorm:update").
		Register("check_etag", callbackUpdateCheckEtag)

	return nil
}

上述这段代码的总来说就是在 GORM 默认的 callbak 之前插入我们自定义的 callback，这些自定义的 callback 会修改将要执行的语句 (db.Statement)。关于 GORM callback 的文档，见这里： Write Plugin。特别提一下，这个文档里说，默认的 callback 在文件 callbacks/callback.go 中的 RegisterDefaultCallbacks 函数中被注册，不过你在 gorm.io/gorm 是找不到这个函数被调用的地方的，因为这个函数是在 driver 中被调用的，可以查看 gorm.io/driver/postgres 的 Dialector 的实现。

本作品采用知识共享署名-非商业性使用-禁止演绎 4.0 国际许可协议进行许可。