Created
November 26, 2025 17:19
-
-
Save mattlord/c9f3f80e1c718726a26789cc82656d84 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| diff --git a/go/vt/sqlparser/ast_funcs.go b/go/vt/sqlparser/ast_funcs.go | |
| index 290c57f9d9..e9af5c03fe 100644 | |
| --- a/go/vt/sqlparser/ast_funcs.go | |
| +++ b/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 | |
| diff --git a/go/vt/sqlparser/ast_funcs_test.go b/go/vt/sqlparser/ast_funcs_test.go | |
| index 7bec47df96..da8ddb7c30 100644 | |
| --- a/go/vt/sqlparser/ast_funcs_test.go | |
| +++ b/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)) | |
| +} | |
| diff --git a/go/vt/vtgate/planbuilder/bypass.go b/go/vt/vtgate/planbuilder/bypass.go | |
| index 52286816a1..d54cf150d9 100644 | |
| --- a/go/vt/vtgate/planbuilder/bypass.go | |
| +++ b/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(), | |
| diff --git a/go/vt/vtgate/planbuilder/testdata/bypass_keyrange_cases.json b/go/vt/vtgate/planbuilder/testdata/bypass_keyrange_cases.json | |
| index b13bafd77f..e1ff7de97a 100644 | |
| --- a/go/vt/vtgate/planbuilder/testdata/bypass_keyrange_cases.json | |
| +++ b/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')" | |
| + } | |
| + } | |
| } | |
| ] | |
| diff --git a/go/vt/vtgate/planbuilder/testdata/bypass_shard_cases.json b/go/vt/vtgate/planbuilder/testdata/bypass_shard_cases.json | |
| index 6f2be325b6..bd0fd9cd7b 100644 | |
| --- a/go/vt/vtgate/planbuilder/testdata/bypass_shard_cases.json | |
| +++ b/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')" | |
| + } | |
| + } | |
| } | |
| ] | |
| \ No newline at end of file |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment