Testing v18 backport of #17503

Signed-off-by: Matt Lord <[email protected]>

Author: mattlord

go/vt/sqlparser/ast_funcs.go

@@ -2135,6 +2135,14 @@ func RemoveKeyspace(in SQLNode) {
 	}, in)
 }
 
+// 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
+	})
+}
+
 // RemoveKeyspaceInTables removes the database qualifier for all table names in the AST
 func RemoveKeyspaceInTables(in SQLNode) {
 	Rewrite(in, nil, func(cursor *Cursor) bool {
@@ -2147,6 +2155,23 @@ func RemoveKeyspaceInTables(in SQLNode) {
 	})
 }
 
+func removeKeyspace(in SQLNode, shouldRemove func(qualifier string) bool) {
+	Rewrite(in, nil, func(cursor *Cursor) bool {
+		switch expr := cursor.Node().(type) {
+		case *ColName:
+			if shouldRemove(expr.Qualifier.Qualifier.String()) {
+				expr.Qualifier.Qualifier = NewIdentifierCS("")
+			}
+		case TableName:
+			if shouldRemove(expr.Qualifier.String()) {
+				expr.Qualifier = NewIdentifierCS("")
+				cursor.Replace(expr)
+			}
+		}
+		return true
+	})
+}
+
 func convertStringToInt(integer string) int {
 	val, _ := strconv.Atoi(integer)
 	return val

go/vt/sqlparser/ast_funcs_test.go

@@ -172,3 +172,18 @@ func TestColumns_Indexes(t *testing.T) {
 		})
 	}
 }
+
+// TestRemoveSpecificKeyspace tests the RemoveSpecificKeyspace function.
+// It removes the specific keyspace from the database qualifier.
+func TestRemoveSpecificKeyspace(t *testing.T) {
+	stmt, err := 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/bypass.go

@@ -49,6 +49,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

@@ -174,5 +174,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": "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')"
+      }
+    }
   }
 ]