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

) close #44497
pingcap · Jul 12, 2023 · a5ad6d3 · a5ad6d3
1 parent da8ef8b
commit a5ad6d3
Show file tree

Hide file tree

Showing 6 changed files with 56 additions and 1 deletion.
diff --git a/br/pkg/lightning/backend/kv/BUILD.bazel b/br/pkg/lightning/backend/kv/BUILD.bazel
@@ -55,7 +55,7 @@ go_test(
     embed = [":kv"],
     flaky = True,
     race = "on",
-    shard_count = 15,
+    shard_count = 16,
     deps = [
         "//br/pkg/lightning/backend/encode",
         "//br/pkg/lightning/common",
@@ -68,6 +68,7 @@ go_test(
         "//parser/ast",
         "//parser/model",
         "//parser/mysql",
+        "//planner/core",
         "//sessionctx",
         "//table",
         "//table/tables",

diff --git a/br/pkg/lightning/backend/kv/base.go b/br/pkg/lightning/backend/kv/base.go
@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/meta/autoid"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/table/tables"
 	"github.com/pingcap/tidb/types"
@@ -266,6 +267,14 @@ func (e *BaseKVEncoder) getActualDatum(col *table.Column, rowID int64, inputDatu
 	case isBadNullValue:
 		err = col.HandleBadNull(&value, e.SessionCtx.Vars.StmtCtx, 0)
 	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.
+			e.SessionCtx.Vars.TxnCtx = new(variable.TransactionContext)
+			defer func() {
+				e.SessionCtx.Vars.TxnCtx = nil
+			}()
+		}
 		value, err = table.GetColDefaultValue(e.SessionCtx, 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
@@ -32,6 +32,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"
@@ -455,6 +456,40 @@ 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(&encode.EncodingConfig{
+		Table: tbl,
+		SessionOptions: encode.SessionOptions{
+			SQLMode: mysql.ModeStrictAllTables,
+			SysVars: map[string]string{
+				"tidb_row_format_version": "2",
+			},
+		},
+		Logger: log.L(),
+	}, nil)
+	require.NoError(t, err)
+
+	strDatumForID := types.NewStringDatum("1")
+	actualDatum, err := lkv.GetActualDatum(encoder, tbl.Cols()[0], 70, &strDatumForID)
+	require.NoError(t, err)
+	require.Equal(t, strDatumForID, actualDatum)
+
+	actualDatum, err = lkv.GetActualDatum(encoder, tbl.Cols()[0], 70, 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, tbl.Cols()[0], 70, 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