From da1c21fd45a4ea5900ac16d2f4a248143f378d18 Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Tue, 19 Oct 2021 14:14:42 +0800
Subject: [PATCH] session: fix set session variable make tidb_snapshot unwork
 (#28677) (#28683)

---
 executor/set.go         |  7 +++++--
 session/session_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/executor/set.go b/executor/set.go
index e303efe1a5520..4689fb11fa99c 100644
--- a/executor/set.go
+++ b/executor/set.go
@@ -181,7 +181,7 @@ func (e *SetExecutor) setSysVariable(ctx context.Context, name string, v *expres
 		}
 	}
 
-	err = e.loadSnapshotInfoSchemaIfNeeded(newSnapshotTS)
+	err = e.loadSnapshotInfoSchemaIfNeeded(name, newSnapshotTS)
 	if err != nil {
 		fallbackOldSnapshotTS()
 		return err
@@ -254,7 +254,10 @@ func (e *SetExecutor) getVarValue(v *expression.VarAssignment, sysVar *variable.
 	return nativeVal.ToString()
 }
 
-func (e *SetExecutor) loadSnapshotInfoSchemaIfNeeded(snapshotTS uint64) error {
+func (e *SetExecutor) loadSnapshotInfoSchemaIfNeeded(name string, snapshotTS uint64) error {
+	if name != variable.TiDBSnapshot && name != variable.TiDBTxnReadTS {
+		return nil
+	}
 	vars := e.ctx.GetSessionVars()
 	if snapshotTS == 0 {
 		vars.SnapshotInfoschema = nil
diff --git a/session/session_test.go b/session/session_test.go
index 52743e98f52aa..02682bd3db288 100644
--- a/session/session_test.go
+++ b/session/session_test.go
@@ -5575,3 +5575,27 @@ func (s *testSessionSuite) TestLocalTemporaryTableUpdate(c *C) {
 		tk.MustQuery("select * from tmp1").Check(testkit.Rows())
 	}
 }
+
+func (s *testSessionSuite) TestFixSetTiDBSnapshotTS(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	safePointName := "tikv_gc_safe_point"
+	safePointValue := "20160102-15:04:05 -0700"
+	safePointComment := "All versions after safe point can be accessed. (DO NOT EDIT)"
+	updateSafePoint := fmt.Sprintf(`INSERT INTO mysql.tidb VALUES ('%[1]s', '%[2]s', '%[3]s')
+	ON DUPLICATE KEY
+	UPDATE variable_value = '%[2]s', comment = '%[3]s'`, safePointName, safePointValue, safePointComment)
+	tk.MustExec(updateSafePoint)
+	tk.MustExec("create database t123")
+	time.Sleep(time.Second)
+	ts := time.Now().Format("2006-1-2 15:04:05")
+	time.Sleep(time.Second)
+	tk.MustExec("drop database t123")
+	err := tk.ExecToErr("use t123")
+	c.Assert(err, NotNil)
+	c.Assert(err.Error(), Matches, ".*Unknown database.*")
+	tk.MustExec(fmt.Sprintf("set @@tidb_snapshot='%s'", ts))
+	tk.MustExec("use t123")
+	// update any session variable and assert whether infoschema is changed
+	tk.MustExec("SET SESSION sql_mode = 'STRICT_TRANS_TABLES,NO_AUTO_CREATE_USER';")
+	tk.MustExec("use t123")
+}