From d7ff21beab5e56a79813a1a31eea9b5438234a95 Mon Sep 17 00:00:00 2001 From: Harshit Gangal Date: Tue, 18 Feb 2025 11:05:08 +0530 Subject: [PATCH] wip Signed-off-by: Harshit Gangal --- go/vt/vtgate/engine/plan.go | 2 +- go/vt/vtgate/engine/plankey_test.go | 66 ++++++++++++++++++ go/vt/vtgate/executor_test.go | 61 +++++++++++++++++ go/vt/vtgate/executorcontext/vcursor_impl.go | 35 ---------- .../executorcontext/vcursor_impl_test.go | 67 ++----------------- 5 files changed, 132 insertions(+), 99 deletions(-) create mode 100644 go/vt/vtgate/engine/plankey_test.go diff --git a/go/vt/vtgate/engine/plan.go b/go/vt/vtgate/engine/plan.go index 18aee6fff6d..a0488e591ca 100644 --- a/go/vt/vtgate/engine/plan.go +++ b/go/vt/vtgate/engine/plan.go @@ -61,7 +61,7 @@ type PlanKey struct { } func (pk PlanKey) DebugString() string { - return fmt.Sprintf("CurrentKeyspace: %s, Query: %s, SetVarComment: %s, Collation: %d", pk.CurrentKeyspace, pk.Query, pk.SetVarComment, pk.Collation) + return fmt.Sprintf("CurrentKeyspace: %s, Destination: %s, Query: %s, SetVarComment: %s, Collation: %d", pk.CurrentKeyspace, pk.Destination, pk.Query, pk.SetVarComment, pk.Collation) } func (pk PlanKey) Hash() theine.HashKey256 { diff --git a/go/vt/vtgate/engine/plankey_test.go b/go/vt/vtgate/engine/plankey_test.go new file mode 100644 index 00000000000..f205cd2199c --- /dev/null +++ b/go/vt/vtgate/engine/plankey_test.go @@ -0,0 +1,66 @@ +package engine + +import ( + "fmt" + "strings" + + "github.com/stretchr/testify/require" + "vitess.io/vitess/go/mysql/collations" + topodatapb "vitess.io/vitess/go/vt/proto/topodata" + vtgatepb "vitess.io/vitess/go/vt/proto/vtgate" + "vitess.io/vitess/go/vt/sqlparser" + "vitess.io/vitess/go/vt/srvtopo" + "vitess.io/vitess/go/vt/vtgate/vindexes" +) + +func TestPlanKey(t *testing.T) { + type testCase struct { + vschema *vindexes.VSchema + targetString string + expectedPlanPrefixKey string + } + + tests := []testCase{{ + vschema: vschemaWith1KS, + targetString: "", + expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "ks1@replica", + expectedPlanPrefixKey: "ks1@replica+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "ks1:-80", + expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+DestinationShard(-80)+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "ks1[deadbeef]", + expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+KsIDsResolved:80-+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "", + expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "ks1@replica", + expectedPlanPrefixKey: "ks1@replica+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", + }} + + for i, tc := range tests { + t.Run(fmt.Sprintf("%d#%s", i, tc.targetString), func(t *testing.T) { + ss := NewSafeSession(&vtgatepb.Session{InTransaction: false}) + ss.SetTargetString(tc.targetString) + cfg := VCursorConfig{ + Collation: collations.CollationUtf8mb4ID, + DefaultTabletType: topodatapb.TabletType_PRIMARY, + } + vc, err := NewVCursorImpl(ss, sqlparser.MarginComments{}, &fakeExecutor{}, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, srvtopo.NewResolver(&FakeTopoServer{}, nil, ""), nil, fakeObserver{}, cfg) + require.NoError(t, err) + vc.vschema = tc.vschema + + var buf strings.Builder + vc.KeyForPlan(context.Background(), "SELECT 1", &buf) + require.Equal(t, tc.expectedPlanPrefixKey, buf.String()) + }) + } +} diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go index acb58bacc71..161afcfebcc 100644 --- a/go/vt/vtgate/executor_test.go +++ b/go/vt/vtgate/executor_test.go @@ -35,6 +35,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/protobuf/proto" + "vitess.io/vitess/go/vt/srvtopo" "vitess.io/vitess/go/streamlog" @@ -61,6 +62,66 @@ import ( "vitess.io/vitess/go/vt/vtgate/vtgateservice" ) +func TestPlanKey(t *testing.T) { + ks1 := &vindexes.Keyspace{Name: "ks1"} + ks1Schema := &vindexes.KeyspaceSchema{Keyspace: ks1} + vschemaWith1KS := &vindexes.VSchema{ + Keyspaces: map[string]*vindexes.KeyspaceSchema{ + ks1.Name: ks1Schema, + }, + } + + type testCase struct { + vschema *vindexes.VSchema + targetString string + expectedPlanPrefixKey string + } + + tests := []testCase{{ + vschema: vschemaWith1KS, + targetString: "", + expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "ks1@replica", + expectedPlanPrefixKey: "ks1@replica+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "ks1:-80", + expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+DestinationShard(-80)+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "ks1[deadbeef]", + expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+KsIDsResolved:80-+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "", + expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", + }, { + vschema: vschemaWith1KS, + targetString: "ks1@replica", + expectedPlanPrefixKey: "ks1@replica+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", + }} + + for i, tc := range tests { + t.Run(fmt.Sprintf("%d#%s", i, tc.targetString), func(t *testing.T) { + ss := NewSafeSession(&vtgatepb.Session{InTransaction: false}) + ss.SetTargetString(tc.targetString) + cfg := VCursorConfig{ + Collation: collations.CollationUtf8mb4ID, + DefaultTabletType: topodatapb.TabletType_PRIMARY, + } + vc, err := NewVCursorImpl(ss, sqlparser.MarginComments{}, &fakeExecutor{}, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, srvtopo.NewResolver(&FakeTopoServer{}, nil, ""), nil, fakeObserver{}, cfg) + require.NoError(t, err) + vc.vschema = tc.vschema + + var buf strings.Builder + vc.KeyForPlan(context.Background(), "SELECT 1", &buf) + require.Equal(t, tc.expectedPlanPrefixKey, buf.String()) + }) + } +} + func TestExecutorResultsExceeded(t *testing.T) { executor, _, _, sbclookup, ctx := createExecutorEnv(t) diff --git a/go/vt/vtgate/executorcontext/vcursor_impl.go b/go/vt/vtgate/executorcontext/vcursor_impl.go index e8af00d7358..ccb9d1a2e8d 100644 --- a/go/vt/vtgate/executorcontext/vcursor_impl.go +++ b/go/vt/vtgate/executorcontext/vcursor_impl.go @@ -19,7 +19,6 @@ package executorcontext import ( "context" "fmt" - "io" "sort" "strings" "sync" @@ -1368,40 +1367,6 @@ func (vc *VCursorImpl) FindMirrorRule(name sqlparser.TableName) (*vindexes.Mirro return mirrorRule, err } -func (vc *VCursorImpl) KeyForPlan(ctx context.Context, query string, buf io.StringWriter) { - _, _ = buf.WriteString(vc.keyspace) - _, _ = buf.WriteString(vindexes.TabletTypeSuffix[vc.tabletType]) - _, _ = buf.WriteString("+Collate:") - _, _ = buf.WriteString(vc.Environment().CollationEnv().LookupName(vc.config.Collation)) - - if vc.destination != nil { - switch vc.destination.(type) { - case key.DestinationKeyspaceID, key.DestinationKeyspaceIDs: - resolved, _, err := vc.ResolveDestinations(ctx, vc.keyspace, nil, []key.Destination{vc.destination}) - if err == nil && len(resolved) > 0 { - shards := make([]string, len(resolved)) - for i := 0; i < len(shards); i++ { - shards[i] = resolved[i].Target.GetShard() - } - sort.Strings(shards) - - _, _ = buf.WriteString("+KsIDsResolved:") - for i, s := range shards { - if i > 0 { - _, _ = buf.WriteString(",") - } - _, _ = buf.WriteString(s) - } - } - default: - _, _ = buf.WriteString("+") - _, _ = buf.WriteString(vc.destination.String()) - } - } - _, _ = buf.WriteString("+Query:") - _, _ = buf.WriteString(query) -} - func (vc *VCursorImpl) GetKeyspace() string { return vc.keyspace } diff --git a/go/vt/vtgate/executorcontext/vcursor_impl_test.go b/go/vt/vtgate/executorcontext/vcursor_impl_test.go index 54173bf63b0..748a77ffb9b 100644 --- a/go/vt/vtgate/executorcontext/vcursor_impl_test.go +++ b/go/vt/vtgate/executorcontext/vcursor_impl_test.go @@ -21,7 +21,6 @@ import ( "errors" "fmt" "strconv" - "strings" "testing" "time" @@ -29,7 +28,6 @@ import ( "vitess.io/vitess/go/streamlog" - "vitess.io/vitess/go/mysql/collations" "vitess.io/vitess/go/sqltypes" binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata" "vitess.io/vitess/go/vt/proto/vschema" @@ -193,15 +191,10 @@ func TestDestinationKeyspace(t *testing.T) { } var ( - ks1 = &vindexes.Keyspace{Name: "ks1"} - ks1Schema = &vindexes.KeyspaceSchema{Keyspace: ks1} - ks2 = &vindexes.Keyspace{Name: "ks2"} - ks2Schema = &vindexes.KeyspaceSchema{Keyspace: ks2} - vschemaWith1KS = &vindexes.VSchema{ - Keyspaces: map[string]*vindexes.KeyspaceSchema{ - ks1.Name: ks1Schema, - }, - } + ks1 = &vindexes.Keyspace{Name: "ks1"} + ks1Schema = &vindexes.KeyspaceSchema{Keyspace: ks1} + ks2 = &vindexes.Keyspace{Name: "ks2"} + ks2Schema = &vindexes.KeyspaceSchema{Keyspace: ks2} ) var vschemaWith2KS = &vindexes.VSchema{ @@ -253,58 +246,6 @@ func TestSetTarget(t *testing.T) { } } -func TestKeyForPlan(t *testing.T) { - type testCase struct { - vschema *vindexes.VSchema - targetString string - expectedPlanPrefixKey string - } - - tests := []testCase{{ - vschema: vschemaWith1KS, - targetString: "", - expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "ks1@replica", - expectedPlanPrefixKey: "ks1@replica+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "ks1:-80", - expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+DestinationShard(-80)+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "ks1[deadbeef]", - expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+KsIDsResolved:80-+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "", - expectedPlanPrefixKey: "ks1@primary+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", - }, { - vschema: vschemaWith1KS, - targetString: "ks1@replica", - expectedPlanPrefixKey: "ks1@replica+Collate:utf8mb4_0900_ai_ci+Query:SELECT 1", - }} - - for i, tc := range tests { - t.Run(fmt.Sprintf("%d#%s", i, tc.targetString), func(t *testing.T) { - ss := NewSafeSession(&vtgatepb.Session{InTransaction: false}) - ss.SetTargetString(tc.targetString) - cfg := VCursorConfig{ - Collation: collations.CollationUtf8mb4ID, - DefaultTabletType: topodatapb.TabletType_PRIMARY, - } - vc, err := NewVCursorImpl(ss, sqlparser.MarginComments{}, &fakeExecutor{}, nil, &fakeVSchemaOperator{vschema: tc.vschema}, tc.vschema, srvtopo.NewResolver(&FakeTopoServer{}, nil, ""), nil, fakeObserver{}, cfg) - require.NoError(t, err) - vc.vschema = tc.vschema - - var buf strings.Builder - vc.KeyForPlan(context.Background(), "SELECT 1", &buf) - require.Equal(t, tc.expectedPlanPrefixKey, buf.String()) - }) - } -} - func TestFirstSortedKeyspace(t *testing.T) { ks1Schema := &vindexes.KeyspaceSchema{Keyspace: &vindexes.Keyspace{Name: "xks1"}} ks2Schema := &vindexes.KeyspaceSchema{Keyspace: &vindexes.Keyspace{Name: "aks2"}}