From b4b88ee84ad7a88ca79033322ce8ae453bbe929a Mon Sep 17 00:00:00 2001 From: Harshit Gangal Date: Fri, 10 Jan 2025 20:07:34 +0530 Subject: [PATCH] remove shard targeted keyspace from database qualifier on the query Signed-off-by: Harshit Gangal --- go/vt/sqlparser/ast_funcs.go | 28 +++++++++++-------- go/vt/sqlparser/ast_test.go | 15 ++++++++++ go/vt/vtgate/planbuilder/builder.go | 4 --- go/vt/vtgate/planbuilder/bypass.go | 2 ++ .../testdata/bypass_keyrange_cases.json | 17 +++++++++++ .../testdata/bypass_shard_cases.json | 17 +++++++++++ 6 files changed, 68 insertions(+), 15 deletions(-) diff --git a/go/vt/sqlparser/ast_funcs.go b/go/vt/sqlparser/ast_funcs.go index 2891d532d16..f0c262a53b5 100644 --- a/go/vt/sqlparser/ast_funcs.go +++ b/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) } diff --git a/go/vt/sqlparser/ast_test.go b/go/vt/sqlparser/ast_test.go index c1484df7cc4..6a2689e59ef 100644 --- a/go/vt/sqlparser/ast_test.go +++ b/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)) +} diff --git a/go/vt/vtgate/planbuilder/builder.go b/go/vt/vtgate/planbuilder/builder.go index ca4ccb7ac5a..065c50a6dfa 100644 --- a/go/vt/vtgate/planbuilder/builder.go +++ b/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 diff --git a/go/vt/vtgate/planbuilder/bypass.go b/go/vt/vtgate/planbuilder/bypass.go index d3384d509c1..6e3f64990d7 100644 --- a/go/vt/vtgate/planbuilder/bypass.go +++ b/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(), diff --git a/go/vt/vtgate/planbuilder/testdata/bypass_keyrange_cases.json b/go/vt/vtgate/planbuilder/testdata/bypass_keyrange_cases.json index b13bafd77f8..e1ff7de97a0 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 02a00444724..80d4dbbd08b 100644 --- a/go/vt/vtgate/planbuilder/testdata/bypass_shard_cases.json +++ b/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')" + } + } } ]