From 4958321b6101301299ccd31b62dad3c83bc9fc91 Mon Sep 17 00:00:00 2001 From: Ti Chi Robot Date: Thu, 16 Feb 2023 15:38:02 +0800 Subject: [PATCH] lightning: support auto_random column in composite primary key (#41463) (#41484) close pingcap/tidb#41454 --- br/pkg/lightning/backend/kv/sql2kv.go | 52 +++++++++---------- br/pkg/lightning/common/BUILD.bazel | 3 ++ br/pkg/lightning/common/util.go | 20 +++++++ br/pkg/lightning/common/util_test.go | 27 ++++++++++ br/pkg/lightning/restore/precheck_impl.go | 3 +- br/pkg/lightning/restore/restore.go | 4 +- br/pkg/lightning/restore/table_restore.go | 4 +- .../data2/auto_random.t.000000001.csv | 3 ++ br/tests/lightning_auto_random_default/run.sh | 27 ++++++++++ 9 files changed, 110 insertions(+), 33 deletions(-) create mode 100644 br/tests/lightning_auto_random_default/data2/auto_random.t.000000001.csv diff --git a/br/pkg/lightning/backend/kv/sql2kv.go b/br/pkg/lightning/backend/kv/sql2kv.go index 9ad552ef5f340..77bce0b05dd24 100644 --- a/br/pkg/lightning/backend/kv/sql2kv.go +++ b/br/pkg/lightning/backend/kv/sql2kv.go @@ -55,10 +55,11 @@ type genCol struct { type autoIDConverter func(int64) int64 type tableKVEncoder struct { - tbl table.Table - se *session - recordCache []types.Datum - genCols []genCol + tbl table.Table + autoRandomColID int64 + se *session + recordCache []types.Datum + genCols []genCol // convert auto id for shard rowid or auto random id base on row id generated by lightning autoIDFn autoIDConverter metrics *metric.Metrics @@ -84,17 +85,16 @@ func NewTableKVEncoder( recordCtx := tables.NewCommonAddRecordCtx(len(cols)) tables.SetAddRecordCtx(se, recordCtx) + var autoRandomColID int64 autoIDFn := func(id int64) int64 { return id } - if meta.PKIsHandle && meta.ContainsAutoRandomBits() { - for _, col := range cols { - if mysql.HasPriKeyFlag(col.GetFlag()) { - shardFmt := autoid.NewShardIDFormat(&col.FieldType, meta.AutoRandomBits, meta.AutoRandomRangeBits) - shard := rand.New(rand.NewSource(options.AutoRandomSeed)).Int63() - autoIDFn = func(id int64) int64 { - return shardFmt.Compose(shard, id) - } - break - } + if meta.ContainsAutoRandomBits() { + col := common.GetAutoRandomColumn(meta) + autoRandomColID = col.ID + + shardFmt := autoid.NewShardIDFormat(&col.FieldType, meta.AutoRandomBits, meta.AutoRandomRangeBits) + shard := rand.New(rand.NewSource(options.AutoRandomSeed)).Int63() + autoIDFn = func(id int64) int64 { + return shardFmt.Compose(shard, id) } } else if meta.ShardRowIDBits > 0 { rd := rand.New(rand.NewSource(options.AutoRandomSeed)) // nolint:gosec @@ -114,11 +114,12 @@ func NewTableKVEncoder( } return &tableKVEncoder{ - tbl: tbl, - se: se, - genCols: genCols, - autoIDFn: autoIDFn, - metrics: metrics, + tbl: tbl, + autoRandomColID: autoRandomColID, + se: se, + genCols: genCols, + autoIDFn: autoIDFn, + metrics: metrics, }, nil } @@ -380,7 +381,7 @@ func (kvcodec *tableKVEncoder) Encode( record = append(record, value) - if isTableAutoRandom(meta) && isPKCol(col.ToInfo()) { + if kvcodec.isAutoRandomCol(col.ToInfo()) { shardFmt := autoid.NewShardIDFormat(&col.FieldType, meta.AutoRandomBits, meta.AutoRandomRangeBits) alloc := kvcodec.tbl.Allocators(kvcodec.se).Get(autoid.AutoRandomType) if err := alloc.Rebase(context.Background(), value.GetInt64()&shardFmt.IncrementalMask(), false); err != nil { @@ -438,18 +439,14 @@ func (kvcodec *tableKVEncoder) Encode( return kvPairs, nil } -func isTableAutoRandom(tblMeta *model.TableInfo) bool { - return tblMeta.PKIsHandle && tblMeta.ContainsAutoRandomBits() +func (kvcodec *tableKVEncoder) isAutoRandomCol(col *model.ColumnInfo) bool { + return kvcodec.tbl.Meta().ContainsAutoRandomBits() && col.ID == kvcodec.autoRandomColID } func isAutoIncCol(colInfo *model.ColumnInfo) bool { return mysql.HasAutoIncrementFlag(colInfo.GetFlag()) } -func isPKCol(colInfo *model.ColumnInfo) bool { - return mysql.HasPriKeyFlag(colInfo.GetFlag()) -} - // GetEncoderIncrementalID return Auto increment id. func GetEncoderIncrementalID(encoder Encoder, id int64) int64 { return encoder.(*tableKVEncoder).autoIDFn(id) @@ -471,7 +468,6 @@ func (kvcodec *tableKVEncoder) getActualDatum(rowID int64, colIndex int, inputDa err error ) - tblMeta := kvcodec.tbl.Meta() cols := kvcodec.tbl.Cols() // Since this method is only called when iterating the columns in the `Encode()` method, @@ -494,7 +490,7 @@ func (kvcodec *tableKVEncoder) getActualDatum(rowID int64, colIndex int, inputDa case isAutoIncCol(col.ToInfo()): // we still need a conversion, e.g. to catch overflow with a TINYINT column. value, err = table.CastValue(kvcodec.se, types.NewIntDatum(rowID), col.ToInfo(), false, false) - case isTableAutoRandom(tblMeta) && isPKCol(col.ToInfo()): + case kvcodec.isAutoRandomCol(col.ToInfo()): var val types.Datum realRowID := kvcodec.autoIDFn(rowID) if mysql.HasUnsignedFlag(col.GetFlag()) { diff --git a/br/pkg/lightning/common/BUILD.bazel b/br/pkg/lightning/common/BUILD.bazel index 3bd871276d733..808832cd7e5a8 100644 --- a/br/pkg/lightning/common/BUILD.bazel +++ b/br/pkg/lightning/common/BUILD.bazel @@ -24,6 +24,7 @@ go_library( "//errno", "//parser/model", "//store/driver/error", + "//table/tables", "//util", "@com_github_go_sql_driver_mysql//:mysql", "@com_github_pingcap_errors//:errors", @@ -99,8 +100,10 @@ go_test( "//br/pkg/errors", "//br/pkg/lightning/log", "//errno", + "//parser", "//store/driver/error", "//testkit/testsetup", + "//util/dbutil", "@com_github_data_dog_go_sqlmock//:go-sqlmock", "@com_github_go_sql_driver_mysql//:mysql", "@com_github_pingcap_errors//:errors", diff --git a/br/pkg/lightning/common/util.go b/br/pkg/lightning/common/util.go index b9bdf564403de..c92b071502453 100644 --- a/br/pkg/lightning/common/util.go +++ b/br/pkg/lightning/common/util.go @@ -37,6 +37,7 @@ import ( "github.com/pingcap/tidb/br/pkg/utils" tmysql "github.com/pingcap/tidb/errno" "github.com/pingcap/tidb/parser/model" + "github.com/pingcap/tidb/table/tables" "go.uber.org/zap" ) @@ -416,3 +417,22 @@ func StringSliceEqual(a, b []string) bool { } return true } + +// GetAutoRandomColumn return the column with auto_random, return nil if the table doesn't have it. +// todo: better put in ddl package, but this will cause import cycle since ddl package import lightning +func GetAutoRandomColumn(tblInfo *model.TableInfo) *model.ColumnInfo { + if !tblInfo.ContainsAutoRandomBits() { + return nil + } + if tblInfo.PKIsHandle { + return tblInfo.GetPkColInfo() + } else if tblInfo.IsCommonHandle { + pk := tables.FindPrimaryIndex(tblInfo) + if pk == nil { + return nil + } + offset := pk.Columns[0].Offset + return tblInfo.Columns[offset] + } + return nil +} diff --git a/br/pkg/lightning/common/util_test.go b/br/pkg/lightning/common/util_test.go index a192ecea11906..781ee046c5fdc 100644 --- a/br/pkg/lightning/common/util_test.go +++ b/br/pkg/lightning/common/util_test.go @@ -31,6 +31,8 @@ import ( "github.com/pingcap/failpoint" "github.com/pingcap/tidb/br/pkg/lightning/common" "github.com/pingcap/tidb/br/pkg/lightning/log" + "github.com/pingcap/tidb/parser" + "github.com/pingcap/tidb/util/dbutil" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -179,3 +181,28 @@ func TestInterpolateMySQLString(t *testing.T) { assert.Equal(t, "'1''23'", common.InterpolateMySQLString("1'23")) assert.Equal(t, "'1''2''''3'", common.InterpolateMySQLString("1'2''3")) } + +func TestGetAutoRandomColumn(t *testing.T) { + tests := []struct { + ddl string + colName string + }{ + {"create table t(c int)", ""}, + {"create table t(c int auto_increment)", ""}, + {"create table t(c bigint auto_random primary key)", "c"}, + {"create table t(a int, c bigint auto_random primary key)", "c"}, + {"create table t(c bigint auto_random, a int, primary key(c,a))", "c"}, + {"create table t(a int, c bigint auto_random, primary key(c,a))", "c"}, + } + p := parser.New() + for _, tt := range tests { + tableInfo, err := dbutil.GetTableInfoBySQL(tt.ddl, p) + require.NoError(t, err) + col := common.GetAutoRandomColumn(tableInfo) + if tt.colName == "" { + require.Nil(t, col, tt.ddl) + } else { + require.Equal(t, tt.colName, col.Name.L, tt.ddl) + } + } +} diff --git a/br/pkg/lightning/restore/precheck_impl.go b/br/pkg/lightning/restore/precheck_impl.go index 3dcd5cce973f9..aff5dc7ae012b 100644 --- a/br/pkg/lightning/restore/precheck_impl.go +++ b/br/pkg/lightning/restore/precheck_impl.go @@ -1001,9 +1001,10 @@ func (ci *schemaCheckItem) SchemaIsValid(ctx context.Context, tableInfo *mydump. core := info.Core defaultCols := make(map[string]struct{}) + autoRandomCol := common.GetAutoRandomColumn(core) for _, col := range core.Columns { // we can extend column the same with columns with default values - if _, isExtendCol := fullExtendColsSet[col.Name.O]; isExtendCol || hasDefault(col) || (info.Core.ContainsAutoRandomBits() && mysql.HasPriKeyFlag(col.GetFlag())) { + if _, isExtendCol := fullExtendColsSet[col.Name.O]; isExtendCol || hasDefault(col) || (autoRandomCol != nil && autoRandomCol.ID == col.ID) { // this column has default value or it's auto random id, so we can ignore it defaultCols[col.Name.L] = struct{}{} } diff --git a/br/pkg/lightning/restore/restore.go b/br/pkg/lightning/restore/restore.go index 1032419ab777f..eb874c521f8d2 100644 --- a/br/pkg/lightning/restore/restore.go +++ b/br/pkg/lightning/restore/restore.go @@ -1804,7 +1804,7 @@ func (tr *TableRestore) restoreTable( web.BroadcastTableCheckpoint(tr.tableName, cp) // rebase the allocator so it exceeds the number of rows. - if tr.tableInfo.Core.PKIsHandle && tr.tableInfo.Core.ContainsAutoRandomBits() { + if tr.tableInfo.Core.ContainsAutoRandomBits() { cp.AllocBase = mathutil.Max(cp.AllocBase, tr.tableInfo.Core.AutoRandID) if err := tr.alloc.Get(autoid.AutoRandomType).Rebase(context.Background(), cp.AllocBase, false); err != nil { return false, err @@ -2518,7 +2518,7 @@ func saveCheckpoint(rc *Controller, t *TableRestore, engineID int32, chunk *chec // or integer primary key), which can only be obtained by reading all data. var base int64 - if t.tableInfo.Core.PKIsHandle && t.tableInfo.Core.ContainsAutoRandomBits() { + if t.tableInfo.Core.ContainsAutoRandomBits() { base = t.alloc.Get(autoid.AutoRandomType).Base() + 1 } else { base = t.alloc.Get(autoid.RowIDAllocType).Base() + 1 diff --git a/br/pkg/lightning/restore/table_restore.go b/br/pkg/lightning/restore/table_restore.go index 6b1372c0bca9e..9a9c365a9b8bc 100644 --- a/br/pkg/lightning/restore/table_restore.go +++ b/br/pkg/lightning/restore/table_restore.go @@ -737,8 +737,8 @@ func (tr *TableRestore) postProcess( rc.alterTableLock.Lock() tblInfo := tr.tableInfo.Core var err error - if tblInfo.PKIsHandle && tblInfo.ContainsAutoRandomBits() { - ft := &tblInfo.GetPkColInfo().FieldType + if tblInfo.ContainsAutoRandomBits() { + ft := &common.GetAutoRandomColumn(tblInfo).FieldType shardFmt := autoid.NewShardIDFormat(ft, tblInfo.AutoRandomBits, tblInfo.AutoRandomRangeBits) maxCap := shardFmt.IncrementalBitsCapacity() err = AlterAutoRandom(ctx, rc.tidbGlue.GetSQLExecutor(), tr.tableName, uint64(tr.alloc.Get(autoid.AutoRandomType).Base())+1, maxCap) diff --git a/br/tests/lightning_auto_random_default/data2/auto_random.t.000000001.csv b/br/tests/lightning_auto_random_default/data2/auto_random.t.000000001.csv new file mode 100644 index 0000000000000..5de7daa1f86f2 --- /dev/null +++ b/br/tests/lightning_auto_random_default/data2/auto_random.t.000000001.csv @@ -0,0 +1,3 @@ +"a","b" +1,11 +2,22 diff --git a/br/tests/lightning_auto_random_default/run.sh b/br/tests/lightning_auto_random_default/run.sh index 41b9798de4560..67f51a6518ec6 100644 --- a/br/tests/lightning_auto_random_default/run.sh +++ b/br/tests/lightning_auto_random_default/run.sh @@ -60,3 +60,30 @@ for backend in tidb local; do run_sql "SELECT max(id & b'000001111111111111111111111111111111111111111111111111111111111') >= $NEXT_AUTO_RAND_VAL as ge FROM auto_random.t" check_contains 'ge: 1' done + +function run_for_auro_random_data2() { + create_table=$1 + run_sql 'DROP DATABASE IF EXISTS auto_random;' + run_sql 'CREATE DATABASE IF NOT EXISTS auto_random;' + run_sql "$create_table" + run_lightning --backend $backend -d "tests/$TEST_NAME/data2" + run_sql 'select count(*) as count from auto_random.t where c > 0' + check_contains "count: 2" + run_sql 'select count(*) as count from auto_random.t where a=1 and b=11' + check_contains "count: 1" + run_sql 'select count(*) as count from auto_random.t where a=2 and b=22' + check_contains "count: 1" +} + +for backend in tidb local; do + if [ "$backend" = 'local' ]; then + check_cluster_version 4 0 0 'local backend' || continue + fi + + run_for_auro_random_data2 'create table auto_random.t(c bigint auto_random primary key, a int, b int)' + run_for_auro_random_data2 'create table auto_random.t(a int, b int, c bigint auto_random primary key)' + # composite key and auto_random is the first column + run_for_auro_random_data2 'create table auto_random.t(c bigint auto_random, a int, b int, primary key(c, a))' + # composite key and auto_random is not the first column + run_for_auro_random_data2 'create table auto_random.t(a int, b int, c bigint auto_random, primary key(c, a))' +done