rest/plugins/snowflake_id.go

package plugins

import (
	"context"
	"github.com/bwmarrin/snowflake"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
	"os"
	"reflect"
	"strconv"
)

var (
	sf *snowflake.Node
)

func init() {
	var err error
	no, _ := strconv.ParseInt(os.Getenv("CC_NODE"), 10, 64)
	if no == 0 {
		no = 1
	}
	if sf, err = snowflake.NewNode(no); err != nil {
		panic(err)
	}
}

//SnowflakeID 自动生成主键ID
func SnowflakeID(db *gorm.DB) {
	var err error
	if db.Statement.Schema != nil {
		if field := db.Statement.Schema.LookUpField("ID"); field != nil {
			if field.DataType == schema.String {
				if db.Statement.ReflectValue.Kind() == reflect.Array || db.Statement.ReflectValue.Kind() == reflect.Slice {
					for i := 0; i < db.Statement.ReflectValue.Len(); i++ {
						if _, zero := field.ValueOf(context.Background(), db.Statement.ReflectValue.Index(i)); zero {
							if err = field.Set(context.Background(), db.Statement.ReflectValue.Index(i), sf.Generate().String()); err != nil {
								_ = db.AddError(err)
							}
						}
					}
				} else {
					if _, zero := field.ValueOf(context.Background(), db.Statement.ReflectValue); zero {
						db.Statement.SetColumn("ID", sf.Generate().String())
					}
				}
			}
		}
	}
}

func RegisterSnowflakeIDCallback(db *gorm.DB) (err error) {
	return db.Callback().Create().Before("gorm:create").Register("snowflake_id", SnowflakeID)
}