junhong_cmp_fiber/pkg/gorm/callback.go

package gorm

import (
	"reflect"

	"github.com/break/junhong_cmp_fiber/pkg/constants"
	"gorm.io/gorm"
)

// RegisterSetCreatorUpdaterCallback 注册 GORM 创建数据时创建人更新人 Callback
func RegisterSetCreatorUpdaterCallback(db *gorm.DB) error {
	err := db.Callback().Create().Before("gorm:create").Register("set_creator_updater", func(tx *gorm.DB) {
		ctx := tx.Statement.Context
		userID, ok := tx.Statement.Context.Value(constants.ContextKeyUserID).(uint)
		if !ok || tx.Statement.Schema == nil {
			return
		}

		creatorField, hasCreator := tx.Statement.Schema.FieldsByName["Creator"]
		updaterField, hasUpdater := tx.Statement.Schema.FieldsByName["Updater"]
		if !hasCreator && !hasUpdater {
			return
		}

		rv := tx.Statement.ReflectValue
		switch rv.Kind() {
		case reflect.Slice, reflect.Array:
			for i := 0; i < rv.Len(); i++ {
				elem := rv.Index(i)
				if elem.Kind() == reflect.Ptr {
					elem = elem.Elem()
				}
				if hasCreator {
					_ = creatorField.Set(ctx, elem, userID)
				}
				if hasUpdater {
					_ = updaterField.Set(ctx, elem, userID)
				}
			}
		case reflect.Struct:
			if hasCreator {
				_ = creatorField.Set(ctx, rv, userID)
			}
			if hasUpdater {
				_ = updaterField.Set(ctx, rv, userID)
			}
		case reflect.Ptr:
			elem := rv.Elem()
			if elem.Kind() == reflect.Struct {
				if hasCreator {
					_ = creatorField.Set(ctx, elem, userID)
				}
				if hasUpdater {
					_ = updaterField.Set(ctx, elem, userID)
				}
			}
		}
	})
	return err
}