go-gorm / sharding Goto Github PK

                type Order struct {
			ID        int64 `gorm:"primarykey"`
			UserID    int64
			ProductID int64
		}

		for i := 0; i < 64; i += 1 {
			table := fmt.Sprintf("orders_%02d", i)
			db.Exec(`DROP TABLE IF EXISTS ` + table)
			db.Exec(`CREATE TABLE ` + table + ` (id int PRIMARY KEY,user_id int,product_id int)`)
		}

		middleware := sharding.Register(sharding.Config{
			ShardingKey:         "user_id",
			NumberOfShards:      64,
			PrimaryKeyGenerator: sharding.PKSnowflake,
		}, "orders")
		err = db.Use(middleware)
		if err != nil {
			panic(err)
		}

		err = db.Create(&Order{UserID: 2}).Error
		if err != nil {
			fmt.Println(err)
		}

		err = db.Create(&Order{UserID: 3}).Error
		if err != nil {
			fmt.Println(err)
		}

		err = db.Create(&Order{UserID: 4}).Error
		if err != nil {
			fmt.Println(err)
		}

func (s *Sharding) compile() error {
...
			// stmt := &gorm.Statement{DB: s.DB}
			stmt := s.DB.Statement
			if err := stmt.Parse(table); err == nil {
				s.configs[stmt.Table] = s._config
			} else {
				return err
			}
...
}

ShardingAlgorithm: func(columnValue interface{}) (suffix string, err error) {
	if id, ok := columnValue.(uint); ok {
		return fmt.Sprintf("_%02d", id%shardingNum), nil
	}
	return "", fmt.Errorf("invalid column value type")
}

func (d *dataobject) TableName() string {
	return fmt.Sprintf("table_name_%02d", d.ID%shardingNum)
}

2023/02/19 20:57:18 testing mysql...
=== RUN   TestGORM

2023/02/19 20:57:18 /Users/zhengjb/go/src/github.com/zishiguo/playground/main_test.go:13 Error 1054 (42S22): Unknown column 'user_follow.deleted_at' in 'where clause'
[2.526ms] [rows:0] SELECT * FROM `user_follow` WHERE user_id = 79535114761158382 AND `user_follow`.`deleted_at` IS NULL
    main_test.go:14: Failed, got error: Error 1054 (42S22): Unknown column 'user_follow.deleted_at' in 'where clause'
--- FAIL: TestGORM (0.00s)

FAIL

Process finished with the exit code 1

// BeginTx Implement ConnPoolBeginner.BeginTx
func (pool *ConnPool) BeginTx(ctx context.Context, opt *sql.TxOptions) (gorm.ConnPool, error) {
	if basePool, ok := pool.ConnPool.(gorm.ConnPoolBeginner); ok {
		return basePool.BeginTx(ctx, opt)
	}
	return pool, nil
}

func (pool ConnPool) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
	ftQuery, stQuery, table, err := pool.sharding.resolve(query, args...)
	if err != nil {
		return nil, err
	}

	pool.sharding.querys.Store("last_query", stQuery)

	if table != "" {
		if r, ok := pool.sharding.configs[table]; ok {
			if r.DoubleWrite {
				pool.ConnPool.ExecContext(ctx, ftQuery, args...)
			}
		}
	}

	return pool.ConnPool.QueryContext(ctx, stQuery, args...)
}

