gorm全局钩子函数修改主键生成规则

一、基本步骤

• 1、安装雪花算法的依赖包

  go get -u github.com/bwmarrin/snowflake
  

• 2、初始化 雪花算法生成主键id

  var node *snowflake.Node
  
  func init() {
  	// 初始化雪花算法节点
  	var err error
  	node, err = snowflake.NewNode(1) // 节点 ID 可以根据实际情况设置
  	if err != nil {
  		panic(err)
  	}
  }
  

• 3、自定义钩子函数

  // 全局钩子：在创建记录时生成雪花算法 ID
  func generateSnowflakeID(db *gorm.DB) {
  	if db.Statement.Schema != nil {
  		for _, field := range db.Statement.Schema.Fields {
  			if field.Name == "ID" && field.PrimaryKey {
  				// 生成雪花算法 ID
  				id := node.Generate().Int64()
  				db.Statement.SetColumn(field.Name, id)
  			}
  		}
  	}
  }
  

• 4、注册钩子函数

  // 创建数据库连接
  	sqlStr := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s",
  		"root",
  		"123456",
  		"localhost",
  		"3306",
  		"test2",
  	)
  	fmt.Println("数据库连接:", sqlStr)
  
  	db, err := gorm.Open(mysql.Open(sqlStr), &gorm.Config{
  		Logger:                                   logger.Default.LogMode(logger.Info),
  		DisableForeignKeyConstraintWhenMigrating: true, // 自动创建表的时候不创建外键
  		NamingStrategy: schema.NamingStrategy{ // 自动创建表时候表名的配置
  			SingularTable: true,
  			// 全部的表名前面加前缀
  			//TablePrefix: "mall_",
  		},
  	})
  	// 注册全局钩子
  	db.Callback().Create().Before("gorm:create").Register("before_create_generate_id", generateSnowflakeID)
  	if err != nil {
  		panic("failed to connect database")
  	}