support force index (XiaoMi#102)

panhongrainbow · Sep 4, 2020 · 2f18fcc · 2f18fcc
1 parent 8aa4fc5
commit 2f18fcc
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 5 deletions.
diff --git a/proxy/plan/decorator_table_name.go b/proxy/plan/decorator_table_name.go
@@ -20,6 +20,7 @@ import (
 	"github.com/XiaoMi/Gaea/parser/ast"
 	"github.com/XiaoMi/Gaea/parser/format"
 	"github.com/XiaoMi/Gaea/proxy/router"
+	"github.com/pingcap/errors"
 )
 
 // TableNameDecorator decorate TableName
@@ -62,10 +63,6 @@ func CreateTableNameDecorator(n *ast.TableName, rule router.Rule, result *RouteR
 		return nil, fmt.Errorf("TableName does not support PartitionNames in sharding")
 	}
 
-	if len(n.IndexHints) != 0 {
-		return nil, fmt.Errorf("TableName does not support IndexHints in sharding")
-	}
-
 	ret := &TableNameDecorator{
 		origin: n,
 		rule:   rule,
@@ -115,6 +112,12 @@ func (t *TableNameDecorator) Restore(ctx *format.RestoreCtx) error {
 		ctx.WriteName(fmt.Sprintf("%s_%04d", t.origin.Name.String(), tableIndex))
 	}
 
+	for _, value := range t.origin.IndexHints {
+		ctx.WritePlain(" ")
+		if err := value.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while splicing IndexHints")
+		}
+	}
 	return nil
 }
 

diff --git a/proxy/plan/plan_select_test.go b/proxy/plan/plan_select_test.go
@@ -1500,7 +1500,7 @@ func TestMycatSelectSubqueryInTableRefs(t *testing.T) {
 		},
 		{
 			db:  "db_mycat",
-			sql: "select id from (select user from tbl_mycat_unknown) as a", //unshard plan 
+			sql: "select id from (select user from tbl_mycat_unknown) as a", //unshard plan
 			sqls: map[string]map[string][]string{
 				"slice-0": {
 					"db_mycat_0": {"SELECT `id` FROM (SELECT `user` FROM (`tbl_mycat_unknown`)) AS `a`"},
@@ -3407,6 +3407,42 @@ func TestSelectMycatOrderByDatabase(t *testing.T) {
 	}
 }
 
+func TestSelectForceIndexDatabase(t *testing.T) {
+	ns, err := preparePlanInfo()
+	if err != nil {
+		t.Fatalf("prepare namespace error: %v", err)
+	}
+
+	tests := []SQLTestcase{
+		{
+			db:  "db_mycat",
+			sql: "select * from tbl_mycat force index(id, name) where id > 100 and name = `zhangsan`",
+			sqls: map[string]map[string][]string{
+				"slice-0": {
+					"db_mycat_0": {
+						"SELECT * FROM `tbl_mycat` FORCE INDEX (`id`, `name`) WHERE `id`>100 AND `name`=`zhangsan`",
+					},
+					"db_mycat_1": {
+						"SELECT * FROM `tbl_mycat` FORCE INDEX (`id`, `name`) WHERE `id`>100 AND `name`=`zhangsan`",
+					},
+				},
+				"slice-1": {
+					"db_mycat_2": {
+						"SELECT * FROM `tbl_mycat` FORCE INDEX (`id`, `name`) WHERE `id`>100 AND `name`=`zhangsan`",
+					},
+					"db_mycat_3": {
+						"SELECT * FROM `tbl_mycat` FORCE INDEX (`id`, `name`) WHERE `id`>100 AND `name`=`zhangsan`",
+					},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.sql, getTestFunc(ns, test))
+	}
+}
+
 func prepareShardKingshardRouter() (*router.Router, error) {
 	nsStr := `
 {