remove shard targeted keyspace from database qualifier on the query

Signed-off-by: Harshit Gangal <harshit@planetscale.com>

go/vt/sqlparser/ast_funcs.go

@@ -2409,29 +2409,35 @@ 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
+// 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 {
-		if tbl, ok := cursor.Node().(TableName); ok && tbl.Qualifier.NotEmpty() {
-			tbl.Qualifier = NewIdentifierCS("")
-			cursor.Replace(tbl)
+		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
 	})
 }
 
-// RemoveKeyspace removes the Qualifier.Qualifier on all ColNames and Qualifier on all TableNames in the AST
-func RemoveKeyspace(in SQLNode) {
+// RemoveSpecificKeyspace removes the Qualifier.Qualifier on all ColNames and Qualifier on all TableNames in the AST
+// when it matches the keyspace provided
+func RemoveSpecificKeyspace(in SQLNode, keyspace string) {
 	Rewrite(in, nil, func(cursor *Cursor) bool {
 		switch expr := cursor.Node().(type) {
 		case *ColName:
-			if expr.Qualifier.Qualifier.NotEmpty() {
+			if expr.Qualifier.Qualifier.String() == keyspace {
 				expr.Qualifier.Qualifier = NewIdentifierCS("")
 			}
 		case TableName:
-			if expr.Qualifier.NotEmpty() {
+			if expr.Qualifier.String() == keyspace {
 				expr.Qualifier = NewIdentifierCS("")
 				cursor.Replace(expr)
 			}

go/vt/sqlparser/ast_test.go

@@ -925,3 +925,18 @@ func TestRemoveKeyspace(t *testing.T) {
 
 	require.Equal(t, "select 1 from unsharded", 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/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')"
+      }
+    }
   }
 ]