[Feature]Support drds createtable SQL statement

go.mod

@@ -4,7 +4,7 @@ go 1.20
 
 require (
 	github.com/appleboy/gin-jwt/v2 v2.9.1
-	github.com/arana-db/parser v0.2.16
+	github.com/arana-db/parser v0.2.17
 	github.com/blang/semver v3.5.1+incompatible
 	github.com/bwmarrin/snowflake v0.3.0
 	github.com/cespare/xxhash/v2 v2.2.0

go.sum

@@ -107,8 +107,8 @@ github.com/appleboy/gin-jwt/v2 v2.9.1 h1:l29et8iLW6omcHltsOP6LLk4s3v4g2FbFs0koxG
 github.com/appleboy/gin-jwt/v2 v2.9.1/go.mod h1:jwcPZJ92uoC9nOUTOKWoN/f6JZOgMSKlFSHw5/FrRUk=
 github.com/appleboy/gofight/v2 v2.1.2 h1:VOy3jow4vIK8BRQJoC/I9muxyYlJ2yb9ht2hZoS3rf4=
 github.com/appleboy/gofight/v2 v2.1.2/go.mod h1:frW+U1QZEdDgixycTj4CygQ48yLTUhplt43+Wczp3rw=
-github.com/arana-db/parser v0.2.16 h1:tBUny7dWSwvINTzdxPdYX7BGZ9TWmJWn9juJL80TSJs=
-github.com/arana-db/parser v0.2.16/go.mod h1:/XA29bplweWSEAjgoM557ZCzhBilSawUlHcZFjOeDAc=
+github.com/arana-db/parser v0.2.17 h1:4wNfSgza2N3pjpwR5jmWLvu4L6Sme6EtoLuZOgwWlsU=
+github.com/arana-db/parser v0.2.17/go.mod h1:/XA29bplweWSEAjgoM557ZCzhBilSawUlHcZFjOeDAc=
 github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=
 github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
 github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY=

pkg/boot/discovery.go

@@ -337,7 +337,7 @@ func (fp *discovery) GetTable(ctx context.Context, tenant, cluster, tableName st
 		return nil, nil
 	}
 
-	return makeVTable(tableName, table)
+	return config.MakeVTable(tableName, table)
 }
 
 func (fp *discovery) loadCluster(tenant, cluster string) (*config.DataSourceCluster, error) {

pkg/boot/watcher.go

@@ -320,7 +320,7 @@ func (d *watcher) onGroupAdd(ctx context.Context, cluster string, group *config.
 		}
 
 		if err := rt.Namespace().EnqueueCommand(namespace.UpsertDB(group.Name, runtime.NewAtomDB(node))); err != nil {
-			log.Errorf("[%s] failed to enqueue upsert-db command: %v", err)
+			log.Errorf("[%s] failed to enqueue upsert-db command: %v", d.tenant, err)
 		}
 	}
 
