diff --git a/go/test/vschemawrapper/vschema_wrapper.go b/go/test/vschemawrapper/vschema_wrapper.go index 3f9f072afc6..27219c6a255 100644 --- a/go/test/vschemawrapper/vschema_wrapper.go +++ b/go/test/vschemawrapper/vschema_wrapper.go @@ -18,7 +18,9 @@ package vschemawrapper import ( "context" + "errors" "fmt" + "sort" "strings" "vitess.io/vitess/go/mysql/collations" @@ -275,17 +277,6 @@ func (vw *VSchemaWrapper) FindTableOrVindex(tab sqlparser.TableName) (*vindexes. return vw.Vcursor.FindTableOrVindex(tab) } -func (vw *VSchemaWrapper) getfirstKeyspace() (ks *vindexes.Keyspace) { - var f string - for name, schema := range vw.V.Keyspaces { - if f == "" || f > name { - f = name - ks = schema.Keyspace - } - } - return -} - func (vw *VSchemaWrapper) getActualKeyspace() string { if vw.Keyspace == nil { return "" @@ -301,15 +292,32 @@ func (vw *VSchemaWrapper) getActualKeyspace() string { } func (vw *VSchemaWrapper) SelectedKeyspace() (*vindexes.Keyspace, error) { - return vw.V.Keyspaces["main"].Keyspace, nil + return vw.AnyKeyspace() } func (vw *VSchemaWrapper) AnyKeyspace() (*vindexes.Keyspace, error) { - return vw.SelectedKeyspace() + ks, found := vw.V.Keyspaces["main"] + if found { + return ks.Keyspace, nil + } + + size := len(vw.V.Keyspaces) + if size == 0 { + return nil, errors.New("no keyspace found in vschema") + } + + // Find the first keyspace in the map alphabetically to get deterministic results + keys := make([]string, size) + for key := range vw.V.Keyspaces { + keys = append(keys, key) + } + sort.Strings(keys) + + return vw.V.Keyspaces[keys[0]].Keyspace, nil } func (vw *VSchemaWrapper) FirstSortedKeyspace() (*vindexes.Keyspace, error) { - return vw.V.Keyspaces["main"].Keyspace, nil + return vw.AnyKeyspace() } func (vw *VSchemaWrapper) TargetString() string { @@ -344,12 +352,12 @@ func (vw *VSchemaWrapper) IsViewsEnabled() bool { // FindMirrorRule finds the mirror rule for the requested keyspace, table // name, and the tablet type in the VSchema. -func (vs *VSchemaWrapper) FindMirrorRule(tab sqlparser.TableName) (*vindexes.MirrorRule, error) { +func (vw *VSchemaWrapper) FindMirrorRule(tab sqlparser.TableName) (*vindexes.MirrorRule, error) { destKeyspace, destTabletType, _, err := topoproto.ParseDestination(tab.Qualifier.String(), topodatapb.TabletType_PRIMARY) if err != nil { return nil, err } - mirrorRule, err := vs.V.FindMirrorRule(destKeyspace, tab.Name.String(), destTabletType) + mirrorRule, err := vw.V.FindMirrorRule(destKeyspace, tab.Name.String(), destTabletType) if err != nil { return nil, err }