Repository: pingcap/tidb-lightning (go)

Issue:
Recognize `block_encryption_mode` from downstream
Lightning as of 4.0.9 always use the default `block_encryption_mode`. Since `block_encryption_mode` affects result of `aes_encrypt()` and `aes_decrypt()`, using the wrong mode may produce incorrect result for stored generated columns.

We should always copy the `block_encryption_mode` from downstream, just like `tidb_row_format_version`.

Pull request description:
backend,restore: duplicate more important system variables from downstream
<!--
Thank you for contributing to TiDB-Lightning! Please read the [CONTRIBUTING](https://github.com/pingcap/tidb-lightning/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->

Fix #504.

### What is changed and how it works?

Copy all sys vars that affects expression evaluation results in additional to `tidb_row_format_version`.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test (will add integration test after #505 is merged)

Side effects

Related changes



Patch:
diff --git a/lightning/backend/session.go b/lightning/backend/session.go
index b5a53f4..3087229 100644
--- a/lightning/backend/session.go
+++ b/lightning/backend/session.go
@@ -176,9 +176,9 @@ type session struct {
 
 // SessionOptions is the initial configuration of the session.
 type SessionOptions struct {
-	SQLMode          mysql.SQLMode
-	Timestamp        int64
-	RowFormatVersion string
+	SQLMode   mysql.SQLMode
+	Timestamp int64
+	SysVars   map[string]string
 	// a seed used for tableKvEncoder's auto random bits value
 	AutoRandomSeed int64
 }
@@ -194,10 +194,13 @@ func newSession(options *SessionOptions) *session {
 	vars.StmtCtx.OverflowAsWarning = !sqlMode.HasStrictMode()
 	vars.StmtCtx.AllowInvalidDate = sqlMode.HasAllowInvalidDatesMode()
 	vars.StmtCtx.IgnoreZeroInDate = !sqlMode.HasStrictMode() || sqlMode.HasAllowInvalidDatesMode()
+	if options.SysVars != nil {
+		for k, v := range options.SysVars {
+			vars.SetSystemVar(k, v)
+		}
+	}
 	vars.StmtCtx.TimeZone = vars.Location()
 	vars.SetSystemVar("timestamp", strconv.FormatInt(options.Timestamp, 10))
-	vars.SetSystemVar(variable.TiDBRowFormatVersion, options.RowFormatVersion)
-	vars.SetSystemVar(variable.MaxAllowedPacket, strconv.FormatUint(variable.MaxOfMaxAllowedPacket, 10))
 	vars.TxnCtx = nil
 
 	s := &session{
diff --git a/lightning/backend/sql2kv.go b/lightning/backend/sql2kv.go
index 90dd865..846dcda 100644
--- a/lightning/backend/sql2kv.go
+++ b/lightning/backend/sql2kv.go
@@ -157,7 +157,7 @@ func collectGeneratedColumns(se *session, meta *model.TableInfo, cols []*table.C
 
 	// order the result by column offset so they match the evaluation order.
 	sort.Slice(genCols, func(i, j int) bool {
-		return cols[i].Offset < cols[j].Offset
+		return cols[genCols[i].index].Offset < cols[genCols[j].index].Offset
 	})
 	return genCols, nil
 }
diff --git a/lightning/glue/glue.go b/lightning/glue/glue.go
index e4e0458..6caf38e 100644
--- a/lightning/glue/glue.go
+++ b/lightning/glue/glue.go
@@ -44,6 +44,7 @@ type SQLExecutor interface {
 	// same underlying connection
 	ExecuteWithLog(ctx context.Context, query string, purpose string, logger log.Logger) error
 	ObtainStringWithLog(ctx context.Context, query string, purpose string, logger log.Logger) (string, error)
+	QueryStringsWithLog(ctx context.Context, query string, purpose string, logger log.Logger) ([][]string, error)
 	Close()
 }
 
@@ -80,6 +81,38 @@ func (e *ExternalTiDBGlue) ObtainStringWithLog(ctx context.Context, query string
 	return s, err
 }
 
+func (e *ExternalTiDBGlue) QueryStringsWithLog(ctx context.Context, query string, purpose string, logger log.Logger) (result [][]string, finalErr error) {
+	finalErr = common.SQLWithRetry{
+		DB:     e.db,
+		Logger: logger,
+	}.Transact(ctx, purpose, func(c context.Context, tx *sql.Tx) (txErr error) {
+		rows, err := tx.QueryContext(c, query)
+		if err != nil {
+			return err
+		}
+		defer rows.Close()
+
+		colNames, err := rows.Columns()
+		if err != nil {
+			return err
+		}
+		for rows.Next() {
+			row := make([]string, len(colNames))
+			refs := make([]interface{}, 0, len(row))
+			for i := range row {
+				refs = append(refs, &row[i])
+			}
+			if err := rows.Scan(refs...); err != nil {
+				return err
+			}
+			result = append(result, row)
+		}
+
+		return rows.Err()
+	})
+	return
+}
+
 func (e *ExternalTiDBGlue) GetDB() (*sql.DB, error) {
 	return e.db, nil
 }
diff --git a/lightning/restore/restore.go b/lightning/restore/restore.go
index baf3e17..20f74b0 100644
--- a/lightning/restore/restore.go
+++ b/lightning/restore/restore.go
@@ -145,7 +145,7 @@ type RestoreController struct {
 	postProcessLock sync.Mutex // a simple way to ensure post-processing is not concurrent without using complicated goroutines
 	alterTableLock  sync.Mutex
 	compactState    int32
-	rowFormatVer    string
+	sysVars         map[string]string
 	tls             *common.TLS
 
 	errorSummaries errorSummaries
@@ -230,7 +230,7 @@ func NewRestoreControllerWithPauser(
 		pauser:        pauser,
 		backend:       backend,
 		tidbGlue:      g,
-		rowFormatVer:  "1",
+		sysVars:       defaultImportantVariables,
 		tls:           tls,
 
 		errorSummaries:    makeErrorSummaries(log.L()),
@@ -362,7 +362,7 @@ func (rc *RestoreController) restoreSchema(ctx context.Context) error {
 
 	go rc.listenCheckpointUpdates()
 
-	rc.rowFormatVer = ObtainRowFormatVersion(ctx, rc.tidbGlue.GetSQLExecutor())
+	rc.sysVars = ObtainImportantVariables(ctx, rc.tidbGlue.GetSQLExecutor())
 
 	// Estimate the number of chunks for progress reporting
 	err = rc.estimateChunkCountIntoMetrics(ctx)
@@ -1928,9 +1928,9 @@ func (cr *chunkRestore) restore(
 ) error {
 	// Create the encoder.
 	kvEncoder, err := rc.backend.NewEncoder(t.encTable, &kv.SessionOptions{
-		SQLMode:          rc.cfg.TiDB.SQLMode,
-		Timestamp:        cr.chunk.Timestamp,
-		RowFormatVersion: rc.rowFormatVer,
+		SQLMode:   rc.cfg.TiDB.SQLMode,
+		Timestamp: cr.chunk.Timestamp,
+		SysVars:   rc.sysVars,
 		// use chunk.PrevRowIDMax as the auto random seed, so it can stay the same value after recover from checkpoint.
 		AutoRandomSeed: cr.chunk.Chunk.PrevRowIDMax,
 	})
diff --git a/lightning/restore/tidb.go b/lightning/restore/tidb.go
index 227b843..d2c0fbb 100644
--- a/lightning/restore/tidb.go
+++ b/lightning/restore/tidb.go
@@ -40,6 +40,22 @@ import (
 	"go.uber.org/zap"
 )
 
+var (
+	// defaultImportantVariables is used in ObtainImportantVariables to retrieve the system
+	// variables from downstream which may affect KV encode result. The values record the default
+	// values if missing.
+	defaultImportantVariables = map[string]string{
+		"tidb_row_format_version": "1",
+		"max_allowed_packet":      "67108864",
+		"div_precision_increment": "4",
+		"time_zone":               "SYSTEM",
+		"lc_time_names":           "en_US",
+		"default_week_format":     "0",
+		"block_encryption_mode":   "aes-128-ecb",
+		"group_concat_max_len":    "1024",
+	}
+)
+
 type TiDBManager struct {
 	db     *sql.DB
 	parser *parser.Parser
@@ -271,17 +287,37 @@ func UpdateGCLifeTime(ctx context.Context, db *sql.DB, gcLifeTime string) error
 	)
 }
 
-func ObtainRowFormatVersion(ctx context.Context, g glue.SQLExecutor) string {
-	rowFormatVersion, err := g.ObtainStringWithLog(
-		ctx,
-		"SELECT @@tidb_row_format_version",
-		"obtain row format version",
-		log.L(),
-	)
+func ObtainImportantVariables(ctx context.Context, g glue.SQLExecutor) map[string]string {
+	var query strings.Builder
+	query.WriteString("SHOW VARIABLES WHERE Variable_name IN ('")
+	first := true
+	for k := range defaultImportantVariables {
+		if first {
+			first = false
+		} else {
+			query.WriteString("','")
+		}
+		query.WriteString(k)
+	}
+	query.WriteString("')")
+	kvs, err := g.QueryStringsWithLog(ctx, query.String(), "obtain system variables", log.L())
 	if err != nil {
-		rowFormatVersion = "1"
+		// error is not fatal
+		log.L().Warn("obtain system variables failed, use default variables instead", log.ShortError(err))
+	}
+
+	// convert result into a map. fill in any missing variables with default values.
+	result := make(map[string]string, len(defaultImportantVariables))
+	for _, kv := range kvs {
+		result[kv[0]] = kv[1]
 	}
-	return rowFormatVersion
+	for k, defV := range defaultImportantVariables {
+		if _, ok := result[k]; !ok {
+			result[k] = defV
+		}
+	}
+
+	return result
 }
 
 func ObtainNewCollationEnabled(ctx context.Context, g glue.SQLExecutor) bool {


Test patch:
diff --git a/lightning/backend/backend_test.go b/lightning/backend/backend_test.go
index e111ad9..46f1c36 100644
--- a/lightning/backend/backend_test.go
+++ b/lightning/backend/backend_test.go
@@ -321,7 +321,7 @@ func (s *backendSuite) TestNewEncoder(c *C) {
 	defer s.tearDownTest()
 
 	encoder := mock.NewMockEncoder(s.controller)
-	options := &kv.SessionOptions{SQLMode: mysql.ModeANSIQuotes, Timestamp: 1234567890, RowFormatVersion: "1"}
+	options := &kv.SessionOptions{SQLMode: mysql.ModeANSIQuotes, Timestamp: 1234567890}
 	s.mockBackend.EXPECT().NewEncoder(nil, options).Return(encoder, nil)
 
 	realEncoder, err := s.mockBackend.NewEncoder(nil, options)
diff --git a/lightning/backend/session_test.go b/lightning/backend/session_test.go
index 605997a..743b60b 100644
--- a/lightning/backend/session_test.go
+++ b/lightning/backend/session_test.go
@@ -30,7 +30,7 @@ func TestKV(t *testing.T) {
 }
 
 func (s *kvSuite) TestSetOption(c *C) {
-	session := newSession(&SessionOptions{SQLMode: mysql.ModeNone, Timestamp: 1234567890, RowFormatVersion: "1"})
+	session := newSession(&SessionOptions{SQLMode: mysql.ModeNone, Timestamp: 1234567890})
 	txn, err := session.Txn(true)
 	c.Assert(err, IsNil)
 	txn.SetOption(tidbkv.Priority, tidbkv.PriorityHigh)
diff --git a/lightning/backend/sql2kv_test.go b/lightning/backend/sql2kv_test.go
index 0a18dec..75955ff 100644
--- a/lightning/backend/sql2kv_test.go
+++ b/lightning/backend/sql2kv_test.go
@@ -29,7 +29,6 @@ import (
 	"github.com/pingcap/tidb/table/tables"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/mock"
-	"github.com/pingcap/tidb/util/timeutil"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 
@@ -84,9 +83,8 @@ func (s *kvSuite) TestEncode(c *C) {
 
 	// Strict mode
 	strictMode, err := NewTableKVEncoder(tbl, &SessionOptions{
-		SQLMode:          mysql.ModeStrictAllTables,
-		Timestamp:        1234567890,
-		RowFormatVersion: "1",
+		SQLMode:   mysql.ModeStrictAllTables,
+		Timestamp: 1234567890,
 	})
 	c.Assert(err, IsNil)
 	pairs, err := strictMode.Encode(logger, rows, 1, []int{0, 1})
@@ -116,9 +114,8 @@ func (s *kvSuite) TestEncode(c *C) {
 	// Mock add record error
 	mockTbl := &mockTable{Table: tbl}
 	mockMode, err := NewTableKVEncoder(mockTbl, &SessionOptions{
-		SQLMode:          mysql.ModeStrictAllTables,
-		Timestamp:        1234567891,
-		RowFormatVersion: "1",
+		SQLMode:   mysql.ModeStrictAllTables,
+		Timestamp: 1234567891,
 	})
 	c.Assert(err, IsNil)
 	pairs, err = mockMode.Encode(logger, rowsWithPk2, 2, []int{0, 1})
@@ -126,9 +123,9 @@ func (s *kvSuite) TestEncode(c *C) {
 
 	// Non-strict mode
 	noneMode, err := NewTableKVEncoder(tbl, &SessionOptions{
-		SQLMode:          mysql.ModeNone,
-		Timestamp:        1234567892,
-		RowFormatVersion: "1",
+		SQLMode:   mysql.ModeNone,
+		Timestamp: 1234567892,
+		SysVars:   map[string]string{"tidb_row_format_version": "1"},
 	})
 	c.Assert(err, IsNil)
 	pairs, err = noneMode.Encode(logger, rows, 1, []int{0, 1})
@@ -156,9 +153,9 @@ func (s *kvSuite) TestEncodeRowFormatV2(c *C) {
 	}
 
 	noneMode, err := NewTableKVEncoder(tbl, &SessionOptions{
-		SQLMode:          mysql.ModeNone,
-		Timestamp:        1234567892,
-		RowFormatVersion: "2",
+		SQLMode:   mysql.ModeNone,
+		Timestamp: 1234567892,
+		SysVars:   map[string]string{"tidb_row_format_version": "2"},
 	})
 	c.Assert(err, IsNil)
 	pairs, err := noneMode.Encode(logger, rows, 1, []int{0, 1})
@@ -199,11 +196,13 @@ func (s *kvSuite) TestEncodeTimestamp(c *C) {
 
 	logger := log.Logger{Logger: zap.NewNop()}
 
-	timeutil.SetSystemTZ("Etc/GMT-8") // force timezone to be UTC+08:00.
 	encoder, err := NewTableKVEncoder(tbl, &SessionOptions{
-		SQLMode:          mysql.ModeStrictAllTables,
-		Timestamp:        1234567893,
-		RowFormatVersion: "1",
+		SQLMode:   mysql.ModeStrictAllTables,
+		Timestamp: 1234567893,
+		SysVars: map[string]string{
+			"tidb_row_format_version": "1",
+			"time_zone":               "+08:00",
+		},
 	})
 	c.Assert(err, IsNil)
 	pairs, err := encoder.Encode(logger, nil, 70, []int{-1, 1})
@@ -234,10 +233,10 @@ func (s *kvSuite) TestDefaultAutoRandoms(c *C) {
 	tbl, err := tables.TableFromMeta(NewPanickingAllocators(0), tblInfo)
 	c.Assert(err, IsNil)
 	encoder, err := NewTableKVEncoder(tbl, &SessionOptions{
-		SQLMode:          mysql.ModeStrictAllTables,
-		Timestamp:        1234567893,
-		RowFormatVersion: "2",
-		AutoRandomSeed:   456,
+		SQLMode:        mysql.ModeStrictAllTables,
+		Timestamp:      1234567893,
+		SysVars:        map[string]string{"tidb_row_format_version": "2"},
+		AutoRandomSeed: 456,
 	})
 	c.Assert(err, IsNil)
 	logger := log.Logger{Logger: zap.NewNop()}
@@ -399,8 +398,7 @@ func (s *benchSQL2KVSuite) SetUpTest(c *C) {
 	// Construct the corresponding KV encoder.
 	tbl, err := tables.TableFromMeta(NewPanickingAllocators(0), tableInfo)
 	c.Assert(err, IsNil)
-	s.encoder, err = NewTableKVEncoder(tbl, &SessionOptions{RowFormatVersion: "2"})
-	c.Assert(err, IsNil)
+	s.encoder, err = NewTableKVEncoder(tbl, &SessionOptions{SysVars: map[string]string{"tidb_row_format_version": "2"}})
 	s.logger = log.Logger{Logger: zap.NewNop()}
 
 	// Prepare the row to insert.
diff --git a/lightning/backend/tidb_test.go b/lightning/backend/tidb_test.go
index 5b0ceb6..f02fda8 100644
--- a/lightning/backend/tidb_test.go
+++ b/lightning/backend/tidb_test.go
@@ -91,7 +91,7 @@ func (s *mysqlSuite) TestWriteRowsReplaceOnDup(c *C) {
 		perms = append(perms, i)
 	}
 	perms = append(perms, -1)
-	encoder, err := s.backend.NewEncoder(s.tbl, &kv.SessionOptions{SQLMode: 0, Timestamp: 1234567890, RowFormatVersion: "1"})
+	encoder, err := s.backend.NewEncoder(s.tbl, &kv.SessionOptions{SQLMode: 0, Timestamp: 1234567890})
 	c.Assert(err, IsNil)
 	row, err := encoder.Encode(logger, []types.Datum{
 		types.NewUintDatum(18446744073709551615),
diff --git a/lightning/restore/restore_test.go b/lightning/restore/restore_test.go
index 4e4896c..92dcb55 100644
--- a/lightning/restore/restore_test.go
+++ b/lightning/restore/restore_test.go
@@ -927,9 +927,8 @@ func (s *chunkRestoreSuite) TestEncodeLoop(c *C) {
 	kvsCh := make(chan []deliveredKVs, 2)
 	deliverCompleteCh := make(chan deliverResult)
 	kvEncoder, err := kv.NewTableKVEncoder(s.tr.encTable, &kv.SessionOptions{
-		SQLMode:          s.cfg.TiDB.SQLMode,
-		Timestamp:        1234567895,
-		RowFormatVersion: "1",
+		SQLMode:   s.cfg.TiDB.SQLMode,
+		Timestamp: 1234567895,
 	})
 	c.Assert(err, IsNil)
 	cfg := config.NewConfig()
@@ -953,9 +952,8 @@ func (s *chunkRestoreSuite) TestEncodeLoopCanceled(c *C) {
 	kvsCh := make(chan []deliveredKVs)
 	deliverCompleteCh := make(chan deliverResult)
 	kvEncoder, err := kv.NewTableKVEncoder(s.tr.encTable, &kv.SessionOptions{
-		SQLMode:          s.cfg.TiDB.SQLMode,
-		Timestamp:        1234567896,
-		RowFormatVersion: "1",
+		SQLMode:   s.cfg.TiDB.SQLMode,
+		Timestamp: 1234567896,
 	})
 	c.Assert(err, IsNil)
 
@@ -972,9 +970,8 @@ func (s *chunkRestoreSuite) TestEncodeLoopForcedError(c *C) {
 	kvsCh := make(chan []deliveredKVs, 2)
 	deliverCompleteCh := make(chan deliverResult)
 	kvEncoder, err := kv.NewTableKVEncoder(s.tr.encTable, &kv.SessionOptions{
-		SQLMode:          s.cfg.TiDB.SQLMode,
-		Timestamp:        1234567897,
-		RowFormatVersion: "1",
+		SQLMode:   s.cfg.TiDB.SQLMode,
+		Timestamp: 1234567897,
 	})
 	c.Assert(err, IsNil)
 
@@ -993,9 +990,8 @@ func (s *chunkRestoreSuite) TestEncodeLoopDeliverErrored(c *C) {
 	kvsCh := make(chan []deliveredKVs)
 	deliverCompleteCh := make(chan deliverResult)
 	kvEncoder, err := kv.NewTableKVEncoder(s.tr.encTable, &kv.SessionOptions{
-		SQLMode:          s.cfg.TiDB.SQLMode,
-		Timestamp:        1234567898,
-		RowFormatVersion: "1",
+		SQLMode:   s.cfg.TiDB.SQLMode,
+		Timestamp: 1234567898,
 	})
 	c.Assert(err, IsNil)
 
diff --git a/lightning/restore/tidb_test.go b/lightning/restore/tidb_test.go
index 4e75c16..4b32834 100644
--- a/lightning/restore/tidb_test.go
+++ b/lightning/restore/tidb_test.go
@@ -370,26 +370,60 @@ func (s *tidbSuite) TestObtainRowFormatVersionSucceed(c *C) {
 	ctx := context.Background()
 
 	s.mockDB.
-		ExpectQuery("\\QSELECT @@tidb_row_format_version\\E").
-		WillReturnRows(sqlmock.NewRows([]string{"@@tidb_row_format_version"}).AddRow("2"))
+		ExpectBegin()
+	s.mockDB.
+		ExpectQuery(`SHOW VARIABLES WHERE Variable_name IN \(.*'tidb_row_format_version'.*\)`).
+		WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).
+			AddRow("tidb_row_format_version", "2").
+			AddRow("max_allowed_packet", "1073741824").
+			AddRow("div_precision_increment", "10").
+			AddRow("time_zone", "-08:00").
+			AddRow("lc_time_names", "ja_JP").
+			AddRow("default_week_format", "1").
+			AddRow("block_encryption_mode", "aes-256-cbc").
+			AddRow("group_concat_max_len", "1073741824"))
+	s.mockDB.
+		ExpectCommit()
 	s.mockDB.
 		ExpectClose()
 
-	version := ObtainRowFormatVersion(ctx, s.tiGlue.GetSQLExecutor())
-	c.Assert(version, Equals, "2")
+	sysVars := ObtainImportantVariables(ctx, s.tiGlue.GetSQLExecutor())
+	c.Assert(sysVars, DeepEquals, map[string]string{
+		"tidb_row_format_version": "2",
+		"max_allowed_packet":      "1073741824",
+		"div_precision_increment": "10",
+		"time_zone":               "-08:00",
+		"lc_time_names":           "ja_JP",
+		"default_week_format":     "1",
+		"block_encryption_mode":   "aes-256-cbc",
+		"group_concat_max_len":    "1073741824",
+	})
 }
 
 func (s *tidbSuite) TestObtainRowFormatVersionFailure(c *C) {
 	ctx := context.Background()
 
 	s.mockDB.
-		ExpectQuery("\\QSELECT @@tidb_row_format_version\\E").
-		WillReturnError(errors.New("ERROR 1193 (HY000): Unknown system variable 'tidb_row_format_version'"))
+		ExpectBegin()
+	s.mockDB.
+		ExpectQuery(`SHOW VARIABLES WHERE Variable_name IN \(.*'tidb_row_format_version'.*\)`).
+		WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).AddRow("time_zone", "+00:00"))
+	s.mockDB.
+		ExpectCommit()
 	s.mockDB.
 		ExpectClose()
 
-	version := ObtainRowFormatVersion(ctx, s.tiGlue.GetSQLExecutor())
-	c.Assert(version, Equals, "1")
+	sysVars := ObtainImportantVariables(ctx, s.tiGlue.GetSQLExecutor())
+	c.Assert(sysVars, DeepEquals, map[string]string{
+		"tidb_row_format_version": "1",
+		"max_allowed_packet":      "67108864",
+		"div_precision_increment": "4",
+		"time_zone":               "+00:00",
+		"lc_time_names":           "en_US",
+		"default_week_format":     "0",
+		"block_encryption_mode":   "aes-128-ecb",
+		"group_concat_max_len":    "1024",
+	})
 }
 
 func (s *tidbSuite) TestObtainNewCollationEnabled(c *C) {
diff --git a/tests/_utils/read_result b/tests/_utils/read_result
new file mode 100755
index 0000000..96c1f90
--- /dev/null
+++ b/tests/_utils/read_result
@@ -0,0 +1,22 @@
+#!/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,
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Read the result of the last run_sql execution.
+# Requires the run_sql returning a single row and column and the column name has no spaces.
+
+set -eu
+TEST_DIR=/tmp/lightning_test_result
+
+tail -1 "$TEST_DIR/sql_res.$TEST_NAME.txt" | cut -d' ' -f 2-
diff --git a/tests/auto_random_default/run.sh b/tests/auto_random_default/run.sh
index 859c55e..04d46ca 100644
--- a/tests/auto_random_default/run.sh
+++ b/tests/auto_random_default/run.sh
@@ -37,10 +37,12 @@ for backend in tidb importer local; do
       check_contains 'inc: 4'
       check_contains 'inc: 5'
       check_contains 'inc: 6'
+      NEXT_AUTO_RAND_VAL=7
     else
       check_contains 'inc: 25'
       check_contains 'inc: 26'
       check_contains 'inc: 27'
+      NEXT_AUTO_RAND_VAL=28
     fi
 
 
@@ -48,11 +50,9 @@ for backend in tidb importer local; do
     check_contains "count: 2"
 
     # auto random base is 4
+    run_sql "SELECT max(id & b'000001111111111111111111111111111111111111111111111111111111111') >= $NEXT_AUTO_RAND_VAL as ge FROM auto_random.t"
+    check_contains 'ge: 0'
     run_sql "INSERT INTO auto_random.t VALUES ();"
-    run_sql "SELECT id & b'000001111111111111111111111111111111111111111111111111111111111' as inc FROM auto_random.t"
-    if [ "$backend" = 'tidb' ]; then
-      check_contains 'inc: 2000001'
-    else
-      check_contains 'inc: 28'
-    fi
+    run_sql "SELECT max(id & b'000001111111111111111111111111111111111111111111111111111111111') >= $NEXT_AUTO_RAND_VAL as ge FROM auto_random.t"
+    check_contains 'ge: 1'
 done
diff --git a/tests/generated_columns/data/gencol.various_types-schema.sql b/tests/generated_columns/data/gencol.various_types-schema.sql
index 839610d..a71cd79 100644
--- a/tests/generated_columns/data/gencol.various_types-schema.sql
+++ b/tests/generated_columns/data/gencol.various_types-schema.sql
@@ -11,5 +11,9 @@ create table various_types (
     bit bit(4) as (int64) stored,
     `set` set('a','b','c') as (enum) stored,
     time timestamp(3) as ('1987-06-05 04:03:02.100') stored,
-    json json as (json_object(string, float32)) stored
+    json json as (json_object(string, float32)) stored,
+    aes blob as (aes_encrypt(`decimal`, 'key', bytes)) stored, -- 0xA876B03CFC8AF93D22D19E2220BD2375, @@block_encryption_mode='aes-256-cbc'
+    -- FIXME: column below disabled due to pingcap/tidb#21510
+    -- week int as (week('2020-02-02')) stored, -- 6, @@default_week_format=4
+    tz varchar(20) as (from_unixtime(1)) stored -- 1969-12-31 16:00:01, @@time_zone='-08:00'
 );
diff --git a/tests/generated_columns/run.sh b/tests/generated_columns/run.sh
index b11f0af..9d02a4c 100644
--- a/tests/generated_columns/run.sh
+++ b/tests/generated_columns/run.sh
@@ -13,9 +13,28 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-set -eu
+set -eux
 
-for BACKEND in 'tidb' 'importer' 'local'; do
+run_sql "SELECT CONCAT('SET GLOBAL time_zone=''', @@time_zone, ''', GLOBAL default_week_format=', @@default_week_format, ', GLOBAL block_encryption_mode=''', @@block_encryption_mode, ''';') cmd;"
+UNDO_CMD=$(read_result)
+undo_set_globals() {
+    run_sql "$UNDO_CMD"
+}
+trap undo_set_globals EXIT
+
+# There is normally a 2 second delay between these SET GLOBAL statements returns
+# and the changes are actually effective. So we have this check-and-retry loop
+# below to ensure Lightning gets our desired global vars.
+run_sql "SET GLOBAL time_zone='-08:00', GLOBAL default_week_format=4, GLOBAL block_encryption_mode='aes-256-cbc'"
+for i in $(seq 3); do
+    sleep 1
+    run_sql "SELECT CONCAT(@@time_zone, ',', @@default_week_format, ',', @@block_encryption_mode) res"
+    if [ "$(read_result)" = '-08:00,4,aes-256-cbc' ]; then
+        break
+    fi
+done
+
+for BACKEND in 'local' 'tidb' 'importer'; do
     if [ "$BACKEND" = 'local' ]; then
         check_cluster_version 4 0 0 'local backend' || continue
     fi
@@ -45,4 +64,8 @@ for BACKEND in 'tidb' 'importer' 'local'; do
     check_contains 'set: c'
     check_contains 'time: 1987-06-05 04:03:02.100'
     check_contains 'json: {"6ad8402ba6610f04d3ec5c9875489a7bc8e259c5": 0.5625}'
+    check_contains 'aes: 0xA876B03CFC8AF93D22D19E2220BD2375'
+    # FIXME: test below disabled due to pingcap/tidb#21510
+    # check_contains 'week: 6'
+    check_contains 'tz: 1969-12-31 16:00:01'
 done