@@ -494,7 +494,7 @@ func (d *watcher) onNodeUpdate(ctx context.Context, node *config.Node) error {
 }
 
 func (d *watcher) onTableDel(_ context.Context, name string) error {
-	db, tbl, err := parseDatabaseAndTable(name)
+	db, tbl, err := config.ParseDatabaseAndTable(name)
 	if err != nil {
 		return errors.WithStack(err)
 	}
@@ -522,7 +522,7 @@ func (d *watcher) onTableDel(_ context.Context, name string) error {
 }
 
 func (d *watcher) onTableAdd(_ context.Context, table *config.Table) error {
-	db, tbl, err := parseDatabaseAndTable(table.Name)
+	db, tbl, err := config.ParseDatabaseAndTable(table.Name)
 	if err != nil {
 		return errors.WithStack(err)
 	}
@@ -537,7 +537,7 @@ func (d *watcher) onTableAdd(_ context.Context, table *config.Table) error {
 		return nil
 	}
 
-	vtab, err := makeVTable(tbl, table)
+	vtab, err := config.MakeVTable(tbl, table)
 	if err != nil {
 		return errors.WithStack(err)
 	}
@@ -547,7 +547,7 @@ func (d *watcher) onTableAdd(_ context.Context, table *config.Table) error {
 }
 
 func (d *watcher) onTableChange(_ context.Context, table *config.Table) error {
-	db, tbl, err := parseDatabaseAndTable(table.Name)
+	db, tbl, err := config.ParseDatabaseAndTable(table.Name)
 	if err != nil {
 		return errors.WithStack(err)
 	}
@@ -561,7 +561,7 @@ func (d *watcher) onTableChange(_ context.Context, table *config.Table) error {
 		return nil
 	}
 
-	vtab, err := makeVTable(tbl, table)
+	vtab, err := config.MakeVTable(tbl, table)
 	if err != nil {
 		return errors.WithStack(err)
 	}

pkg/boot/misc.go → pkg/config/boot_misc.go

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package boot
+package config
 
 import (
 	"fmt"
@@ -30,7 +30,6 @@ import (
 )
 
 import (
-	"github.com/arana-db/arana/pkg/config"
 	"github.com/arana-db/arana/pkg/proto/rule"
 	"github.com/arana-db/arana/pkg/util/math"
 )
@@ -47,7 +46,7 @@ func getTopologyRegexp() *regexp.Regexp {
 	return _regexpTopology
 }
 
-func parseTopology(input string) (format string, begin, end int, err error) {
+func ParseTopology(input string) (format string, begin, end int, err error) {
 	mats := getTopologyRegexp().FindAllStringSubmatch(input, -1)
 
 	if len(mats) < 1 {
@@ -83,7 +82,7 @@ func parseTopology(input string) (format string, begin, end int, err error) {
 	return
 }
 
-func makeVTable(tableName string, table *config.Table) (*rule.VTable, error) {
+func MakeVTable(tableName string, table *Table) (*rule.VTable, error) {
 	var (
 		vt                 rule.VTable
 		topology           rule.Topology
@@ -95,12 +94,12 @@ func makeVTable(tableName string, table *config.Table) (*rule.VTable, error) {
 
 	if table.Topology != nil {
 		if len(table.Topology.DbPattern) > 0 {
-			if dbFormat, dbBegin, dbEnd, err = parseTopology(table.Topology.DbPattern); err != nil {
+			if dbFormat, dbBegin, dbEnd, err = ParseTopology(table.Topology.DbPattern); err != nil {
 				return nil, errors.WithStack(err)
 			}
 		}
 		if len(table.Topology.TblPattern) > 0 {
-			if tbFormat, tbBegin, tbEnd, err = parseTopology(table.Topology.TblPattern); err != nil {
+			if tbFormat, tbBegin, tbEnd, err = ParseTopology(table.Topology.TblPattern); err != nil {
 				return nil, errors.WithStack(err)
 			}
 		}
@@ -170,7 +169,7 @@ var (
 	_fullTableNameRegexpOnce sync.Once
 )
 
-func parseDatabaseAndTable(name string) (db, tbl string, err error) {
+func ParseDatabaseAndTable(name string) (db, tbl string, err error) {
 	_fullTableNameRegexpOnce.Do(func() {
 		_fullTableNameRegexp = regexp.MustCompile(`^([a-zA-Z_][a-zA-Z0-9_-]+)\.([a-zA-Z_][a-zA-Z0-9_-]+)$`)
 	})
@@ -186,7 +185,7 @@ func parseDatabaseAndTable(name string) (db, tbl string, err error) {
 	return
 }
 
-func toSharder(input *config.Rule) (rule.ShardComputer, error) {
+func toSharder(input *Rule) (rule.ShardComputer, error) {
 	columns := make([]string, 0, len(input.Columns))
 	for i := range input.Columns {
 		columns = append(columns, input.Columns[i].Name)
@@ -205,8 +204,8 @@ func getRender(format string) func(int) string {
 	}
 }
 
-func toShardMetadata(rules []*config.Rule, defaultSteps int) ([]*rule.ShardMetadata, error) {
-	toShardColumn := func(ru *config.ColumnRule, dst *[]*rule.ShardColumn, defaultSteps int) {
+func toShardMetadata(rules []*Rule, defaultSteps int) ([]*rule.ShardMetadata, error) {
+	toShardColumn := func(ru *ColumnRule, dst *[]*rule.ShardColumn, defaultSteps int) {
 		unit := rule.Unum
 		switch strings.ToLower(ru.Type) {
 		case "string", "str":

pkg/boot/misc_test.go → pkg/config/boot_misc_test.go

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package boot
+package config
 
 import (
 	"testing"
@@ -36,7 +36,7 @@ func TestParseDatabaseAndTable(t *testing.T) {
 		{"fake-db.fake_table", "fake-db", "fake_table"},
 	} {
 		t.Run(it.name, func(t *testing.T) {
-			db, table, err := parseDatabaseAndTable(it.name)
+			db, table, err := ParseDatabaseAndTable(it.name)
 			assert.NoError(t, err)
 			assert.Equal(t, it.db, db)
 			assert.Equal(t, it.table, table)

pkg/runtime/ast/ast.go

@@ -1754,6 +1754,7 @@ func (cc *convCtx) convCreateTableStmt(stmt *ast.CreateTableStmt) Statement {
 		Cols:        stmt.Cols,
 		Constraints: stmt.Constraints,
 		Options:     stmt.Options,
+		Partition:   stmt.Partition,
 	}
 }
 

pkg/runtime/ast/create_table.go

@@ -42,7 +42,7 @@ type CreateTableStmt struct {
 	Cols        []*ast.ColumnDef
 	Constraints []*ast.Constraint
 	Options     []*ast.TableOption
-	// Partition      *PartitionOptions
+	Partition   *ast.PartitionOptions
 	// OnDuplicate    OnDuplicateKeyHandlingType
 	// Select         ResultSetNode
 }

pkg/runtime/optimize/ddl/create_table.go

@@ -19,16 +19,21 @@ package ddl
 
 import (
 	"context"
+	"fmt"
+	"strings"
 )
 
 import (
+	"github.com/arana-db/arana/pkg/admin"
+	"github.com/arana-db/arana/pkg/config"
 	"github.com/arana-db/arana/pkg/proto"
 	"github.com/arana-db/arana/pkg/proto/rule"
 	"github.com/arana-db/arana/pkg/runtime/ast"
 	"github.com/arana-db/arana/pkg/runtime/optimize"
 	"github.com/arana-db/arana/pkg/runtime/plan/ddl"
 	"github.com/arana-db/arana/pkg/runtime/plan/dml"
 	"github.com/arana-db/arana/pkg/util/log"
+	parse "github.com/arana-db/parser/ast"
 )
 
 func init() {
@@ -44,6 +49,12 @@ func optimizeCreateTable(ctx context.Context, o *optimize.Optimizer) (proto.Plan
 	)
 	vt, ok := o.Rule.VTable(stmt.Table.Suffix())
 	fullScan = ok
+	if !fullScan {
+		if stmt.Partition != nil && stmt.Partition.PartitionFrom == 1 {
+			vt, ok = drdsCreateTable(ctx, o)
+			fullScan = ok
+		}
+	}
 
 	log.Debugf("compute shards: result=%s, isFullScan=%v", shards, fullScan)
 
@@ -83,3 +94,78 @@ func optimizeCreateTable(ctx context.Context, o *optimize.Optimizer) (proto.Plan
 
 	return tmpPlan, nil
 }
+
+func drdsCreateTable(ctx context.Context, o *optimize.Optimizer) (*rule.VTable, bool) {
+	stmt := o.Stmt.(*ast.CreateTableStmt)
+
+	dbName := ctx.Value(proto.ContextKeySchema{}).(string)
+	tbName := stmt.Table.Suffix()
+
+	dbColume := ""
+	switch node := stmt.Partition.PartitionMethod.Expr.(type) {
+	case *parse.ColumnNameExpr:
+		dbColume = parse.ColumnNameExpr(*node).Name.Name.O
+	}
+	tbColume := ""
+	switch node := stmt.Partition.Sub.Expr.(type) {
+	case *parse.ColumnNameExpr:
+		tbColume = parse.ColumnNameExpr(*node).Name.Name.O
+	}
+	if strings.Compare(dbColume, "") == 0 || strings.Compare(tbColume, "") == 0 {
+		return nil, false
+	}
+
+	newTb := &config.Table{
+		Name: dbName + "." + tbName,
+		Sequence: &config.Sequence{
+			Type: "snowflake",
+		},
+		DbRules: []*config.Rule{
+			{
+				Columns: []*config.ColumnRule{{Name: dbColume}},
+				Expr: strings.ReplaceAll(dbColume, dbColume, "$0") + " % " +
+					fmt.Sprintf("%d", stmt.Partition.Num),
+			},
+		},
+		TblRules: []*config.Rule{
+			{
+				Columns: []*config.ColumnRule{{Name: tbColume}},
+				Expr: strings.ReplaceAll(tbColume, tbColume, "$0") + " % " +
+					fmt.Sprintf("%d", stmt.Partition.Sub.Num),
+			},
+		},
+		Topology: &config.Topology{
+			DbPattern:  dbName + fmt.Sprintf("_${0000..%04d}", stmt.Partition.Num-1),
+			TblPattern: tbName + fmt.Sprintf("_${0000..%04d}", stmt.Partition.Num*stmt.Partition.Sub.Num-1),
+		},
+		Attributes: map[string]string{
+			"allow_full_scan": "true",
+			"sqlMaxLimit":     "-1",
+		},
+	}
+	newVt, err := config.MakeVTable(tbName, newTb)
+	if err == nil {
+		o.Rule.SetVTable(tbName, newVt)
+
+		tenant := ctx.Value(proto.ContextKeyTenant{}).(string)
+		op, _ := config.NewTenantOperator(config.GetStoreOperate())
+		srv := &admin.MyConfigService{
+			TenantOp: op,
+		}
+		tbDTO := &admin.TableDTO{
+			Name:           newTb.Name,
+			Sequence:       newTb.Sequence,
+			DbRules:        newTb.DbRules,
+			TblRules:       newTb.TblRules,
+			Topology:       newTb.Topology,
+			ShadowTopology: newTb.ShadowTopology,
+			Attributes:     newTb.Attributes,
+		}
+		e := srv.UpsertTable(ctx, tenant, dbName, tbName, tbDTO)
+		if e == nil {
+			return newVt, true
+		}
+	}
+
+	return nil, false
+}