|
|
|
package database
|
|
|
|
|
|
|
|
import (
|
|
|
|
"fmt"
|
|
|
|
"github.com/jinzhu/now"
|
|
|
|
"github.com/zeromicro/go-zero/core/stores/redis"
|
|
|
|
"gorm.io/gorm"
|
|
|
|
"gorm.io/gorm/clause"
|
|
|
|
"gorm.io/gorm/migrator"
|
|
|
|
"gorm.io/gorm/schema"
|
|
|
|
"strings"
|
|
|
|
"time"
|
|
|
|
)
|
|
|
|
|
|
|
|
var (
|
|
|
|
// PartitionByRangeTime 按unix时间戳分区
|
|
|
|
PartitionByRangeTime = 1
|
|
|
|
// PartitionByHash 按系统的hash值分区
|
|
|
|
PartitionByHash = 2
|
|
|
|
// PartitionByList 按List包含值分区
|
|
|
|
PartitionByList = 3
|
|
|
|
)
|
|
|
|
|
|
|
|
type PartitionTable interface {
|
|
|
|
TableName() string
|
|
|
|
}
|
|
|
|
|
|
|
|
type PartitionMigrator struct {
|
|
|
|
ServiceName string
|
|
|
|
DB *gorm.DB
|
|
|
|
Redis *redis.Redis
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewPartitionMigrator(serviceName string, db *gorm.DB, redis *redis.Redis) *PartitionMigrator {
|
|
|
|
return &PartitionMigrator{
|
|
|
|
DB: db,
|
|
|
|
ServiceName: serviceName,
|
|
|
|
Redis: redis,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (c *PartitionMigrator) AutoMigrate(t PartitionTable, option ...PartitionOptionFunc) error {
|
|
|
|
options := NewPartitionOptions()
|
|
|
|
for i := range option {
|
|
|
|
option[i](options)
|
|
|
|
}
|
|
|
|
|
|
|
|
tableName := t.TableName()
|
|
|
|
if !c.DB.Migrator().HasTable(tableName) {
|
|
|
|
migrator := Migrator{migrator.Migrator{
|
|
|
|
migrator.Config{
|
|
|
|
CreateIndexAfterCreateTable: true,
|
|
|
|
DB: c.DB,
|
|
|
|
Dialector: c.DB.Dialector,
|
|
|
|
},
|
|
|
|
}}
|
|
|
|
if err := migrator.CreatePartitionTable(options, t); err != nil {
|
|
|
|
panic(err)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
rk := fmt.Sprintf("%s:auto-partition:%s", c.ServiceName, tableName)
|
|
|
|
lock := redis.NewRedisLock(c.Redis, rk)
|
|
|
|
ok, err := lock.Acquire()
|
|
|
|
if !ok || err != nil {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
defer lock.Release()
|
|
|
|
switch options.Type {
|
|
|
|
case PartitionByRangeTime:
|
|
|
|
begin := options.TimeBegin
|
|
|
|
end := options.TimeEnd
|
|
|
|
for {
|
|
|
|
if begin.Unix() > end.Unix() {
|
|
|
|
break
|
|
|
|
}
|
|
|
|
pTable := fmt.Sprintf("%s_%s", tableName, options.FormatTimeSubFunc(begin))
|
|
|
|
sql := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s PARTITION OF %s FOR VALUES FROM (%d) TO (%d);",
|
|
|
|
pTable, tableName, begin.Unix(), begin.AddDate(0, options.TimeSpanMonth, 0).Unix())
|
|
|
|
tx := c.DB.Exec(sql)
|
|
|
|
if tx.Error != nil {
|
|
|
|
return tx.Error
|
|
|
|
}
|
|
|
|
c.log(t, pTable)
|
|
|
|
begin = begin.AddDate(0, options.TimeSpanMonth, 0)
|
|
|
|
}
|
|
|
|
break
|
|
|
|
case PartitionByHash:
|
|
|
|
for i := 0; i < options.Modulus; i++ {
|
|
|
|
pTable := fmt.Sprintf("%s_%d", tableName, i)
|
|
|
|
sql := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s PARTITION OF %s FOR VALUES WITH (MODULUS %d, REMAINDER %d);",
|
|
|
|
pTable, tableName, options.Modulus, i)
|
|
|
|
tx := c.DB.Exec(sql)
|
|
|
|
if tx.Error != nil {
|
|
|
|
return tx.Error
|
|
|
|
}
|
|
|
|
c.log(t, pTable)
|
|
|
|
}
|
|
|
|
break
|
|
|
|
case PartitionByList:
|
|
|
|
for i := 0; i < len(options.ListRange); i++ {
|
|
|
|
pTable := fmt.Sprintf("%s_%d", tableName, i)
|
|
|
|
sql := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s PARTITION OF %s FOR VALUES IN %s;",
|
|
|
|
pTable, tableName, InArgs(options.ListRange[i]))
|
|
|
|
tx := c.DB.Exec(sql)
|
|
|
|
if tx.Error != nil {
|
|
|
|
return tx.Error
|
|
|
|
}
|
|
|
|
c.log(t, pTable)
|
|
|
|
}
|
|
|
|
break
|
|
|
|
default:
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (c *PartitionMigrator) log(t PartitionTable, pTable string) {
|
|
|
|
fmt.Println("【自动分区】 create partition table", pTable, "on table", t.TableName())
|
|
|
|
}
|
|
|
|
|
|
|
|
type PartitionOptions struct {
|
|
|
|
// 分区类型 1:Hash 2:RangeTime
|
|
|
|
Type int
|
|
|
|
// 分区列
|
|
|
|
Column string
|
|
|
|
|
|
|
|
// Hash分区
|
|
|
|
Modulus int
|
|
|
|
|
|
|
|
// List 范围
|
|
|
|
ListRange []interface{}
|
|
|
|
|
|
|
|
// Range时间分区
|
|
|
|
TimeBegin time.Time
|
|
|
|
TimeEnd time.Time
|
|
|
|
TimeSpanMonth int
|
|
|
|
FormatTimeSubFunc func(time.Time) string
|
|
|
|
|
|
|
|
// 禁用PrimaryKey生成
|
|
|
|
// 分区字段有函数表达式的，需要禁用掉PrimaryKey,使用自定义的唯一ID生成规则
|
|
|
|
DisablePrimaryKey bool
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewPartitionOptions() *PartitionOptions {
|
|
|
|
return &PartitionOptions{
|
|
|
|
Type: PartitionByRangeTime,
|
|
|
|
FormatTimeSubFunc: func(t time.Time) string {
|
|
|
|
return t.Format("200601")
|
|
|
|
},
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (c *PartitionOptions) Sql() string {
|
|
|
|
if c.Type == PartitionByHash {
|
|
|
|
return fmt.Sprintf("PARTITION BY HASH(%s)", c.Column)
|
|
|
|
}
|
|
|
|
if c.Type == PartitionByRangeTime {
|
|
|
|
return fmt.Sprintf("PARTITION BY RANGE(%s)", c.Column)
|
|
|
|
}
|
|
|
|
if c.Type == PartitionByList {
|
|
|
|
return fmt.Sprintf("PARTITION BY LIST(%s)", c.Column)
|
|
|
|
}
|
|
|
|
return ""
|
|
|
|
}
|
|
|
|
|
|
|
|
type PartitionOptionFunc func(*PartitionOptions)
|
|
|
|
|
|
|
|
func WithPartitionType(t int) PartitionOptionFunc {
|
|
|
|
return func(options *PartitionOptions) {
|
|
|
|
options.Type = t
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func WithPartitionColumn(c string) PartitionOptionFunc {
|
|
|
|
return func(options *PartitionOptions) {
|
|
|
|
options.Column = c
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func WithPartitionHash(modulus int) PartitionOptionFunc {
|
|
|
|
return func(options *PartitionOptions) {
|
|
|
|
options.Modulus = modulus
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func WithPartitionRangeTime(begin, end time.Time, spanMonth int) PartitionOptionFunc {
|
|
|
|
return func(options *PartitionOptions) {
|
|
|
|
options.TimeBegin = begin
|
|
|
|
options.TimeEnd = end
|
|
|
|
options.TimeSpanMonth = spanMonth
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func WithPartitionList(list ...interface{}) PartitionOptionFunc {
|
|
|
|
return func(options *PartitionOptions) {
|
|
|
|
options.ListRange = list
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func WithDisablePrimaryKey(disablePrimaryKey bool) PartitionOptionFunc {
|
|
|
|
return func(options *PartitionOptions) {
|
|
|
|
options.DisablePrimaryKey = disablePrimaryKey
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func Date(date string) time.Time {
|
|
|
|
return now.MustParse(date)
|
|
|
|
}
|
|
|
|
|
|
|
|
type Migrator struct {
|
|
|
|
migrator.Migrator
|
|
|
|
}
|
|
|
|
|
|
|
|
// CreatePartitionTable create table in database for values
|
|
|
|
func (m Migrator) CreatePartitionTable(options *PartitionOptions, values ...interface{}) error {
|
|
|
|
for _, value := range m.ReorderModels(values, false) {
|
|
|
|
tx := m.DB.Session(&gorm.Session{})
|
|
|
|
if err := m.RunWithValue(value, func(stmt *gorm.Statement) (errr error) {
|
|
|
|
var (
|
|
|
|
createTableSQL = "CREATE TABLE ? ("
|
|
|
|
values = []interface{}{m.CurrentTable(stmt)}
|
|
|
|
hasPrimaryKeyInDataType bool
|
|
|
|
)
|
|
|
|
|
|
|
|
for _, dbName := range stmt.Schema.DBNames {
|
|
|
|
field := stmt.Schema.FieldsByDBName[dbName]
|
|
|
|
if !field.IgnoreMigration {
|
|
|
|
createTableSQL += "? ?"
|
|
|
|
hasPrimaryKeyInDataType = hasPrimaryKeyInDataType || strings.Contains(strings.ToUpper(string(field.DataType)), "PRIMARY KEY")
|
|
|
|
values = append(values, clause.Column{Name: dbName}, m.DB.Migrator().FullDataTypeOf(field))
|
|
|
|
createTableSQL += ","
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if !hasPrimaryKeyInDataType && len(stmt.Schema.PrimaryFields) > 0 && !options.DisablePrimaryKey {
|
|
|
|
createTableSQL += "PRIMARY KEY ?,"
|
|
|
|
primaryKeys := []interface{}{}
|
|
|
|
for _, field := range stmt.Schema.PrimaryFields {
|
|
|
|
primaryKeys = append(primaryKeys, clause.Column{Name: field.DBName})
|
|
|
|
}
|
|
|
|
|
|
|
|
values = append(values, primaryKeys)
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, idx := range stmt.Schema.ParseIndexes() {
|
|
|
|
if m.CreateIndexAfterCreateTable {
|
|
|
|
defer func(value interface{}, name string) {
|
|
|
|
if errr == nil {
|
|
|
|
errr = tx.Migrator().CreateIndex(value, name)
|
|
|
|
}
|
|
|
|
}(value, idx.Name)
|
|
|
|
} else {
|
|
|
|
if idx.Class != "" {
|
|
|
|
createTableSQL += idx.Class + " "
|
|
|
|
}
|
|
|
|
createTableSQL += "INDEX ? ?"
|
|
|
|
|
|
|
|
if idx.Comment != "" {
|
|
|
|
createTableSQL += fmt.Sprintf(" COMMENT '%s'", idx.Comment)
|
|
|
|
}
|
|
|
|
|
|
|
|
if idx.Option != "" {
|
|
|
|
createTableSQL += " " + idx.Option
|
|
|
|
}
|
|
|
|
|
|
|
|
createTableSQL += ","
|
|
|
|
values = append(values, clause.Column{Name: idx.Name}, tx.Migrator().(migrator.BuildIndexOptionsInterface).BuildIndexOptions(idx.Fields, stmt))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if !m.DB.DisableForeignKeyConstraintWhenMigrating && !m.DB.IgnoreRelationshipsWhenMigrating {
|
|
|
|
for _, rel := range stmt.Schema.Relationships.Relations {
|
|
|
|
if rel.Field.IgnoreMigration {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
if constraint := rel.ParseConstraint(); constraint != nil {
|
|
|
|
if constraint.Schema == stmt.Schema {
|
|
|
|
sql, vars := buildConstraint(constraint)
|
|
|
|
createTableSQL += sql + ","
|
|
|
|
values = append(values, vars...)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, chk := range stmt.Schema.ParseCheckConstraints() {
|
|
|
|
createTableSQL += "CONSTRAINT ? CHECK (?),"
|
|
|
|
values = append(values, clause.Column{Name: chk.Name}, clause.Expr{SQL: chk.Constraint})
|
|
|
|
}
|
|
|
|
|
|
|
|
createTableSQL = strings.TrimSuffix(createTableSQL, ",")
|
|
|
|
|
|
|
|
createTableSQL += ")"
|
|
|
|
|
|
|
|
if options != nil {
|
|
|
|
createTableSQL += options.Sql()
|
|
|
|
}
|
|
|
|
|
|
|
|
if tableOption, ok := m.DB.Get("gorm:table_options"); ok {
|
|
|
|
createTableSQL += fmt.Sprint(tableOption)
|
|
|
|
}
|
|
|
|
|
|
|
|
errr = tx.Exec(createTableSQL, values...).Error
|
|
|
|
return errr
|
|
|
|
}); err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func buildConstraint(constraint *schema.Constraint) (sql string, results []interface{}) {
|
|
|
|
sql = "CONSTRAINT ? FOREIGN KEY ? REFERENCES ??"
|
|
|
|
if constraint.OnDelete != "" {
|
|
|
|
sql += " ON DELETE " + constraint.OnDelete
|
|
|
|
}
|
|
|
|
|
|
|
|
if constraint.OnUpdate != "" {
|
|
|
|
sql += " ON UPDATE " + constraint.OnUpdate
|
|
|
|
}
|
|
|
|
|
|
|
|
var foreignKeys, references []interface{}
|
|
|
|
for _, field := range constraint.ForeignKeys {
|
|
|
|
foreignKeys = append(foreignKeys, clause.Column{Name: field.DBName})
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, field := range constraint.References {
|
|
|
|
references = append(references, clause.Column{Name: field.DBName})
|
|
|
|
}
|
|
|
|
results = append(results, clause.Table{Name: constraint.Name}, foreignKeys, clause.Table{Name: constraint.ReferenceSchema.Table}, references)
|
|
|
|
return
|
|
|
|
} |
...
|
...
|
|
