From b824e351b96d3e8644e9134ac2dbfc3d93d887fd Mon Sep 17 00:00:00 2001
From: dsdashun <dsdashun@qq.com>
Date: Thu, 12 May 2022 14:04:35 +0800
Subject: [PATCH] lightning: support null value for auto-incr column on local
 backend (#34552)

close pingcap/tidb#34208
---
 br/pkg/lightning/backend/kv/sql2kv.go         | 97 +++++++++++++------
 br/pkg/lightning/backend/kv/sql2kv_test.go    | 92 +++++++++++++++++-
 br/tests/lightning_auto_columns/config.toml   |  2 +
 .../lightning_auto_cols-schema-create.sql     |  1 +
 ...lightning_auto_cols.t_auto_incr-schema.sql |  3 +
 .../data/lightning_auto_cols.t_auto_incr.sql  |  8 ++
 ...ghtning_auto_cols.t_auto_random-schema.sql |  3 +
 .../lightning_auto_cols.t_auto_random.sql     |  8 ++
 br/tests/lightning_auto_columns/run.sh        | 38 ++++++++
 9 files changed, 219 insertions(+), 33 deletions(-)
 create mode 100644 br/tests/lightning_auto_columns/config.toml
 create mode 100644 br/tests/lightning_auto_columns/data/lightning_auto_cols-schema-create.sql
 create mode 100644 br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_incr-schema.sql
 create mode 100644 br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_incr.sql
 create mode 100644 br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_random-schema.sql
 create mode 100644 br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_random.sql
 create mode 100755 br/tests/lightning_auto_columns/run.sh

diff --git a/br/pkg/lightning/backend/kv/sql2kv.go b/br/pkg/lightning/backend/kv/sql2kv.go
index cfee7eceaea7a..4dc80e0c17ce2 100644
--- a/br/pkg/lightning/backend/kv/sql2kv.go
+++ b/br/pkg/lightning/backend/kv/sql2kv.go
@@ -356,50 +356,27 @@ func (kvcodec *tableKVEncoder) Encode(
 	}
 
 	meta := kvcodec.tbl.Meta()
-	isAutoRandom := meta.PKIsHandle && meta.ContainsAutoRandomBits()
 	for i, col := range cols {
+		var theDatum *types.Datum = nil
 		j := columnPermutation[i]
-		isAutoIncCol := mysql.HasAutoIncrementFlag(col.GetFlag())
-		isPk := mysql.HasPriKeyFlag(col.GetFlag())
-		switch {
-		case j >= 0 && j < len(row):
-			value, err = table.CastValue(kvcodec.se, row[j], col.ToInfo(), false, false)
-			if err == nil {
-				err = col.HandleBadNull(&value, kvcodec.se.vars.StmtCtx)
-			}
-		case isAutoIncCol:
-			// 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 isAutoRandom && isPk:
-			var val types.Datum
-			realRowID := kvcodec.autoIDFn(rowID)
-			if mysql.HasUnsignedFlag(col.GetFlag()) {
-				val = types.NewUintDatum(uint64(realRowID))
-			} else {
-				val = types.NewIntDatum(realRowID)
-			}
-			value, err = table.CastValue(kvcodec.se, val, col.ToInfo(), false, false)
-		case col.IsGenerated():
-			// inject some dummy value for gen col so that MutRowFromDatums below sees a real value instead of nil.
-			// if MutRowFromDatums sees a nil it won't initialize the underlying storage and cause SetDatum to panic.
-			value = types.GetMinValue(&col.FieldType)
-		default:
-			value, err = table.GetColDefaultValue(kvcodec.se, col.ToInfo())
+		if j >= 0 && j < len(row) {
+			theDatum = &row[j]
 		}
+		value, err = kvcodec.getActualDatum(rowID, i, theDatum)
 		if err != nil {
 			return nil, logKVConvertFailed(logger, row, j, col.ToInfo(), err)
 		}
 
 		record = append(record, value)
 
-		if isAutoRandom && isPk {
+		if isTableAutoRandom(meta) && isPKCol(col.ToInfo()) {
 			incrementalBits := autoRandomIncrementBits(col, 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)
 			}
 		}
-		if isAutoIncCol {
+		if isAutoIncCol(col.ToInfo()) {
 			alloc := kvcodec.tbl.Allocators(kvcodec.se).Get(autoid.AutoIncrementType)
 			if err := alloc.Rebase(context.Background(), getAutoRecordID(value, &col.FieldType), false); err != nil {
 				return nil, errors.Trace(err)
@@ -450,6 +427,68 @@ func (kvcodec *tableKVEncoder) Encode(
 	return kvPairs, nil
 }
 
+func isTableAutoRandom(tblMeta *model.TableInfo) bool {
+	return tblMeta.PKIsHandle && tblMeta.ContainsAutoRandomBits()
+}
+
+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,
+	// we can assume that the `colIndex` always have a valid input
+	col := cols[colIndex]
+
+	isBadNullValue := false
+	if inputDatum != nil {
+		value, err = table.CastValue(kvcodec.se, *inputDatum, col.ToInfo(), false, false)
+		if err != nil {
+			return value, err
+		}
+		if err := col.CheckNotNull(&value); err == nil {
+			return value, nil // the most normal case
+		}
+		isBadNullValue = true
+	}
+	// handle special values
+	switch {
+	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()):
+		var val types.Datum
+		realRowID := kvcodec.autoIDFn(rowID)
+		if mysql.HasUnsignedFlag(col.GetFlag()) {
+			val = types.NewUintDatum(uint64(realRowID))
+		} else {
+			val = types.NewIntDatum(realRowID)
+		}
+		value, err = table.CastValue(kvcodec.se, val, col.ToInfo(), false, false)
+	case col.IsGenerated():
+		// inject some dummy value for gen col so that MutRowFromDatums below sees a real value instead of nil.
+		// if MutRowFromDatums sees a nil it won't initialize the underlying storage and cause SetDatum to panic.
+		value = types.GetMinValue(&col.FieldType)
+	case isBadNullValue:
+		err = col.HandleBadNull(&value, kvcodec.se.vars.StmtCtx)
+	default:
+		value, err = table.GetColDefaultValue(kvcodec.se, col.ToInfo())
+	}
+	return value, err
+}
+
 // get record value for auto-increment field
 //
 // See: https://github.com/pingcap/tidb/blob/47f0f15b14ed54fc2222f3e304e29df7b05e6805/executor/insert_common.go#L781-L852
diff --git a/br/pkg/lightning/backend/kv/sql2kv_test.go b/br/pkg/lightning/backend/kv/sql2kv_test.go
index db749e7dfd5eb..a202b39bf56fb 100644
--- a/br/pkg/lightning/backend/kv/sql2kv_test.go
+++ b/br/pkg/lightning/backend/kv/sql2kv_test.go
@@ -322,11 +322,17 @@ func TestEncodeDoubleAutoIncrement(t *testing.T) {
 		},
 	})
 	require.NoError(t, err)
-	pairs, err := encoder.Encode(logger, []types.Datum{
-		types.NewStringDatum("1"),
+
+	strDatumForID := types.NewStringDatum("1")
+	actualDatum, err := encoder.(*tableKVEncoder).getActualDatum(70, 0, &strDatumForID)
+	require.NoError(t, err)
+	require.Equal(t, types.NewFloat64Datum(1.0), actualDatum)
+
+	pairsExpect, err := encoder.Encode(logger, []types.Datum{
+		types.NewFloat64Datum(1.0),
 	}, 70, []int{0, -1}, "1.csv", 1234)
 	require.NoError(t, err)
-	require.Equal(t, pairs, &KvPairs{pairs: []common.KvPair{
+	require.Equal(t, &KvPairs{pairs: []common.KvPair{
 		{
 			Key:   []uint8{0x74, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x5f, 0x72, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x46},
 			Val:   []uint8{0x80, 0x0, 0x1, 0x0, 0x0, 0x0, 0x1, 0x8, 0x0, 0xbf, 0xf0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
@@ -337,10 +343,88 @@ func TestEncodeDoubleAutoIncrement(t *testing.T) {
 			Val:   []uint8{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x46},
 			RowID: 70,
 		},
-	}})
+	}}, pairsExpect)
+
+	pairs, err := encoder.Encode(logger, []types.Datum{
+		types.NewStringDatum("1"),
+	}, 70, []int{0, -1}, "1.csv", 1234)
+	require.NoError(t, err)
+
+	require.Equal(t, pairsExpect, pairs)
 	require.Equal(t, tbl.Allocators(encoder.(*tableKVEncoder).se).Get(autoid.AutoIncrementType).Base(), int64(70))
 }
 
+func TestEncodeMissingAutoValue(t *testing.T) {
+	logger := log.Logger{Logger: zap.NewNop()}
+
+	var rowID int64 = 70
+	type testTableInfo struct {
+		AllocType  autoid.AllocatorType
+		CreateStmt string
+	}
+
+	for _, testTblInfo := range []testTableInfo{
+		{
+			AllocType:  autoid.AutoIncrementType,
+			CreateStmt: "create table t (id integer primary key auto_increment);",
+		},
+		{
+			AllocType:  autoid.AutoRandomType,
+			CreateStmt: "create table t (id integer primary key auto_random(3));",
+		},
+	} {
+		tblInfo := mockTableInfo(t, testTblInfo.CreateStmt)
+		if testTblInfo.AllocType == autoid.AutoRandomType {
+			// seems parser can't parse auto_random properly.
+			tblInfo.AutoRandomBits = 3
+		}
+		tbl, err := tables.TableFromMeta(NewPanickingAllocators(0), tblInfo)
+		require.NoError(t, err)
+
+		encoder, err := NewTableKVEncoder(tbl, &SessionOptions{
+			SQLMode: mysql.ModeStrictAllTables,
+			SysVars: map[string]string{
+				"tidb_row_format_version": "2",
+			},
+		})
+		require.NoError(t, err)
+
+		realRowID := encoder.(*tableKVEncoder).autoIDFn(rowID)
+
+		var nullDatum types.Datum
+		nullDatum.SetNull()
+
+		expectIDDatum := types.NewIntDatum(realRowID)
+		actualIDDatum, err := encoder.(*tableKVEncoder).getActualDatum(rowID, 0, nil)
+		require.NoError(t, err)
+		require.Equal(t, expectIDDatum, actualIDDatum)
+
+		actualIDDatum, err = encoder.(*tableKVEncoder).getActualDatum(rowID, 0, &nullDatum)
+		require.NoError(t, err)
+		require.Equal(t, expectIDDatum, actualIDDatum)
+
+		pairsExpect, err := encoder.Encode(logger, []types.Datum{
+			types.NewIntDatum(realRowID),
+		}, rowID, []int{0}, "1.csv", 1234)
+		require.NoError(t, err)
+
+		// test insert a NULL value on auto_xxxx column, and it is set to NOT NULL
+		pairs, err := encoder.Encode(logger, []types.Datum{
+			nullDatum,
+		}, rowID, []int{0}, "1.csv", 1234)
+		require.NoError(t, err)
+		require.Equalf(t, pairsExpect, pairs, "test table info: %+v", testTblInfo)
+		require.Equalf(t, rowID, tbl.Allocators(encoder.(*tableKVEncoder).se).Get(testTblInfo.AllocType).Base(), "test table info: %+v", testTblInfo)
+
+		// test insert a row without specifying the auto_xxxx column
+		pairs, err = encoder.Encode(logger, []types.Datum{}, rowID, []int{0}, "1.csv", 1234)
+		require.NoError(t, err)
+		require.Equalf(t, pairsExpect, pairs, "test table info: %+v", testTblInfo)
+		require.Equalf(t, rowID, tbl.Allocators(encoder.(*tableKVEncoder).se).Get(testTblInfo.AllocType).Base(), "test table info: %+v", testTblInfo)
+
+	}
+}
+
 func mockTableInfo(t *testing.T, createSQL string) *model.TableInfo {
 	parser := parser.New()
 	node, err := parser.ParseOneStmt(createSQL, "", "")
diff --git a/br/tests/lightning_auto_columns/config.toml b/br/tests/lightning_auto_columns/config.toml
new file mode 100644
index 0000000000000..261c42def21e0
--- /dev/null
+++ b/br/tests/lightning_auto_columns/config.toml
@@ -0,0 +1,2 @@
+[tikv-importer]
+backend = "local"
diff --git a/br/tests/lightning_auto_columns/data/lightning_auto_cols-schema-create.sql b/br/tests/lightning_auto_columns/data/lightning_auto_cols-schema-create.sql
new file mode 100644
index 0000000000000..396acc94b7839
--- /dev/null
+++ b/br/tests/lightning_auto_columns/data/lightning_auto_cols-schema-create.sql
@@ -0,0 +1 @@
+create schema lightning_auto_cols;
diff --git a/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_incr-schema.sql b/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_incr-schema.sql
new file mode 100644
index 0000000000000..db58c901e1d40
--- /dev/null
+++ b/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_incr-schema.sql
@@ -0,0 +1,3 @@
+CREATE TABLE t_auto_incr (
+id bigint PRIMARY KEY AUTO_INCREMENT,
+c char(40) NOT NULL DEFAULT '');
diff --git a/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_incr.sql b/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_incr.sql
new file mode 100644
index 0000000000000..40fbec9a3e719
--- /dev/null
+++ b/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_incr.sql
@@ -0,0 +1,8 @@
+INSERT INTO t_auto_incr (id, c) VALUES
+(1, 'normal_pk_01');
+INSERT INTO t_auto_incr VALUES
+(NULL, 'null_pk_02');
+INSERT INTO t_auto_incr VALUES
+(NULL, 'null_pk_03');
+INSERT INTO t_auto_incr (id, c) VALUES
+(4, 'normal_pk_04');
diff --git a/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_random-schema.sql b/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_random-schema.sql
new file mode 100644
index 0000000000000..972759243b85f
--- /dev/null
+++ b/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_random-schema.sql
@@ -0,0 +1,3 @@
+CREATE TABLE t_auto_random (
+id bigint PRIMARY KEY AUTO_RANDOM(3),
+c char(40) NOT NULL DEFAULT '');
diff --git a/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_random.sql b/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_random.sql
new file mode 100644
index 0000000000000..3a0fe765278f9
--- /dev/null
+++ b/br/tests/lightning_auto_columns/data/lightning_auto_cols.t_auto_random.sql
@@ -0,0 +1,8 @@
+INSERT INTO t_auto_random (id, c) VALUES
+(1, 'normal_pk_01');
+INSERT INTO t_auto_random VALUES
+(NULL, 'null_pk_02');
+INSERT INTO t_auto_random VALUES
+(NULL, 'null_pk_03');
+INSERT INTO t_auto_random (id, c) VALUES
+(4, 'normal_pk_04');
diff --git a/br/tests/lightning_auto_columns/run.sh b/br/tests/lightning_auto_columns/run.sh
new file mode 100755
index 0000000000000..471fa179d0973
--- /dev/null
+++ b/br/tests/lightning_auto_columns/run.sh
@@ -0,0 +1,38 @@
+#!/bin/sh
+#
+# Copyright 2020 PingCAP, Inc.
+#
+# 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.
+
+set -eux
+
+run_sql 'DROP DATABASE IF EXISTS lightning_auto_cols;'
+run_lightning
+
+run_sql "SELECT CONCAT_WS(':', id, c) AS row_data FROM lightning_auto_cols.t_auto_incr;"
+check_contains "row_data: 1:normal_pk_01"
+check_contains "row_data: 2:null_pk_02"
+check_contains "row_data: 3:null_pk_03"
+check_contains "row_data: 4:normal_pk_04"
+run_sql "SELECT COUNT(*) AS row_count FROM lightning_auto_cols.t_auto_incr;"
+check_contains "row_count: 4"
+
+run_sql "SELECT CONCAT_WS(':', id, c) AS row_data FROM lightning_auto_cols.t_auto_random;"
+check_contains "row_data: 1:normal_pk_01"
+check_contains ":null_pk_02"
+check_not_contains "row_data: 0:null_pk_02"
+check_contains ":null_pk_03"
+check_not_contains "row_data: 0:null_pk_03"
+check_contains "row_data: 4:normal_pk_04"
+run_sql "SELECT COUNT(*) AS row_count FROM lightning_auto_cols.t_auto_random;"
+check_contains "row_count: 4"