70 lines
1.7 KiB
Go
70 lines
1.7 KiB
Go
// Package snowid 基于雪花算法的全局唯一 ID 生成器（进程内单例，需先 Init）。
|
|
//
|
|
// 运行时：main 里 Init 后调用 Generate/GenerateString。
|
|
// GORM 模型：IDMethod 交给 gorm.io/gen 的 WithMethod 注入，
|
|
// 生成的模型自带 BeforeCreate 钩子（Id 为 0 时自动填充 Generate()）。
|
|
package snowid
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/bwmarrin/snowflake"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var (
|
|
node *snowflake.Node
|
|
once sync.Once
|
|
initErr error
|
|
)
|
|
|
|
// Init 初始化雪花ID生成器
|
|
// nodeID: 节点ID，范围 0-1023
|
|
func Init(nodeID int64) error {
|
|
once.Do(func() {
|
|
node, initErr = snowflake.NewNode(nodeID)
|
|
})
|
|
return initErr
|
|
}
|
|
|
|
// InitWithStartTime 使用自定义起始时间初始化雪花ID生成器
|
|
// nodeID: 节点ID，范围 0-1023
|
|
// startTime: 自定义起始时间（Unix毫秒时间戳）
|
|
func InitWithStartTime(nodeID int64, startTime int64) error {
|
|
once.Do(func() {
|
|
snowflake.Epoch = startTime
|
|
node, initErr = snowflake.NewNode(nodeID)
|
|
})
|
|
return initErr
|
|
}
|
|
|
|
// Generate 生成一个新的雪花ID
|
|
func Generate() int64 {
|
|
if node == nil {
|
|
panic("snowid: not initialized, please call Init() first")
|
|
}
|
|
return node.Generate().Int64()
|
|
}
|
|
|
|
// GenerateString 生成字符串形式的雪花ID
|
|
func GenerateString() string {
|
|
if node == nil {
|
|
panic("snowid: not initialized, please call Init() first")
|
|
}
|
|
return node.Generate().String()
|
|
}
|
|
|
|
// IDMethod GORM 模型方法模板：交给 gen.WithMethod 注入后，
|
|
// 生成的模型在 Create 前自动生成雪花ID主键（字段名须为 Id）。
|
|
type IDMethod struct {
|
|
Id int64
|
|
}
|
|
|
|
// BeforeCreate 创建前钩子，自动生成雪花ID
|
|
func (m *IDMethod) BeforeCreate(tx *gorm.DB) error {
|
|
if m.Id == 0 {
|
|
m.Id = Generate()
|
|
}
|
|
return nil
|
|
}
|