Skip to content

Commit

Permalink
Merge pull request #9038 from planetscale/collation-typing
Browse files Browse the repository at this point in the history
gen4: add collation name to the type definition
  • Loading branch information
systay authored Oct 22, 2021
2 parents 32b74e6 + 4149a50 commit e8cc6a1
Show file tree
Hide file tree
Showing 18 changed files with 259 additions and 69 deletions.
10 changes: 5 additions & 5 deletions go/mysql/collations/8bit.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ type simpletables struct {
}

type Collation_8bit_bin struct {
id uint
id ID
name string
simpletables
}
Expand All @@ -39,7 +39,7 @@ func (c *Collation_8bit_bin) init() {}
func (c *Collation_8bit_bin) Name() string {
return c.name
}
func (c *Collation_8bit_bin) Id() uint {
func (c *Collation_8bit_bin) Id() ID {
return c.id
}

Expand Down Expand Up @@ -69,7 +69,7 @@ func (c *Collation_8bit_bin) WeightStringLen(numBytes int) int {
}

type Collation_8bit_simple_ci struct {
id uint
id ID
name string
simpletables
}
Expand All @@ -80,7 +80,7 @@ func (c *Collation_8bit_simple_ci) Name() string {
return c.name
}

func (c *Collation_8bit_simple_ci) Id() uint {
func (c *Collation_8bit_simple_ci) Id() ID {
return c.id
}

Expand Down Expand Up @@ -142,7 +142,7 @@ type Collation_binary struct{}

func (c *Collation_binary) init() {}

func (c *Collation_binary) Id() uint {
func (c *Collation_binary) Id() ID {
return 63
}

Expand Down
20 changes: 17 additions & 3 deletions go/mysql/collations/collation.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@ import (
// Generate mysqldata.go from the JSON information dumped from MySQL
//go:generate go run ./tools/makemysqldata/

type ID uint16

const Unknown ID = 0

// Collation implements a MySQL-compatible collation. It defines how to compare
// for sorting order and equality two strings with the same encoding.
type Collation interface {
Expand All @@ -33,7 +37,7 @@ type Collation interface {
// Id returns the numerical identifier for this collation. This is the same
// value that is returned by MySQL in a query's headers to identify the collation
// for a given column
Id() uint
Id() ID

// Name is the full name of this collation, in the form of "ENCODING_LANG_SENSITIVITY"
Name() string
Expand Down Expand Up @@ -110,7 +114,7 @@ func minInt(i1, i2 int) int {
}

var collationsByName = make(map[string]Collation)
var collationsById = make(map[uint]Collation)
var collationsById = make(map[ID]Collation)

func register(c Collation) {
duplicatedCharset := func(old Collation) {
Expand Down Expand Up @@ -138,9 +142,19 @@ func LookupByName(name string) Collation {
return csi
}

// LookupIDByName returns the collation ID for the given name
func LookupIDByName(name string) ID {
csi := collationsByName[name]
if csi == nil {
return Unknown
}

return csi.Id()
}

// LookupById returns the collation with the given numerical identifier. The collation
// is initialized if it's the first time being accessed.
func LookupById(id uint) Collation {
func LookupById(id ID) Collation {
csi := collationsById[id]
if csi != nil {
csi.init()
Expand Down
8 changes: 4 additions & 4 deletions go/mysql/collations/remote.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import (
// used as a fallback or as a way to test our native implementations.
type RemoteCollation struct {
name string
id uint
id ID

prefix string
suffix string
Expand All @@ -33,7 +33,7 @@ type RemoteCollation struct {
err error
}

func makeRemoteCollation(conn *mysql.Conn, collid uint, collname string) *RemoteCollation {
func makeRemoteCollation(conn *mysql.Conn, collid ID, collname string) *RemoteCollation {
coll := &RemoteCollation{
name: collname,
id: collid,
Expand All @@ -52,7 +52,7 @@ func makeRemoteCollation(conn *mysql.Conn, collid uint, collname string) *Remote
}

func RemoteByName(conn *mysql.Conn, collname string) *RemoteCollation {
var collid uint
var collid ID
if known, ok := collationsByName[collname]; ok {
collid = known.Id()
}
Expand All @@ -67,7 +67,7 @@ func (c *RemoteCollation) LastError() error {

func (c *RemoteCollation) init() {}

func (c *RemoteCollation) Id() uint {
func (c *RemoteCollation) Id() ID {
return c.id
}

Expand Down
10 changes: 5 additions & 5 deletions go/mysql/collations/uca.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ type CollationUCA interface {

type Collation_utf8mb4_uca_0900 struct {
name string
id uint
id ID

weights uca.WeightTable
tailoring []uca.WeightPatch
Expand Down Expand Up @@ -72,7 +72,7 @@ func (c *Collation_utf8mb4_uca_0900) Name() string {
return c.name
}

func (c *Collation_utf8mb4_uca_0900) Id() uint {
func (c *Collation_utf8mb4_uca_0900) Id() ID {
return c.id
}

Expand Down Expand Up @@ -168,7 +168,7 @@ func (c *Collation_utf8mb4_0900_bin) Encoding() encoding.Encoding {
return encoding.Encoding_utf8mb4{}
}

func (c *Collation_utf8mb4_0900_bin) Id() uint {
func (c *Collation_utf8mb4_0900_bin) Id() ID {
return 309
}

Expand Down Expand Up @@ -197,7 +197,7 @@ func (c *Collation_utf8mb4_0900_bin) WeightStringLen(numBytes int) int {

type Collation_uca_legacy struct {
name string
id uint
id ID

charset encoding.Encoding
weights uca.WeightTable
Expand Down Expand Up @@ -226,7 +226,7 @@ func (c *Collation_uca_legacy) UnicodeWeightsTable() (uca.WeightTable, uca.Table
return c.uca.Weights()
}

func (c *Collation_uca_legacy) Id() uint {
func (c *Collation_uca_legacy) Id() ID {
return c.id
}

Expand Down
4 changes: 2 additions & 2 deletions go/mysql/collations/utf8.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ type Collation_utf8mb4_general_ci struct {

func (c *Collation_utf8mb4_general_ci) init() {}

func (c *Collation_utf8mb4_general_ci) Id() uint {
func (c *Collation_utf8mb4_general_ci) Id() ID {
return 45
}

Expand Down Expand Up @@ -123,7 +123,7 @@ type Collation_utf8mb4_bin struct{}

func (c *Collation_utf8mb4_bin) init() {}

func (c *Collation_utf8mb4_bin) Id() uint {
func (c *Collation_utf8mb4_bin) Id() ID {
return 46
}

Expand Down
7 changes: 5 additions & 2 deletions go/mysql/schema.go
Original file line number Diff line number Diff line change
Expand Up @@ -97,15 +97,18 @@ select table_schema, table_name, column_name, ordinal_position, character_set_na
from information_schema.columns
where table_schema = database()`

// fetchColumns are the columns we fetch
fetchColumns = "table_name, column_name, data_type, collation_name"

// FetchUpdatedTables queries fetches all information about updated tables
FetchUpdatedTables = `select table_name, column_name, data_type
FetchUpdatedTables = `select ` + fetchColumns + `
from _vt.schemacopy
where table_schema = database() and
table_name in ::tableNames
order by table_name, ordinal_position`

// FetchTables queries fetches all information about tables
FetchTables = `select table_name, column_name, data_type
FetchTables = `select ` + fetchColumns + `
from _vt.schemacopy
where table_schema = database()
order by table_name, ordinal_position`
Expand Down
3 changes: 3 additions & 0 deletions go/vt/vtgate/engine/ordered_aggregate.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@ import (
"fmt"
"strconv"

"vitess.io/vitess/go/mysql/collations"

"vitess.io/vitess/go/vt/sqlparser"

"google.golang.org/protobuf/proto"
Expand Down Expand Up @@ -66,6 +68,7 @@ type GroupByParams struct {
WeightStringCol int
Expr sqlparser.Expr
FromGroupBy bool
CollationID collations.ID
}

// String returns a string. Used for plan descriptions
Expand Down
124 changes: 124 additions & 0 deletions go/vt/vtgate/planbuilder/collations_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
/*
Copyright 2021 The Vitess Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package planbuilder

import (
"fmt"
"testing"

"github.com/stretchr/testify/require"

"vitess.io/vitess/go/mysql/collations"
"vitess.io/vitess/go/vt/vtgate/engine"
)

// collationInTable allows us to set a collation on a column
type collationInTable struct {
ks, table, collationName string
offsetInTable int
}

type collationTestCase struct {
query string
check func(t *testing.T, colls []collationInTable, primitive engine.Primitive)
collations []collationInTable
}

func (tc *collationTestCase) run(t *testing.T) {
vschemaWrapper := &vschemaWrapper{
v: loadSchema(t, "schema_test.json"),
sysVarEnabled: true,
version: Gen4,
}

tc.addCollationsToSchema(vschemaWrapper)
plan, err := TestBuilder(tc.query, vschemaWrapper, vschemaWrapper.currentDb())
require.NoError(t, err)
tc.check(t, tc.collations, plan.Instructions)
}

func (tc *collationTestCase) addCollationsToSchema(vschema *vschemaWrapper) {
for _, collation := range tc.collations {
vschema.v.Keyspaces[collation.ks].Tables[collation.table].Columns[collation.offsetInTable].CollationName = collation.collationName
}
}

func TestOrderedAggregateCollations(t *testing.T) {
testCases := []collationTestCase{
{
collations: []collationInTable{{ks: "user", table: "user", collationName: "utf8mb4_bin", offsetInTable: 2}},
query: "select textcol1 from user group by textcol1",
check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) {
oa, isOA := primitive.(*engine.OrderedAggregate)
require.True(t, isOA, "should be an OrderedAggregate")
require.Equal(t, collations.LookupByName(colls[0].collationName).Id(), oa.GroupByKeys[0].CollationID)
},
},
{
collations: []collationInTable{{ks: "user", table: "user", collationName: "utf8mb4_bin", offsetInTable: 2}},
query: "select distinct textcol1 from user",
check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) {
oa, isOA := primitive.(*engine.OrderedAggregate)
require.True(t, isOA, "should be an OrderedAggregate")
require.Equal(t, collations.LookupByName(colls[0].collationName).Id(), oa.GroupByKeys[0].CollationID)
},
},
{
collations: []collationInTable{
{ks: "user", table: "user", collationName: "utf8mb4_bin", offsetInTable: 2},
{ks: "user", table: "user", collationName: "utf8mb4_bin", offsetInTable: 4},
},
query: "select textcol1, textcol2 from user group by textcol1, textcol2",
check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) {
oa, isOA := primitive.(*engine.OrderedAggregate)
require.True(t, isOA, "should be an OrderedAggregate")
require.Equal(t, collations.LookupByName(colls[0].collationName).Id(), oa.GroupByKeys[0].CollationID)
require.Equal(t, collations.LookupByName(colls[1].collationName).Id(), oa.GroupByKeys[1].CollationID)
},
},
{
collations: []collationInTable{
{ks: "user", table: "user", collationName: "utf8mb4_bin", offsetInTable: 4},
},
query: "select count(*), textcol2 from user group by textcol2",
check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) {
oa, isOA := primitive.(*engine.OrderedAggregate)
require.True(t, isOA, "should be an OrderedAggregate")
require.Equal(t, collations.LookupByName(colls[0].collationName).Id(), oa.GroupByKeys[0].CollationID)
},
},
{
collations: []collationInTable{
{ks: "user", table: "user", collationName: "utf8mb4_bin", offsetInTable: 4},
},
query: "select count(*) as c, textcol2 from user group by textcol2 order by c",
check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) {
memSort, isMemSort := primitive.(*engine.MemorySort)
require.True(t, isMemSort, "should be a MemorySort")
oa, isOA := memSort.Input.(*engine.OrderedAggregate)
require.True(t, isOA, "should be an OrderedAggregate")
require.Equal(t, collations.LookupByName(colls[0].collationName).Id(), oa.GroupByKeys[0].CollationID)
},
},
}

for i, tc := range testCases {
t.Run(fmt.Sprintf("%d %s", i+1, tc.query), func(t *testing.T) {
tc.run(t)
})
}
}
Loading

0 comments on commit e8cc6a1

Please sign in to comment.