Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history 
Signed-off-by: Harshit Gangal <harshit@planetscale.com>
  • Loading branch information
@harshit-gangal
harshit-gangal committed Feb 18, 2025
1 parent 5085d6c commit d7ff21b
Show file tree
Hide file tree
Showing 5 changed files with 132 additions and 99 deletions.
2 changes: 1 addition & 1 deletion go/vt/vtgate/engine/plan.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
66 changes: 66 additions & 0 deletions go/vt/vtgate/engine/plankey_test.go
View file
Original file line number Diff line number Diff line change
@@ -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())
})
}
}
61 changes: 61 additions & 0 deletions go/vt/vtgate/executor_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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)

Expand Down
35 changes: 0 additions & 35 deletions go/vt/vtgate/executorcontext/vcursor_impl.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ package executorcontext
import (
"context"
"fmt"
"io"
"sort"
"strings"
"sync"
Expand Down Expand Up @@ -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
}
Expand Down
67 changes: 4 additions & 63 deletions go/vt/vtgate/executorcontext/vcursor_impl_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,15 +21,13 @@ import (
"errors"
"fmt"
"strconv"
"strings"
"testing"
"time"

"github.com/stretchr/testify/require"

"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"
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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"}}
Expand Down

0 comments on commit d7ff21b

Please sign in to comment.