From fd0c134c8a4a43664ef3add188fb9586008ddf85 Mon Sep 17 00:00:00 2001 From: csynineyang Date: Fri, 29 Dec 2023 17:12:15 +0800 Subject: [PATCH 1/4] Support drds createtable SQL statement --- pkg/boot/discovery.go | 2 +- pkg/boot/watcher.go | 12 +-- pkg/{boot/misc.go => config/boot_misc.go} | 26 +++--- .../misc_test.go => config/boot_misc_test.go} | 6 +- pkg/runtime/ast/ast.go | 1 + pkg/runtime/ast/create_table.go | 2 +- pkg/runtime/optimize/ddl/create_table.go | 86 +++++++++++++++++++ 7 files changed, 107 insertions(+), 28 deletions(-) rename pkg/{boot/misc.go => config/boot_misc.go} (89%) rename pkg/{boot/misc_test.go => config/boot_misc_test.go} (94%) diff --git a/pkg/boot/discovery.go b/pkg/boot/discovery.go index 0ba794b6..6becca7e 100644 --- a/pkg/boot/discovery.go +++ b/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) { diff --git a/pkg/boot/watcher.go b/pkg/boot/watcher.go index 80467187..5bc698fe 100644 --- a/pkg/boot/watcher.go +++ b/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) } diff --git a/pkg/boot/misc.go b/pkg/config/boot_misc.go similarity index 89% rename from pkg/boot/misc.go rename to pkg/config/boot_misc.go index 0886b73b..6df6305f 100644 --- a/pkg/boot/misc.go +++ b/pkg/config/boot_misc.go @@ -15,7 +15,7 @@ * limitations under the License. */ -package boot +package config import ( "fmt" @@ -23,16 +23,10 @@ import ( "strconv" "strings" "sync" -) -import ( - "github.com/pkg/errors" -) - -import ( - "github.com/arana-db/arana/pkg/config" "github.com/arana-db/arana/pkg/proto/rule" "github.com/arana-db/arana/pkg/util/math" + "github.com/pkg/errors" ) var ( @@ -47,7 +41,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 +77,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 +89,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 +164,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 +180,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 +199,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": diff --git a/pkg/boot/misc_test.go b/pkg/config/boot_misc_test.go similarity index 94% rename from pkg/boot/misc_test.go rename to pkg/config/boot_misc_test.go index b70d68ba..b5616344 100644 --- a/pkg/boot/misc_test.go +++ b/pkg/config/boot_misc_test.go @@ -15,13 +15,11 @@ * limitations under the License. */ -package boot +package config import ( "testing" -) -import ( "github.com/stretchr/testify/assert" ) @@ -36,7 +34,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) diff --git a/pkg/runtime/ast/ast.go b/pkg/runtime/ast/ast.go index 8eaa693a..e353ce92 100644 --- a/pkg/runtime/ast/ast.go +++ b/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, } } diff --git a/pkg/runtime/ast/create_table.go b/pkg/runtime/ast/create_table.go index c8315624..3648d2f7 100644 --- a/pkg/runtime/ast/create_table.go +++ b/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 } diff --git a/pkg/runtime/optimize/ddl/create_table.go b/pkg/runtime/optimize/ddl/create_table.go index 81e1ef97..f0ad0ba2 100644 --- a/pkg/runtime/optimize/ddl/create_table.go +++ b/pkg/runtime/optimize/ddl/create_table.go @@ -19,9 +19,13 @@ 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" @@ -29,6 +33,7 @@ import ( "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 +} From 71ee0830942a3441ae017a215be4f4b8cf7bf106 Mon Sep 17 00:00:00 2001 From: csynineyang Date: Mon, 8 Jan 2024 16:27:54 +0800 Subject: [PATCH 2/4] fix: import arana-db/parser v0.2.17 --- go.mod | 2 +- go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index c023d0f8..9ac929d3 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 56c965e8..e5f3f293 100644 --- a/go.sum +++ b/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= From b3444634328c58b7f7d7ed559e8fc128db97d13d Mon Sep 17 00:00:00 2001 From: csynineyang Date: Mon, 8 Jan 2024 17:13:11 +0800 Subject: [PATCH 3/4] fix: format style --- pkg/config/boot_misc.go | 8 +++++++- pkg/config/boot_misc_test.go | 2 ++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/pkg/config/boot_misc.go b/pkg/config/boot_misc.go index 6df6305f..621bd080 100644 --- a/pkg/config/boot_misc.go +++ b/pkg/config/boot_misc.go @@ -23,10 +23,16 @@ import ( "strconv" "strings" "sync" +) +import ( + "github.com/pkg/errors" +) + +import ( + "github.com/arana-db/arana/pkg/config" "github.com/arana-db/arana/pkg/proto/rule" "github.com/arana-db/arana/pkg/util/math" - "github.com/pkg/errors" ) var ( diff --git a/pkg/config/boot_misc_test.go b/pkg/config/boot_misc_test.go index b5616344..a85afdee 100644 --- a/pkg/config/boot_misc_test.go +++ b/pkg/config/boot_misc_test.go @@ -19,7 +19,9 @@ package config import ( "testing" +) +import ( "github.com/stretchr/testify/assert" ) From c9968079366bdb9d2d8230c9a020ac232c7384b4 Mon Sep 17 00:00:00 2001 From: csynineyang Date: Mon, 8 Jan 2024 17:19:46 +0800 Subject: [PATCH 4/4] fix: format style --- pkg/config/boot_misc.go | 1 - 1 file changed, 1 deletion(-) diff --git a/pkg/config/boot_misc.go b/pkg/config/boot_misc.go index 621bd080..2ec1c2d0 100644 --- a/pkg/config/boot_misc.go +++ b/pkg/config/boot_misc.go @@ -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" )