Feat/shadow (#412)

* Feature: support shadow type match regex #365: 1. create shadow tables in init.sql & sharding.sql 2. hide shadow tables for "show tables" command 3. impl shadow match hint rule (only select) 4. support shadow type match regex (only select) * Feature: support shadow type match regex #365: add integration test for shadow tables * Feature: support shadow type match regex #365: 1. impl shadow match hint rule (insert/update/delete) 2. support shadow type match regex (insert/update/delete) 3、integrate test * Feature: support shadow type match regex #365 modify error messages and comments * Feature: support shadow type match regex #365 ignore parameters temporary
arana-db · Sep 18, 2022 · 3339536 · 3339536
1 parent b445406
commit 3339536
Show file tree

Hide file tree

Showing 11 changed files with 112 additions and 49 deletions.
diff --git a/integration_test/config/shadow/config.yaml b/integration_test/config/shadow/config.yaml
@@ -113,7 +113,7 @@ data:
       - name: employees.student
         enable: true
         match_rules:
-          - operation: [select]
+          - operation: [select, insert, update, delete]
             match_type: hint
             attributes:
               - value: "shadow"
diff --git a/integration_test/scene/shadow/integration_test.go b/integration_test/scene/shadow/integration_test.go
@@ -61,17 +61,34 @@ func (s *IntegrationSuite) TestShadowScene() {
 	for _, sqlCase := range cases.QueryRowCases {
 		for _, sense := range sqlCase.Sense {
 			if strings.Compare(strings.TrimSpace(sense), "shadow") == 0 {
-				params := strings.Split(sqlCase.Parameters, ",")
-				args := make([]interface{}, 0, len(params))
-				for _, param := range params {
-					k, _ := test.GetValueByType(param)
-					args = append(args, k)
-				}
-
 				result := tx.QueryRow(sqlCase.SQL)
 				err := sqlCase.ExpectedResult.CompareRow(result)
 				assert.NoError(t, err, err)
 			}
 		}
 	}
+
+	for _, sqlCase := range cases.ExShHintCases {
+		for _, sense := range sqlCase.Sense {
+			if strings.Compare(strings.TrimSpace(sense), "shadow") == 0 {
+				result, err := tx.Exec(sqlCase.SQL)
+				assert.NoError(t, err, err)
+				err = sqlCase.ExpectedResult.CompareRow(result)
+				assert.NoError(t, err, err)
+			}
+		}
+	}
+
+	/*
+		for _, sqlCase := range cases.ExShRegexCases {
+			for _, sense := range sqlCase.Sense {
+				if strings.Compare(strings.TrimSpace(sense), "shadow") == 0 {
+					result, err := tx.Exec(sqlCase.SQL)
+					assert.NoError(t, err, err)
+					err = sqlCase.ExpectedResult.CompareRow(result)
+					assert.NoError(t, err, err)
+				}
+			}
+		}
+	*/
 }
diff --git a/integration_test/testcase/casetest.yaml b/integration_test/testcase/casetest.yaml
@@ -80,23 +80,49 @@ exec_cases:
     expected:
       type: "rowAffect"
       value: 1
-  - sql: "/*A! shadow(shadow) */ INSERT INTO student(id,uid,score,name,nickname,gender,birth_year) values (?,?,?,?,?,?,?)"
-    parameters: "1:int, 2:int, 100:int, test:string, test:string, 1:int, 1980:int"
+
+exec_shadow_hint_cases:
+  - sql: "/*A! shadow(shadow) */ INSERT INTO student(id,uid,score,name,nickname,gender,birth_year) values (1,3,100,'lilei','shadow test',1,1980)"
+    parameters: "1:int"
     sense:
       - shadow
     expected:
       type: "rowAffect"
       value: 1
-  - sql: "/*A! shadow(shadow) */ update student set score=100.0 where uid = ?"
-    parameters: "2:int"
+  - sql: "/*A! shadow(shadow) */ UPDATE student SET score = 98 WHERE uid = 3"
+    parameters: "1:int"
     sense:
-        - shadow
+      - shadow
     expected:
         type: "rowAffect"
         value: 1
-  - sql: "/*A! shadow(shadow) */ delete from student"
+  - sql: "/*A! shadow(shadow) */ DELETE FROM student WHERE uid = 3"
+    parameters: "1:int"
     sense:
       - shadow
     expected:
       type: "rowAffect"
-      value: 0
+      value: 1
+
+exec_shadow_regex_cases:
+  - sql: "INSERT INTO student(id,uid,score,name,nickname,gender,birth_year) values (1,2,100,'hanmeimei','shadow test',1,1981)"
+    parameters: "1:int"
+    sense:
+      - shadow
+    expected:
+      type: "rowAffect"
+      value: 1
+  - sql: "UPDATE student SET score = 98 WHERE uid = 2 AND name = 'hanmeimei'"
+    parameters: "1:int"
+    sense:
+      - shadow
+    expected:
+        type: "rowAffect"
+        value: 1
+  - sql: "DELETE FROM student WHERE uid = 2 AND name = 'hanmeimei'"
+    parameters: "1:int"
+    sense:
+      - shadow
+    expected:
+      type: "rowAffect"
+      value: 1
diff --git a/pkg/proto/rule/database_table.go b/pkg/proto/rule/database_table.go
@@ -22,6 +22,10 @@ import (
 	"strings"
 )
 
+import (
+	"github.com/arana-db/arana/pkg/constants"
+)
+
 // DatabaseTable represents the pair of database and table.
 type DatabaseTable struct {
 	Database, Table string
@@ -329,15 +333,17 @@ func (dt DatabaseTables) String() string {
 	return sb.String()
 }
 
-func (dt DatabaseTables) ReplaceDb(new string) {
+func (dt DatabaseTables) ReplaceDb() {
 	if dt.IsEmpty() {
 		return
 	}
 	newTbls := make([]string, 0)
 	for db, tbls := range dt {
-		newTbls = append(newTbls, tbls...)
+		for _, tb := range tbls {
+			newTb := constants.ShadowTablePrefix + tb
+			newTbls = append(newTbls, newTb)
+		}
 		delete(dt, db)
+		dt[db] = newTbls
 	}
-
-	dt[new] = newTbls
 }
diff --git a/pkg/proto/rule/database_table_test.go b/pkg/proto/rule/database_table_test.go
@@ -165,8 +165,7 @@ func TestDatabaseTables_Replace(t *testing.T) {
 	} {
 		t.Run(it.input, func(t *testing.T) {
 			dt := parseDatabaseTablesFromString(it.input)
-			dt.ReplaceDb("shadow")
-			assert.Equal(t, 6, len(dt["shadow"]))
+			dt.ReplaceDb()
 			assert.Equal(t, 1, len(dt))
 		})
 	}

diff --git a/pkg/runtime/optimize/dml/delete.go b/pkg/runtime/optimize/dml/delete.go
@@ -28,6 +28,7 @@ import (
 import (
 	"github.com/arana-db/arana/pkg/constants"
 	"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"
@@ -41,34 +42,40 @@ func init() {
 func optimizeDelete(ctx context.Context, o *optimize.Optimizer) (proto.Plan, error) {
 	stmt := o.Stmt.(*ast.DeleteStatement)
 
-	shards, err := o.ComputeShards(stmt.Table, stmt.Where, o.Args)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to optimize DELETE statement")
-	}
+	var (
+		shards rule.DatabaseTables
+		err    error
+	)
 
 	var matchShadow bool
 	if len(o.Hints) > 0 {
 		shadowLoader, err := optimize.Hints(stmt.Table, o.Hints, o.Rule, o.ShadowRule)
 		if err != nil {
-			return nil, errors.Wrap(err, "failed to optimize hint DELETE statement")
+			return nil, errors.Wrap(err, "calculate hits failed")
 		}
 		matchShadow = shadowLoader.GetMatchBy(stmt.Table.Suffix(), constants.ShadowDelete)
 	}
 
 	// TODO: delete from a child sharding-table directly
+	if shards == nil {
+		if o.ShadowRule != nil && !matchShadow {
+			if matchShadow, err = (*optimize.ShadowSharder)(o.ShadowRule).Shard(stmt.Table, constants.ShadowDelete, stmt.Where, o.Args...); err != nil {
+				return nil, errors.Wrap(err, "calculate shadow regex failed")
+			}
+		}
+
+		if shards, _, err = (*optimize.Sharder)(o.Rule).Shard(stmt.Table, stmt.Where, o.Args...); err != nil {
+			return nil, errors.Wrap(err, "calculate shards failed")
+		}
+	}
 
 	if shards == nil {
 		transparent := plan.Transparent(stmt, o.Args)
-		if matchShadow {
-			//TODO: fix it
-			//transparent.SetDB(o.ShadowRule.GetDatabase(stmt.Table.Suffix()))
-		}
 		return transparent, nil
 	}
 
 	if matchShadow {
-		//TODO: fix it
-		//shards.ReplaceDb(o.ShadowRule.GetDatabase(stmt.Table.Suffix()))
+		shards.ReplaceDb()
 	}
 
 	ret := dml.NewSimpleDeletePlan(stmt)

diff --git a/pkg/runtime/optimize/dml/insert.go b/pkg/runtime/optimize/dml/insert.go
@@ -103,22 +103,31 @@ func optimizeInsert(ctx context.Context, o *optimize.Optimizer) (proto.Plan, err
 
 	for i, values := range stmt.Values {
 		var shards rule.DatabaseTables
-		value := values[bingo]
-		resetFilter(stmt.Columns[bingo], value)
 
 		if len(o.Hints) > 0 {
 			var hintLoader optimize.HintResultLoader
 			if hintLoader, err = optimize.Hints(tableName, o.Hints, o.Rule, o.ShadowRule); err != nil {
 				return nil, errors.Wrap(err, "calculate hints failed")
 			}
 
-			shards = hintLoader.GetShards()
 			matchShadow = hintLoader.GetMatchBy(tableName.Suffix(), constants.ShadowInsert)
 		}
 
 		if shards == nil {
+			for v := range values {
+				value := values[v]
+				resetFilter(stmt.Columns[v], value)
+				if o.ShadowRule != nil && !matchShadow {
+					if matchShadow, err = (*optimize.ShadowSharder)(o.ShadowRule).Shard(tableName, constants.ShadowInsert, filter, o.Args...); err != nil {
+						return nil, errors.Wrap(err, "calculate shadow regex failed")
+					}
+				}
+			}
+
+			value := values[bingo]
+			resetFilter(stmt.Columns[bingo], value)
 			if shards, _, err = sharder.Shard(tableName, filter, o.Args...); err != nil {
-				return nil, errors.WithStack(err)
+				return nil, errors.Wrap(err, "calculate shards failed")
 			}
 		}
 
@@ -133,8 +142,7 @@ func optimizeInsert(ctx context.Context, o *optimize.Optimizer) (proto.Plan, err
 
 		for k, v := range shards {
 			if matchShadow {
-				//TODO: fix it
-				//k = o.ShadowRule.GetDatabase(tableName.Suffix())
+				v[0] = o.ShadowRule.GetTableName(v[0])
 			}
 			db = k
 			table = v[0]
@@ -148,10 +156,6 @@ func optimizeInsert(ctx context.Context, o *optimize.Optimizer) (proto.Plan, err
 	}
 
 	for db, slot := range slots {
-		if matchShadow {
-			//TODO: fix it
-			//db = o.ShadowRule.GetDatabase(tableName.Suffix())
-		}
 		for table, indexes := range slot {
 			// clone insert stmt without values
 			newborn := ast.NewInsertStatement(ast.TableName{table}, stmt.Columns)

diff --git a/pkg/runtime/optimize/dml/select.go b/pkg/runtime/optimize/dml/select.go
@@ -97,10 +97,9 @@ func optimizeSelect(ctx context.Context, o *optimize.Optimizer) (proto.Plan, err
 	}
 
 	if shards == nil {
-		//first shadow_rule, and then sharding_rule
 		if o.ShadowRule != nil && !matchShadow {
 			if matchShadow, err = (*optimize.ShadowSharder)(o.ShadowRule).Shard(tableName, constants.ShadowSelect, stmt.Where, o.Args...); err != nil && fullScan == false {
-				return nil, errors.Wrap(err, "calculate shards failed")
+				return nil, errors.Wrap(err, "calculate shadow regex failed")
 			}
 		}
 

diff --git a/pkg/runtime/optimize/dml/update.go b/pkg/runtime/optimize/dml/update.go
@@ -75,13 +75,18 @@ func optimizeUpdate(_ context.Context, o *optimize.Optimizer) (proto.Plan, error
 			if hintLoader, err = optimize.Hints(table, o.Hints, o.Rule, o.ShadowRule); err != nil {
 				return nil, errors.Wrap(err, "calculate hints failed")
 			}
-			shards = hintLoader.GetShards()
 			matchShadow = hintLoader.GetMatchBy(table.Suffix(), constants.ShadowUpdate)
 		}
 
 		if shards == nil {
+			if o.ShadowRule != nil && !matchShadow {
+				if matchShadow, err = (*optimize.ShadowSharder)(o.ShadowRule).Shard(table, constants.ShadowUpdate, stmt.Where, o.Args...); err != nil {
+					return nil, errors.Wrap(err, "calculate shadow regex failed")
+				}
+			}
+
 			if shards, fullScan, err = (*optimize.Sharder)(o.Rule).Shard(table, where, o.Args...); err != nil {
-				return nil, errors.Wrap(err, "failed to update")
+				return nil, errors.Wrap(err, "calculate shards failed")
 			}
 		}
 	}
@@ -103,8 +108,7 @@ func optimizeUpdate(_ context.Context, o *optimize.Optimizer) (proto.Plan, error
 	}
 
 	if matchShadow {
-		//TODO: fix it
-		//shards.ReplaceDb(o.ShadowRule.GetDatabase(stmt.Table.Suffix()))
+		shards.ReplaceDb()
 	}
 
 	ret := dml.NewUpdatePlan(stmt)

diff --git a/scripts/sharding.sql b/scripts/sharding.sql
@@ -137,4 +137,3 @@ CREATE TABLE IF NOT EXISTS `employees_0000_r`.`__shadow_student_0006` LIKE `empl
 CREATE TABLE IF NOT EXISTS `employees_0000_r`.`__shadow_student_0007` LIKE `employees_0000`.`student_0000`;
 
 INSERT INTO employees_0000.student_0001(id,uid,name,score,nickname,gender,birth_year,created_at,modified_at) VALUES (1, 1, 'arana', 95, 'Awesome Arana', 0, 2021, NOW(), NOW());
-INSERT INTO employees_0000.__shadow_student_0002(id,uid,name,score,nickname,gender,birth_year,created_at,modified_at) VALUES (2, 2, 'hanmeimei', 97, 'Shadow Arana', 0, 2021, NOW(), NOW());
diff --git a/test/dataset.go b/test/dataset.go
@@ -82,6 +82,8 @@ type (
 		ExecCases      []*Case `yaml:"exec_cases"`
 		QueryRowsCases []*Case `yaml:"query_rows_cases"`
 		QueryRowCases  []*Case `yaml:"query_row_cases"`
+		ExShHintCases  []*Case `yaml:"exec_shadow_hint_cases"`
+		ExShRegexCases []*Case `yaml:"exec_shadow_regex_cases"`
 	}
 
 	Case struct {