Remove shard targeted keyspace from database qualifier on the query (#…

…17503)

Signed-off-by: Harshit Gangal <[email protected]>

go/vt/sqlparser/ast_funcs.go

@@ -2429,48 +2429,38 @@ func RemoveKeyspaceInCol(in SQLNode) {
 	}, in)
 }
 
-// RemoveKeyspaceInTables removes the Qualifier on all TableNames in the AST
-func RemoveKeyspaceInTables(in SQLNode) {
-	// Walk will only return an error if we return an error from the inner func. safe to ignore here
-	Rewrite(in, nil, func(cursor *Cursor) bool {
-		if tbl, ok := cursor.Node().(TableName); ok && tbl.Qualifier.NotEmpty() {
-			tbl.Qualifier = NewIdentifierCS("")
-			cursor.Replace(tbl)
-		}
-
-		return true
+// RemoveKeyspace removes the keyspace qualifier from all ColName and TableName
+func RemoveKeyspace(in SQLNode) {
+	removeKeyspace(in, func(_ string) bool {
+		return true // Always remove
 	})
 }
 
-// RemoveKeyspace removes the Qualifier.Qualifier on all ColNames and Qualifier on all TableNames in the AST
-func RemoveKeyspace(in SQLNode) {
-	Rewrite(in, nil, func(cursor *Cursor) bool {
-		switch expr := cursor.Node().(type) {
-		case *ColName:
-			if expr.Qualifier.Qualifier.NotEmpty() {
-				expr.Qualifier.Qualifier = NewIdentifierCS("")
-			}
-		case TableName:
-			if expr.Qualifier.NotEmpty() {
-				expr.Qualifier = NewIdentifierCS("")
-				cursor.Replace(expr)
-			}
-		}
-		return true
+// RemoveSpecificKeyspace removes the keyspace qualifier from all ColName and TableName
+// when it matches the keyspace provided
+func RemoveSpecificKeyspace(in SQLNode, keyspace string) {
+	removeKeyspace(in, func(qualifier string) bool {
+		return qualifier == keyspace // Remove only if it matches the provided keyspace
 	})
 }
 
-// RemoveKeyspaceIgnoreSysSchema removes the Qualifier.Qualifier on all ColNames and Qualifier on all TableNames in the AST
-// except for the system schema.
+// RemoveKeyspaceIgnoreSysSchema removes the keyspace qualifier from all ColName and TableName
+// except for the system schema qualifier.
 func RemoveKeyspaceIgnoreSysSchema(in SQLNode) {
+	removeKeyspace(in, func(qualifier string) bool {
+		return qualifier != "" && !SystemSchema(qualifier) // Remove if it's not empty and not a system schema
+	})
+}
+
+func removeKeyspace(in SQLNode, shouldRemove func(qualifier string) bool) {
 	Rewrite(in, nil, func(cursor *Cursor) bool {
 		switch expr := cursor.Node().(type) {
 		case *ColName:
-			if expr.Qualifier.Qualifier.NotEmpty() && !SystemSchema(expr.Qualifier.Qualifier.String()) {
+			if shouldRemove(expr.Qualifier.Qualifier.String()) {
 				expr.Qualifier.Qualifier = NewIdentifierCS("")
 			}
 		case TableName:
-			if expr.Qualifier.NotEmpty() && !SystemSchema(expr.Qualifier.String()) {
+			if shouldRemove(expr.Qualifier.String()) {
 				expr.Qualifier = NewIdentifierCS("")
 				cursor.Replace(expr)
 			}

go/vt/sqlparser/ast_funcs_test.go

@@ -219,3 +219,28 @@ func TestExtractTables(t *testing.T) {
 		})
 	}
 }
+
+// TestRemoveKeyspace tests the RemoveKeyspaceIgnoreSysSchema function.
+// It removes all the keyspace except system schema.
+func TestRemoveKeyspaceIgnoreSysSchema(t *testing.T) {
+	stmt, err := NewTestParser().Parse("select 1 from uks.unsharded join information_schema.tables")
+	require.NoError(t, err)
+	RemoveKeyspaceIgnoreSysSchema(stmt)
+
+	require.Equal(t, "select 1 from unsharded join information_schema.`tables`", String(stmt))
+}
+
+// TestRemoveSpecificKeyspace tests the RemoveSpecificKeyspace function.
+// It removes the specific keyspace from the database qualifier.
+func TestRemoveSpecificKeyspace(t *testing.T) {
+	stmt, err := NewTestParser().Parse("select 1 from uks.unsharded")
+	require.NoError(t, err)
+
+	// does not match
+	RemoveSpecificKeyspace(stmt, "ks2")
+	require.Equal(t, "select 1 from uks.unsharded", String(stmt))
+
+	// match
+	RemoveSpecificKeyspace(stmt, "uks")
+	require.Equal(t, "select 1 from unsharded", String(stmt))
+}

go/vt/sqlparser/ast_test.go

@@ -917,11 +917,3 @@ func TestCloneComments(t *testing.T) {
 		assert.Equal(t, "b", val)
 	}
 }
-
-func TestRemoveKeyspace(t *testing.T) {
-	stmt, err := NewTestParser().Parse("select 1 from uks.unsharded")
-	require.NoError(t, err)
-	RemoveKeyspaceIgnoreSysSchema(stmt)
-
-	require.Equal(t, "select 1 from unsharded", String(stmt))
-}

go/vt/vtgate/planbuilder/builder.go

@@ -43,10 +43,6 @@ const (
 	Gen4Left2Right = querypb.ExecuteOptions_Gen4Left2Right
 )
 
-var (
-	plannerVersions = []plancontext.PlannerVersion{Gen4, Gen4GreedyOnly, Gen4Left2Right}
-)
-
 type (
 	planResult struct {
 		primitive engine.Primitive

go/vt/vtgate/planbuilder/bypass.go

@@ -56,6 +56,8 @@ func buildPlanForBypass(stmt sqlparser.Statement, _ *sqlparser.ReservedVars, vsc
 		}
 	}
 
+	sqlparser.RemoveSpecificKeyspace(stmt, keyspace.Name)
+
 	send := &engine.Send{
 		Keyspace:             keyspace,
 		TargetDestination:    vschema.Destination(),

go/vt/vtgate/planbuilder/testdata/bypass_keyrange_cases.json

@@ -164,5 +164,22 @@
         "Query": "create /* test */ table t1(id bigint, primary key(id)) /* comments */"
       }
     }
+  },
+  {
+    "comment": "remove the matching keyspace from shard targeted query",
+    "query": "select count(*), col from `main`.unsharded join vt_main.t1 where exists (select 1 from main.t2 join information_schema.tables where table_name = 't3')",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select count(*), col from `main`.unsharded join vt_main.t1 where exists (select 1 from main.t2 join information_schema.tables where table_name = 't3')",
+      "Instructions": {
+        "OperatorType": "Send",
+        "Keyspace": {
+          "Name": "main",
+          "Sharded": false
+        },
+        "TargetDestination": "ExactKeyRange(-)",
+        "Query": "select count(*), col from unsharded join vt_main.t1 where exists (select 1 from t2 join information_schema.`tables` where table_name = 't3')"
+      }
+    }
   }
 ]

go/vt/vtgate/planbuilder/testdata/bypass_shard_cases.json

@@ -251,5 +251,22 @@
         "QueryTimeout": 100
       }
     }
+  },
+  {
+    "comment": "remove the matching keyspace from shard targeted query",
+    "query": "select count(*), col from `main`.unsharded join vt_main.t1 where exists (select 1 from main.t2 join information_schema.tables where table_name = 't3')",
+    "plan": {
+      "QueryType": "SELECT",
+      "Original": "select count(*), col from `main`.unsharded join vt_main.t1 where exists (select 1 from main.t2 join information_schema.tables where table_name = 't3')",
+      "Instructions": {
+        "OperatorType": "Send",
+        "Keyspace": {
+          "Name": "main",
+          "Sharded": false
+        },
+        "TargetDestination": "Shard(-80)",
+        "Query": "select count(*), col from unsharded join vt_main.t1 where exists (select 1 from t2 join information_schema.`tables` where table_name = 't3')"
+      }
+    }
   }
 ]