func NewReadWriteSplits(cm map[string]ReadWriteSplitConfig) (db *gorm.DB) {
	defaultConfig := cm["default"]
	defaultMasterConfig := defaultConfig.Master

	dialector := NewDialectorWithPool(&defaultMasterConfig)

	db, err := gorm.Open(dialector, gormConfig(&defaultMasterConfig))
	if err != nil {
		log.Panicf("database connection failed. database name: %s, err: %+v", defaultMasterConfig.Name, err)
	}
	db.Set("gorm:table_options", "CHARSET=utf8mb4")
	// default
	resolver := dbresolver.Register(dbresolver.Config{
		Replicas: func(cs []Config) []gorm.Dialector {
			var ds []gorm.Dialector
			for _, c := range cs {
				ds = append(ds, NewDialectorWithPool(&c))
			}
			return ds
		}(defaultConfig.Slaves),
		TraceResolverMode: true,
	})
	// other
	for name, c := range cm {
		if name == "default" {
			break
		}
		resolver = resolver.Register(dbresolver.Config{
			Sources: func(c Config) []gorm.Dialector {
				ds := []gorm.Dialector{
					NewDialector(&c),
				}
				return ds
			}(c.Master),
			Replicas: func(cs []Config) []gorm.Dialector {
				var ds []gorm.Dialector
				for _, c := range cs {
					ds = append(ds, NewDialector(&c))
				}
				return ds
			}(c.Slaves),
			TraceResolverMode: true,
		}, name).SetMaxOpenConns(c.Master.MaxOpenConn).
			SetMaxIdleConns(c.Master.MaxIdleConn).
			SetConnMaxLifetime(c.Master.ConnMaxLifeTime).
			SetConnMaxIdleTime(c.Master.ConnMaxIdleTime)
	}

	db.Use(resolver)

	return db
}

func (d *repository) GetUser(ctx context.Context, uid int64) (*model.User, error) {
	var data model.User
	resolverName := data.GetDBResolverName(uid)
	db := d.orm.WithContext(ctx).Clauses(dbresolver.Use(string(resolverName)))
	err := db.Use(sharding.Register(sharding.Config{
		ShardingKey:         "id",
		NumberOfShards:      16,
		PrimaryKeyGenerator: sharding.PKSnowflake,
		ShardingAlgorithmByPrimaryKey: func(id int64) (suffix string) {
			hashed, _ := utils.Md5(fmt.Sprintf("%d", id))
			fmt.Printf("id:%d, hashed:%s\n", id, hashed)
			return fmt.Sprintf("_%s", hashed[:2])
		},
	}, model.User{}))
	if err != nil {
		return nil, errors.Wrapf(err, "[repo.user] sharding.Register err")
	}
	err = db.First(&data, uid).Error
	if errors.Is(err, ErrNotFound) {
		return nil, ErrNotFound
	} else if err != nil {
		return nil, errors.Wrapf(err, "[repo.user] query db err")
	}
	return &data, nil
}

2023/02/13 20:13:11 /Users/zhengjb/go/src/gitlab.y5ops.com/jin/hertz_demo/biz/repository/user_repo.go:30 sharding key or id required, and use operator =
[0.141ms] [rows:0] [replica] SELECT * FROM `user` WHERE `user`.`id` = 79535114761158382 ORDER BY `user`.`id` LIMIT 1

2023/02/19 21:14:12 testing mysql...
panic: interface conversion: sharding.ShardingMigrator is not migrator.BuildIndexOptionsInterface: missing method BuildIndexOptions

goroutine 1 [running]:
gorm.io/gorm/migrator.Migrator.CreateTable.func1(0x14000484380)
	/Users/zhengjb/go/src/github.com/zishiguo/playground/gorm/migrator/migrator.go:247 +0x568
gorm.io/gorm/migrator.Migrator.RunWithValue({{0x0?, 0x14000271500?, {0x105414100?, 0x1400023c000?}}}, {0x105330f80?, 0x1400023c230}, 0x140000e3540)
	/Users/zhengjb/go/src/github.com/zishiguo/playground/gorm/migrator/migrator.go:62 +0x13c
gorm.io/gorm/migrator.Migrator.CreateTable({{0x70?, 0x14000271500?, {0x105414100?, 0x1400023c000?}}}, {0x140004803c0?, 0x140004803c0?, 0x0?})
	/Users/zhengjb/go/src/github.com/zishiguo/playground/gorm/migrator/migrator.go:198 +0x114
