lightning: fix uuid panic when expression index enabled (#44516) (#44548

) close #44497
pingcap · Jun 9, 2023 · 43cf5ea · 43cf5ea
1 parent 07a6d94
commit 43cf5ea
Show file tree

Hide file tree

Showing 6 changed files with 50 additions and 0 deletions.
diff --git a/br/pkg/lightning/backend/kv/BUILD.bazel b/br/pkg/lightning/backend/kv/BUILD.bazel
@@ -64,6 +64,7 @@ go_test(
         "//parser/ast",
         "//parser/model",
         "//parser/mysql",
+        "//planner/core",
         "//sessionctx",
         "//table",
         "//table/tables",

diff --git a/br/pkg/lightning/backend/kv/sql2kv.go b/br/pkg/lightning/backend/kv/sql2kv.go
@@ -511,6 +511,14 @@ func (kvcodec *tableKVEncoder) getActualDatum(rowID int64, colIndex int, inputDa
 	case isBadNullValue:
 		err = col.HandleBadNull(&value, kvcodec.se.vars.StmtCtx)
 	default:
+		// copy from the following GetColDefaultValue function, when this is true it will use getColDefaultExprValue
+		if col.DefaultIsExpr {
+			// the expression rewriter requires a non-nil TxnCtx.
+			kvcodec.se.vars.TxnCtx = new(variable.TransactionContext)
+			defer func() {
+				kvcodec.se.vars.TxnCtx = nil
+			}()
+		}
 		value, err = table.GetColDefaultValue(kvcodec.se, col.ToInfo())
 	}
 	return value, err

diff --git a/br/pkg/lightning/backend/kv/sql2kv_test.go b/br/pkg/lightning/backend/kv/sql2kv_test.go
@@ -31,6 +31,7 @@ import (
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/parser/mysql"
+	_ "github.com/pingcap/tidb/planner/core" // to setup expression.EvalAstExpr. Otherwise we cannot parse the default value
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/table/tables"
@@ -430,6 +431,36 @@ func TestEncodeMissingAutoValue(t *testing.T) {
 	}
 }
 
+func TestEncodeExpressionColumn(t *testing.T) {
+	tblInfo := mockTableInfo(t, "create table t (id varchar(40) not null DEFAULT uuid(), unique key `u_id` (`id`));")
+	tbl, err := tables.TableFromMeta(lkv.NewPanickingAllocators(0), tblInfo)
+	require.NoError(t, err)
+
+	encoder, err := lkv.NewTableKVEncoder(tbl, &lkv.SessionOptions{
+		SQLMode: mysql.ModeStrictAllTables,
+		SysVars: map[string]string{
+			"tidb_row_format_version": "2",
+		},
+	}, nil, log.L())
+	require.NoError(t, err)
+
+	strDatumForID := types.NewStringDatum("1")
+	actualDatum, err := lkv.GetActualDatum(encoder, 70, 0, &strDatumForID)
+	require.NoError(t, err)
+	require.Equal(t, strDatumForID, actualDatum)
+
+	actualDatum, err = lkv.GetActualDatum(encoder, 70, 0, nil)
+	require.NoError(t, err)
+	require.Equal(t, types.KindString, actualDatum.Kind())
+	require.Len(t, actualDatum.GetString(), 36) // uuid length
+
+	actualDatum2, err := lkv.GetActualDatum(encoder, 70, 0, nil)
+	require.NoError(t, err)
+	require.Equal(t, types.KindString, actualDatum2.Kind())
+	require.Len(t, actualDatum2.GetString(), 36)
+	require.NotEqual(t, actualDatum.GetString(), actualDatum2.GetString()) // check different uuid
+}
+
 func mockTableInfo(t *testing.T, createSQL string) *model.TableInfo {
 	parser := parser.New()
 	node, err := parser.ParseOneStmt(createSQL, "", "")

diff --git a/br/tests/lightning_generated_columns/data/gencol.uuid-schema.sql b/br/tests/lightning_generated_columns/data/gencol.uuid-schema.sql
@@ -0,0 +1,5 @@
+CREATE TABLE `uuid` (
+  `a` bigint(30) NOT NULL,
+  `b` varchar(40) NOT NULL DEFAULT uuid(),
+  UNIQUE KEY `c` (`b`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;
diff --git a/br/tests/lightning_generated_columns/data/gencol.uuid.0.sql b/br/tests/lightning_generated_columns/data/gencol.uuid.0.sql
@@ -0,0 +1 @@
+insert into uuid (a) values (1), (10), (100), (1000);
diff --git a/br/tests/lightning_generated_columns/run.sh b/br/tests/lightning_generated_columns/run.sh
@@ -92,4 +92,8 @@ for BACKEND in 'local' 'tidb'; do
     check_contains 'id: 2'
     check_contains 'a: ABC'
     check_contains 'b: CDSFDS'
+
+    run_sql 'ADMIN CHECK TABLE gencol.uuid'
+    run_sql 'SELECT /*+ use_index(gencol.uuid, c) */ COUNT(DISTINCT(b)) FROM gencol.uuid'
+    check_contains '4'
 done
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		insert into uuid (a) values (1), (10), (100), (1000);