From 5209ac2fa1a3db278661588370786b2991728320 Mon Sep 17 00:00:00 2001
From: Song Gao <disxiaofei@163.com>
Date: Thu, 15 Sep 2022 18:15:00 +0800
Subject: [PATCH] sessionctx: support session var tidb_last_plan_replayer_token
 (#37851)

ref pingcap/tidb#37798
---
 executor/executor_test.go        | 13 +++++++------
 executor/plan_replayer.go        |  1 +
 server/plan_replayer_test.go     |  8 ++++++++
 sessionctx/variable/session.go   |  3 +++
 sessionctx/variable/sysvar.go    |  5 +++++
 sessionctx/variable/tidb_vars.go |  3 +++
 6 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/executor/executor_test.go b/executor/executor_test.go
index 8cfd71df003ca..e01eac6cf31c8 100644
--- a/executor/executor_test.go
+++ b/executor/executor_test.go
@@ -152,17 +152,18 @@ func TestPlanReplayer(t *testing.T) {
 	tk.MustExec("drop table if exists t")
 	tk.MustExec("create table t(a int, b int, index idx_a(a))")
 	tk.MustExec("alter table t set tiflash replica 1")
-	tk.MustExec("plan replayer dump explain select * from t where a=10")
-	tk.MustExec("plan replayer dump explain select /*+ read_from_storage(tiflash[t]) */ * from t")
+	tk.MustQuery("plan replayer dump explain select * from t where a=10")
+	tk.MustQuery("plan replayer dump explain select /*+ read_from_storage(tiflash[t]) */ * from t")
 
 	tk.MustExec("create table t1 (a int)")
 	tk.MustExec("create table t2 (a int)")
 	tk.MustExec("create definer=`root`@`127.0.0.1` view v1 as select * from t1")
 	tk.MustExec("create definer=`root`@`127.0.0.1` view v2 as select * from v1")
-	tk.MustExec("plan replayer dump explain with tmp as (select a from t1 group by t1.a) select * from tmp, t2 where t2.a=tmp.a;")
-	tk.MustExec("plan replayer dump explain select * from t1 where t1.a > (with cte1 as (select 1) select count(1) from cte1);")
-	tk.MustExec("plan replayer dump explain select * from v1")
-	tk.MustExec("plan replayer dump explain select * from v2")
+	tk.MustQuery("plan replayer dump explain with tmp as (select a from t1 group by t1.a) select * from tmp, t2 where t2.a=tmp.a;")
+	tk.MustQuery("plan replayer dump explain select * from t1 where t1.a > (with cte1 as (select 1) select count(1) from cte1);")
+	tk.MustQuery("plan replayer dump explain select * from v1")
+	tk.MustQuery("plan replayer dump explain select * from v2")
+	require.True(t, len(tk.Session().GetSessionVars().LastPlanReplayerToken) > 0)
 }
 
 func TestShow(t *testing.T) {
diff --git a/executor/plan_replayer.go b/executor/plan_replayer.go
index b213717d137c9..e780a223ff9f1 100644
--- a/executor/plan_replayer.go
+++ b/executor/plan_replayer.go
@@ -155,6 +155,7 @@ func (e *PlanReplayerExec) Next(ctx context.Context, req *chunk.Chunk) error {
 	}
 	req.AppendString(0, res)
 	e.endFlag = true
+	e.ctx.GetSessionVars().LastPlanReplayerToken = res
 	return nil
 }
 
diff --git a/server/plan_replayer_test.go b/server/plan_replayer_test.go
index 2ea91677d8438..2a00bc0db04de 100644
--- a/server/plan_replayer_test.go
+++ b/server/plan_replayer_test.go
@@ -136,5 +136,13 @@ func prepareData4PlanReplayer(t *testing.T, client *testServerClient, statHandle
 	var filename string
 	err = rows.Scan(&filename)
 	require.NoError(t, err)
+	rows.Close()
+	rows = tk.MustQuery("select @@tidb_last_plan_replayer_token")
+	require.True(t, rows.Next(), "unexpected data")
+	var filename2 string
+	err = rows.Scan(&filename2)
+	require.NoError(t, err)
+	rows.Close()
+	require.Equal(t, filename, filename2)
 	return filename
 }
diff --git a/sessionctx/variable/session.go b/sessionctx/variable/session.go
index ecb4b5b93c638..e7a2a96bd68b0 100644
--- a/sessionctx/variable/session.go
+++ b/sessionctx/variable/session.go
@@ -1270,6 +1270,9 @@ type SessionVars struct {
 	// ranges would exceed the limit, it chooses less accurate ranges such as full range. 0 indicates that there is no
 	// memory limit for ranges.
 	RangeMaxSize int64
+
+	// LastPlanReplayerToken indicates the last plan replayer token
+	LastPlanReplayerToken string
 }
 
 // GetPreparedStmtByName returns the prepared statement specified by stmtName.
diff --git a/sessionctx/variable/sysvar.go b/sessionctx/variable/sysvar.go
index b4b374c0e8761..3f27851e85bc2 100644
--- a/sessionctx/variable/sysvar.go
+++ b/sessionctx/variable/sysvar.go
@@ -333,6 +333,11 @@ var defaultSysVars = []*SysVar{
 		}
 		return string(info), nil
 	}},
+	{Scope: ScopeSession, Name: TiDBLastPlanReplayerToken, Value: "", ReadOnly: true,
+		GetSession: func(s *SessionVars) (string, error) {
+			return s.LastPlanReplayerToken, nil
+		},
+	},
 	/* The system variables below have INSTANCE scope  */
 	{Scope: ScopeInstance, Name: TiDBLogFileMaxDays, Value: strconv.Itoa(config.GetGlobalConfig().Log.File.MaxDays), Type: TypeInt, MinValue: 0, MaxValue: math.MaxInt32, SetGlobal: func(s *SessionVars, val string) error {
 		maxAge, err := strconv.ParseInt(val, 10, 32)
diff --git a/sessionctx/variable/tidb_vars.go b/sessionctx/variable/tidb_vars.go
index 7b77423ba7db0..e029a3ab8a3e0 100644
--- a/sessionctx/variable/tidb_vars.go
+++ b/sessionctx/variable/tidb_vars.go
@@ -90,6 +90,9 @@ const (
 	// TiDBLastDDLInfo is used to get the last ddl info within the current session.
 	TiDBLastDDLInfo = "tidb_last_ddl_info"
 
+	// TiDBLastPlanReplayerToken is used to get the last plan replayer token within the current session
+	TiDBLastPlanReplayerToken = "tidb_last_plan_replayer_token"
+
 	// TiDBConfig is a read-only variable that shows the config of the current server.
 	TiDBConfig = "tidb_config"