Add to support single table with flexible support for multiple partition keys.

README.md

@@ -77,6 +77,8 @@ err = db.Exec("DELETE FROM orders WHERE product_id = 3").Error
 fmt.Println(err) // ErrMissingShardingKey
 ```
 
+The example demonstrating a single table supporting multiple partitioning strategies is（单表支持多种分表策略的例子在这里）[here](./test/sharding_test.go).
+
 The full example is [here](./examples/order.go).
 
 > 🚨 NOTE: Gorm config `PrepareStmt: true` is not supported for now.

conn_pool.go

@@ -28,19 +28,26 @@ func (pool ConnPool) ExecContext(ctx context.Context, query string, args ...any)
 		curTime = time.Now()
 	)
 
-	ftQuery, stQuery, table, err := pool.sharding.resolve(query, args...)
+	ftQuery, stQuery, table, err := pool.sharding.resolve(ctx, 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 {
+		key := table
+		key, err = pool.sharding.getConfigKey(ctx, table)
+		if err != nil {
+			return nil, err
+		}
+		if r, ok := pool.sharding.configs[key]; ok {
 			if r.DoubleWrite {
 				pool.sharding.Logger.Trace(ctx, curTime, func() (sql string, rowsAffected int64) {
 					result, _ := pool.ConnPool.ExecContext(ctx, ftQuery, args...)
-					rowsAffected, _ = result.RowsAffected()
+					if result != nil {
+						rowsAffected, _ = result.RowsAffected()
+					}
 					return pool.sharding.Explain(ftQuery, args...), rowsAffected
 				}, pool.sharding.Error)
 			}
@@ -50,7 +57,9 @@ func (pool ConnPool) ExecContext(ctx context.Context, query string, args ...any)
 	var result sql.Result
 	result, err = pool.ConnPool.ExecContext(ctx, stQuery, args...)
 	pool.sharding.Logger.Trace(ctx, curTime, func() (sql string, rowsAffected int64) {
-		rowsAffected, _ = result.RowsAffected()
+		if result != nil {
+			rowsAffected, _ = result.RowsAffected()
+		}
 		return pool.sharding.Explain(stQuery, args...), rowsAffected
 	}, pool.sharding.Error)
 
@@ -63,7 +72,7 @@ func (pool ConnPool) QueryContext(ctx context.Context, query string, args ...any
 		curTime = time.Now()
 	)
 
-	_, stQuery, _, err := pool.sharding.resolve(query, args...)
+	_, stQuery, _, err := pool.sharding.resolve(ctx, query, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -80,7 +89,7 @@ func (pool ConnPool) QueryContext(ctx context.Context, query string, args ...any
 }
 
 func (pool ConnPool) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row {
-	_, query, _, _ = pool.sharding.resolve(query, args...)
+	_, query, _, _ = pool.sharding.resolve(ctx, query, args...)
 	pool.sharding.querys.Store("last_query", query)
 
 	return pool.ConnPool.QueryRowContext(ctx, query, args...)

sharding.go

@@ -1,6 +1,7 @@
 package sharding
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"hash/crc32"
@@ -15,15 +16,20 @@ import (
 )
 
 var (
-	ErrMissingShardingKey = errors.New("sharding key or id required, and use operator =")
-	ErrInvalidID          = errors.New("invalid id format")
-	ErrInsertDiffSuffix   = errors.New("can not insert different suffix table in one query ")
+	ErrMissingShardingKey           = errors.New("sharding key or id required, and use operator =")
+	ErrInvalidID                    = errors.New("invalid id format")
+	ErrInsertDiffSuffix             = errors.New("can not insert different suffix table in one query ")
+	ErrShardingKeyNotExistInContext = errors.New("the value passed in the context is not the sharding key")
+	ErrMissingTableName             = errors.New("table name is required")
 )
 
 var (
 	ShardingIgnoreStoreKey = "sharding_ignore"
 )
 
+// ContextKeyForShardingKey is the context key for sharding key.
+const ContextKeyForShardingKey = "sharding_key"
+
 type Sharding struct {
 	*gorm.DB
 	ConnPool       *ConnPool
@@ -46,6 +52,10 @@ type Config struct {
 	// For example, for a product order table, you may want to split the rows by `user_id`.
 	ShardingKey string
 
+	// logical table name.Suport multiple table names with same sharding key.
+	// For example, for user and order table, you may want to shard by `user_id`.
+	TableNames []string
+
 	// NumberOfShards specifies how many tables you want to sharding.
 	NumberOfShards uint
 
@@ -110,6 +120,56 @@ func Register(config Config, tables ...any) *Sharding {
 	}
 }
 
+// enables sharding for a single table with flexible support for multiple partition keys.
+func RegisterWithKeys(configs []Config) (*Sharding, error) {
+	mapConfig := make(map[string]Config, len(configs))
+	for _, config := range configs {
+		for _, tableName := range config.TableNames {
+			configKey, err := generateConfigsKey(tableName, config.ShardingKey)
+			if err != nil {
+				return nil, err
+			}
+			mapConfig[configKey] = config
+		}
+	}
+	return &Sharding{
+		configs: mapConfig,
+	}, nil
+}
+
+// generates the key for the sharding config.
+func generateConfigsKey(tableName, shardingKey string) (string, error) {
+	// Table name cannot be empty
+	if tableName == "" {
+		return "", ErrMissingTableName
+	}
+	if shardingKey == "" {
+		return "", ErrMissingShardingKey
+	}
+	return fmt.Sprintf("%s_%s", tableName, shardingKey), nil
+}
+
+// get the configs key for using it to get the sharding config.
+func (s *Sharding) getConfigKey(ctx context.Context, tableName string) (string, error) {
+	configKey := tableName
+	if shardingKey, ok := ctx.Value(ContextKeyForShardingKey).(string); ok {
+		// If sharding key is set in context, use it to get the sharding config.
+		configKey = fmt.Sprintf("%s_%s", tableName, shardingKey)
+	} else {
+		// If sharding key is not set in context, use the table name as the key.
+		return configKey, nil
+	}
+
+	// check if the sharding key exists in the configs.
+	_, exis := s.configs[configKey]
+	if !exis {
+		return "", ErrShardingKeyNotExistInContext
+	}
+
+	// If sharding key is not set in context, use the table name as the key.
+	return configKey, nil
+}
+
 func (s *Sharding) compile() error {
 	if s.configs == nil {
 		s.configs = make(map[string]Config)
@@ -297,7 +357,7 @@ func (s *Sharding) switchConn(db *gorm.DB) {
 }
 
 // resolve split the old query to full table query and sharding table query
-func (s *Sharding) resolve(query string, args ...any) (ftQuery, stQuery, tableName string, err error) {
+func (s *Sharding) resolve(ctx context.Context, query string, args ...any) (ftQuery, stQuery, tableName string, err error) {
 	ftQuery = query
 	stQuery = query
 	if len(s.configs) == 0 {
@@ -344,7 +404,12 @@ func (s *Sharding) resolve(query string, args ...any) (ftQuery, stQuery, tableNa
 	}
 
 	tableName = table.Name.Name
-	r, ok := s.configs[tableName]
+	key := tableName
+	key, err = s.getConfigKey(ctx, tableName)
+	if err != nil {
+		return
+	}
+	r, ok := s.configs[key]
 	if !ok {
 		return
 	}