gorm.io/gorm/migrator.Migrator.AutoMigrate({{0x0?, 0x14000270720?, {0x105414100?, 0x1400023c000?}}}, {0x1400021bad0?, 0x105311aa0?, 0x1400023c001?})
	/Users/zhengjb/go/src/github.com/zishiguo/playground/gorm/migrator/migrator.go:112 +0x244
gorm.io/sharding.ShardingMigrator.AutoMigrate({{0x105417ca8, 0x14000270690}, 0x140002368f0, {0x105413e40, 0x14000010bd0}}, {0x1400021b770?, 0x0?, 0x0?})
	/Users/zhengjb/go/pkg/mod/gorm.io/sharding@v0.5.2/dialector.go:45 +0x308
gorm.io/gorm.(*DB).AutoMigrate(0x140000017a0?, {0x1400021b770, 0x1, 0x1})
	/Users/zhengjb/go/src/github.com/zishiguo/playground/gorm/migrator.go:28 +0x44
gorm.io/playground.RunMigrations()
	/Users/zhengjb/go/src/github.com/zishiguo/playground/db.go:113 +0x194
gorm.io/playground.init.0()
	/Users/zhengjb/go/src/github.com/zishiguo/playground/db.go:37 +0x108

Process finished with the exit code 1

_ = db.Use(sharding.Register(sharding.Config{
		ShardingKey:         "ymd",
		NumberOfShards:      100,
		PrimaryKeyGenerator: sharding.PKCustom,
		ShardingAlgorithm: func(columnValue any) (suffix string, err error) {
			parentId := cast.ToInt(columnValue)
			suffix = fmt.Sprintf("_%v", parentId)
			return
		},
		PrimaryKeyGeneratorFn: func(tableIdx int64) int64 {
			return 0
		},
		// This case for show up give notifications, audit_logs table use same sharding rule.
	}, ChatMsg{}))

qb.Where(fmt.Sprintf("(MATCH(body) AGAINST('%s'))", req.Key))

type User struct{
  ID int64 `gorm:"column:id;primaryKey"`
  UserId int64 `gorm:"column:user_id"`
}

shard, err := sharding.Register(sharding.Config{
  ShardingKey:         "user_id",
  NumberOfShards:      1,
  PrimaryKeyGenerator: sharding.PKSnowflake,
  ShardingAlgorithm:   getShardFn(1),
}, "users"))
if err != nil {
  return nil, err
}

if err := db.Use(shard); err != nil {
  return nil, err
}

result := db.Create(&user) // pass pointer of data to Create

user.ID             // 0

db.Use(sharding.Register(sharding.Config{
		ShardingKey: "user_id",
		ShardingAlgorithm: func(value interface{}) (suffix string, err error) {
				if uid, ok := value.(int64); ok {
						return fmt.Sprintf("_%02d", uid%4), nil
				}
				return "", errors.New("invalid user_id")
		},
		PrimaryKeyGenerate: func(tableIdx int64) interface{} {
				// use Built in keygen for generate a sequence primary key with table index
				return keygen.Next(tableIdx)
		}
}, "orders").Register(sharding.Config{
		ShardingKey: "user_id",
		ShardingAlgorithm: func(value interface{}) (suffix string, err error) {
				if uid, ok := value.(int64); ok {
						return fmt.Sprintf("_%02d", uid%4), nil
				}
				return "", errors.New("invalid user_id")
		},
		PrimaryKeyGenerate: func(tableIdx int64) interface{} {
				return uuid.UUID()
		}
		// This case for show up give notifications, activities table use same sharding rule.
}, Notification{}, Acitivity{}))

func (s *Sharding) switchConn(db *gorm.DB) {
	// Support ignore sharding in some case, like:
	// When DoubleWrite is enabled, we need to query database schema
	// information by table name during the migration.
	if _, ok := db.Get(ShardingIgnoreStoreKey); !ok {
		s.ConnPool = &ConnPool{ConnPool: db.Statement.ConnPool, sharding: s}
		db.Statement.ConnPool = s.ConnPool
	}
}

