vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/collations_test.go (about) 1 /* 2 Copyright 2021 The Vitess Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package planbuilder 18 19 import ( 20 "fmt" 21 "testing" 22 23 "github.com/stretchr/testify/require" 24 25 "vitess.io/vitess/go/mysql/collations" 26 "vitess.io/vitess/go/vt/vtgate/engine" 27 ) 28 29 // collationInTable allows us to set a collation on a column 30 type collationInTable struct { 31 ks, table, collationName string 32 colName string 33 } 34 35 type collationTestCase struct { 36 query string 37 check func(t *testing.T, colls []collationInTable, primitive engine.Primitive) 38 collations []collationInTable 39 } 40 41 func (tc *collationTestCase) run(t *testing.T) { 42 vschemaWrapper := &vschemaWrapper{ 43 v: loadSchema(t, "vschemas/schema.json", false), 44 sysVarEnabled: true, 45 version: Gen4, 46 } 47 48 tc.addCollationsToSchema(vschemaWrapper) 49 plan, err := TestBuilder(tc.query, vschemaWrapper, vschemaWrapper.currentDb()) 50 require.NoError(t, err) 51 tc.check(t, tc.collations, plan.Instructions) 52 } 53 54 func (tc *collationTestCase) addCollationsToSchema(vschema *vschemaWrapper) { 55 for _, collation := range tc.collations { 56 tbl := vschema.v.Keyspaces[collation.ks].Tables[collation.table] 57 for i, c := range tbl.Columns { 58 if c.Name.EqualString(collation.colName) { 59 tbl.Columns[i].CollationName = collation.collationName 60 break 61 } 62 } 63 } 64 } 65 66 func TestOrderedAggregateCollations(t *testing.T) { 67 collid := func(collname string) collations.ID { 68 return collations.Local().LookupByName(collname).ID() 69 } 70 testCases := []collationTestCase{ 71 { 72 collations: []collationInTable{{ks: "user", table: "user", collationName: "utf8mb4_bin", colName: "textcol1"}}, 73 query: "select textcol1 from user group by textcol1", 74 check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) { 75 oa, isOA := primitive.(*engine.OrderedAggregate) 76 require.True(t, isOA, "should be an OrderedAggregate") 77 require.Equal(t, collid(colls[0].collationName), oa.GroupByKeys[0].CollationID) 78 }, 79 }, 80 { 81 collations: []collationInTable{{ks: "user", table: "user", collationName: "utf8mb4_bin", colName: "textcol1"}}, 82 query: "select distinct textcol1 from user", 83 check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) { 84 oa, isOA := primitive.(*engine.OrderedAggregate) 85 require.True(t, isOA, "should be an OrderedAggregate") 86 require.Equal(t, collid(colls[0].collationName), oa.GroupByKeys[0].CollationID) 87 }, 88 }, 89 { 90 collations: []collationInTable{ 91 {ks: "user", table: "user", collationName: "utf8mb4_bin", colName: "textcol1"}, 92 {ks: "user", table: "user", collationName: "utf8mb4_bin", colName: "textcol2"}, 93 }, 94 query: "select textcol1, textcol2 from user group by textcol1, textcol2", 95 check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) { 96 oa, isOA := primitive.(*engine.OrderedAggregate) 97 require.True(t, isOA, "should be an OrderedAggregate") 98 require.Equal(t, collid(colls[0].collationName), oa.GroupByKeys[0].CollationID) 99 require.Equal(t, collid(colls[1].collationName), oa.GroupByKeys[1].CollationID) 100 }, 101 }, 102 { 103 collations: []collationInTable{ 104 {ks: "user", table: "user", collationName: "utf8mb4_bin", colName: "textcol2"}, 105 }, 106 query: "select count(*), textcol2 from user group by textcol2", 107 check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) { 108 oa, isOA := primitive.(*engine.OrderedAggregate) 109 require.True(t, isOA, "should be an OrderedAggregate") 110 require.Equal(t, collid(colls[0].collationName), oa.GroupByKeys[0].CollationID) 111 }, 112 }, 113 { 114 collations: []collationInTable{ 115 {ks: "user", table: "user", collationName: "utf8mb4_bin", colName: "textcol2"}, 116 }, 117 query: "select count(*) as c, textcol2 from user group by textcol2 order by c", 118 check: func(t *testing.T, colls []collationInTable, primitive engine.Primitive) { 119 memSort, isMemSort := primitive.(*engine.MemorySort) 120 require.True(t, isMemSort, "should be a MemorySort") 121 oa, isOA := memSort.Input.(*engine.OrderedAggregate) 122 require.True(t, isOA, "should be an OrderedAggregate") 123 require.Equal(t, collid(colls[0].collationName), oa.GroupByKeys[0].CollationID) 124 }, 125 }, 126 } 127 128 for i, tc := range testCases { 129 t.Run(fmt.Sprintf("%d %s", i+1, tc.query), func(t *testing.T) { 130 tc.run(t) 131 }) 132 } 133 }