From 020f93660ace1f9b1ae6cca23406f55c7427dcf9 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Wed, 29 Mar 2023 16:54:54 +0800
Subject: [PATCH] lightning: support auto_random column in composite primary
 key (#41463) (#41482)

close pingcap/tidb#41454
---
 br/pkg/lightning/backend/kv/sql2kv.go         | 53 +++++++++----------
 br/pkg/lightning/common/util.go               | 20 +++++++
 br/pkg/lightning/common/util_test.go          | 27 ++++++++++
 br/pkg/lightning/restore/check_info.go        |  3 +-
 br/pkg/lightning/restore/restore.go           |  4 +-
 br/pkg/lightning/restore/table_restore.go     |  5 +-
 .../data2/auto_random.t.000000001.csv         |  3 ++
 br/tests/lightning_auto_random_default/run.sh | 27 ++++++++++
 8 files changed, 108 insertions(+), 34 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 4dc80e0c17ce2..1a99b04c3e99e 100644
--- a/br/pkg/lightning/backend/kv/sql2kv.go
+++ b/br/pkg/lightning/backend/kv/sql2kv.go
@@ -57,10 +57,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
 }
@@ -74,17 +75,15 @@ func NewTableKVEncoder(tbl table.Table, options *SessionOptions) (Encoder, error
 	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()) {
-				incrementalBits := autoRandomIncrementBits(col, int(meta.AutoRandomBits))
-				autoRandomBits := rand.New(rand.NewSource(options.AutoRandomSeed)).Int63n(1<<meta.AutoRandomBits) << incrementalBits // nolint:gosec
-				autoIDFn = func(id int64) int64 {
-					return autoRandomBits | id
-				}
-				break
-			}
+	if meta.ContainsAutoRandomBits() {
+		col := common.GetAutoRandomColumn(meta)
+		autoRandomColID = col.ID
+		incrementalBits := autoRandomIncrementBits(col, int(meta.AutoRandomBits))
+		autoRandomBits := rand.New(rand.NewSource(options.AutoRandomSeed)).Int63n(1<<meta.AutoRandomBits) << incrementalBits // nolint:gosec
+		autoIDFn = func(id int64) int64 {
+			return autoRandomBits | id
 		}
 	} else if meta.ShardRowIDBits > 0 {
 		rd := rand.New(rand.NewSource(options.AutoRandomSeed)) // nolint:gosec
@@ -104,14 +103,15 @@ func NewTableKVEncoder(tbl table.Table, options *SessionOptions) (Encoder, error
 	}
 
 	return &tableKVEncoder{
-		tbl:      tbl,
-		se:       se,
-		genCols:  genCols,
-		autoIDFn: autoIDFn,
+		tbl:             tbl,
+		autoRandomColID: autoRandomColID,
+		se:              se,
+		genCols:         genCols,
+		autoIDFn:        autoIDFn,
 	}, nil
 }
 
-func autoRandomIncrementBits(col *table.Column, randomBits int) int {
+func autoRandomIncrementBits(col *model.ColumnInfo, randomBits int) int {
 	typeBitsLength := mysql.DefaultLengthOfMysqlTypes[col.GetType()] * 8
 	incrementalBits := typeBitsLength - randomBits
 	hasSignBit := !mysql.HasUnsignedFlag(col.GetFlag())
@@ -369,8 +369,8 @@ func (kvcodec *tableKVEncoder) Encode(
 
 		record = append(record, value)
 
-		if isTableAutoRandom(meta) && isPKCol(col.ToInfo()) {
-			incrementalBits := autoRandomIncrementBits(col, int(meta.AutoRandomBits))
+		if kvcodec.isAutoRandomCol(col.ToInfo()) {
+			incrementalBits := autoRandomIncrementBits(col.ToInfo(), int(meta.AutoRandomBits))
 			alloc := kvcodec.tbl.Allocators(kvcodec.se).Get(autoid.AutoRandomType)
 			if err := alloc.Rebase(context.Background(), value.GetInt64()&((1<<incrementalBits)-1), false); err != nil {
 				return nil, errors.Trace(err)
@@ -427,25 +427,20 @@ 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())
-}
-
 func (kvcodec *tableKVEncoder) getActualDatum(rowID int64, colIndex int, inputDatum *types.Datum) (types.Datum, error) {
 	var (
 		value types.Datum
 		err   error
 	)
 
-	tblMeta := kvcodec.tbl.Meta()
 	cols := kvcodec.tbl.Cols()
 
 	// Since this method is only called when iterating the columns in the `Encode()` method,
@@ -468,7 +463,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/util.go b/br/pkg/lightning/common/util.go
index 57afc1fb7eac0..87dec1994c703 100644
--- a/br/pkg/lightning/common/util.go
+++ b/br/pkg/lightning/common/util.go
@@ -36,6 +36,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"
 )
 
@@ -388,3 +389,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..3ef5eaf29cc0c 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) clustered)", "c"},
+		{"create table t(a int, c bigint auto_random, primary key(c,a) clustered)", "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/check_info.go b/br/pkg/lightning/restore/check_info.go
index be0fdd5e85008..1fe58dd716925 100644
--- a/br/pkg/lightning/restore/check_info.go
+++ b/br/pkg/lightning/restore/check_info.go
@@ -662,8 +662,9 @@ func (rc *Controller) SchemaIsValid(ctx context.Context, tableInfo *mydump.MDTab
 	colCountFromTiDB := len(info.Core.Columns)
 	core := info.Core
 	defaultCols := make(map[string]struct{})
+	autoRandomCol := common.GetAutoRandomColumn(core)
 	for _, col := range core.Columns {
-		if hasDefault(col) || (info.Core.ContainsAutoRandomBits() && mysql.HasPriKeyFlag(col.GetFlag())) {
+		if 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 92e02e3c4cd67..44c2c9c41a69e 100644
--- a/br/pkg/lightning/restore/restore.go
+++ b/br/pkg/lightning/restore/restore.go
@@ -1623,7 +1623,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
@@ -2303,7 +2303,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 d875cbd0af87d..03bb94faafcb9 100644
--- a/br/pkg/lightning/restore/table_restore.go
+++ b/br/pkg/lightning/restore/table_restore.go
@@ -688,11 +688,12 @@ func (tr *TableRestore) postProcess(
 		rc.alterTableLock.Lock()
 		tblInfo := tr.tableInfo.Core
 		var err error
-		if tblInfo.PKIsHandle && tblInfo.ContainsAutoRandomBits() {
+		if tblInfo.ContainsAutoRandomBits() {
+			ft := &common.GetAutoRandomColumn(tblInfo).FieldType
 			var maxAutoRandom, autoRandomTotalBits uint64
 			autoRandomTotalBits = 64
 			autoRandomBits := tblInfo.AutoRandomBits // range from (0, 15]
-			if !tblInfo.IsAutoRandomBitColUnsigned() {
+			if !mysql.HasUnsignedFlag(ft.GetFlag()) {
 				// if auto_random is signed, leave one extra bit
 				autoRandomTotalBits = 63
 			}
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..5ef3d56cb3ed7 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) clustered)'
+    # 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) clustered)'
+done