func (s *Sharding) switchConn(db *gorm.DB) {
	// Support ignore sharding in some case, like:
	// When DoubleWrite is enabled, we need to query database schema
	// information by table name during the migration.
	if _, ok := db.Get(ShardingIgnoreStoreKey); !ok {
		s.mutex.Lock()
		if db.Statement.ConnPool != nil {
			if _, ok := db.Statement.ConnPool.(ConnPool); !ok {
				db.Statement.ConnPool = &ConnPool{ConnPool: db.Statement.ConnPool, sharding: s}
			}
		}
		s.mutex.Unlock()
	}
}

package main

import (
	"fmt"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"
	"gorm.io/sharding"
)

type Order struct {
	ID        int64 `gorm:"primarykey"`
	UserID    int64
	ProductID int64
}

func main() {
	connection := fmt.Sprintf("test:test@tcp(127.0.0.1:3306)/?parseTime=true")
	db, err := gorm.Open(mysql.New(mysql.Config{
		DSN: connection,
	}), &gorm.Config{
		NamingStrategy: schema.NamingStrategy{
			SingularTable: true,
		},
		Logger: logger.Default.LogMode(logger.Silent),
	})
	if err != nil {
		fmt.Println("err", err)
	}

	db.Exec("USE " + "MyTable")

	for i := 0; i < 4; i += 1 {
		table := fmt.Sprintf("orders_%d", i)
		db.Exec(`DROP TABLE IF EXISTS ` + table)
		db.Table(table).AutoMigrate(&Order{})
	}
	// db.AutoMigrate(&Order{})

	middleware := sharding.Register(sharding.Config{
		ShardingKey:         "user_id",
		NumberOfShards:      4,
		PrimaryKeyGenerator: sharding.PKSnowflake,
	}, "orders")
	db.Use(middleware)

	// this record will insert to orders_02
        // !!!!!!!!!!!!!!!!!!  it returm err: Error 1146: Table 'MyTable.order' doesn't exist
	// err = db.Create(&Order{UserID: 2}).Error
	err = db.Table("orders").Create(&Order{UserID: 2}).Error
	if err != nil {
		fmt.Println("err1", err)
	}

	// this record will insert to orders_03
	err = db.Exec("INSERT INTO orders(user_id) VALUES(?)", int64(3)).Error
	if err != nil {
		fmt.Println("err2", err)
	}

	// this will throw ErrMissingShardingKey error
	err = db.Exec("INSERT INTO orders(product_id) VALUES(1)").Error
	fmt.Println("err3", err)

	// this will redirect query to orders_02
	var orders []Order
        // !!!!!!!!!!!!!!!!!!  it returm err: Error 1146: Table 'MyTable.order' doesn't exist
	// err = db.Model(&Order{}).Where("user_id", int64(3)).Find(&orders).Error
	err = db.Table("orders").Where("user_id", int64(2)).Find(&orders).Error
	if err != nil {
		fmt.Println("err4", err)
	}
	fmt.Printf("%#v\n", orders)

	// Raw SQL also supported
	db.Raw("SELECT * FROM orders WHERE user_id = ?", int64(3)).Scan(&orders)
	fmt.Printf("%#v\n", orders)

	// this will throw ErrMissingShardingKey error
        // !!!!!!!!!!!!!!!!!!  it returm err: Error 1146: Table 'MyTable.order' doesn't exist
	// err = db.Model(&Order{}).Where("product_id", "1").Find(&orders).Error
	err = db.Table("orders").Where("product_id", "1").Find(&orders).Error
	fmt.Println("err5", err)

	// Update and Delete are similar to create and query
	err = db.Exec("UPDATE orders SET product_id = ? WHERE user_id = ?", 2, int64(3)).Error
	fmt.Println("err6", err) // nil
	err = db.Exec("DELETE FROM orders WHERE product_id = 3").Error
	fmt.Println("err7", err) // ErrMissingShardingKey
}