From fc26bbe61d60f1c59bb6e955656dc5af1d8a02d1 Mon Sep 17 00:00:00 2001
From: lance6716 <lance6716@gmail.com>
Date: Mon, 29 Jul 2024 13:38:13 +0800
Subject: [PATCH] executor: `runWithSystemSession` also copy snapshot status

Signed-off-by: lance6716 <lance6716@gmail.com>
---
 pkg/executor/set.go  | 10 +++++++---
 pkg/executor/show.go |  4 ++++
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/pkg/executor/set.go b/pkg/executor/set.go
index 78ee2e3181430..b5922906c974a 100644
--- a/pkg/executor/set.go
+++ b/pkg/executor/set.go
@@ -316,7 +316,11 @@ func (e *SetExecutor) loadSnapshotInfoSchemaIfNeeded(name string, snapshotTS uin
 	if name != variable.TiDBSnapshot && name != variable.TiDBTxnReadTS {
 		return nil
 	}
-	vars := e.Ctx().GetSessionVars()
+	return loadSnapshotInfoSchemaIfNeeded(e.Ctx(), snapshotTS)
+}
+
+func loadSnapshotInfoSchemaIfNeeded(sctx sessionctx.Context, snapshotTS uint64) error {
+	vars := sctx.GetSessionVars()
 	if snapshotTS == 0 {
 		vars.SnapshotInfoschema = nil
 		return nil
@@ -324,12 +328,12 @@ func (e *SetExecutor) loadSnapshotInfoSchemaIfNeeded(name string, snapshotTS uin
 	logutil.BgLogger().Info("load snapshot info schema",
 		zap.Uint64("conn", vars.ConnectionID),
 		zap.Uint64("SnapshotTS", snapshotTS))
-	dom := domain.GetDomain(e.Ctx())
+	dom := domain.GetDomain(sctx)
 	snapInfo, err := dom.GetSnapshotInfoSchema(snapshotTS)
 	if err != nil {
 		return err
 	}
 
-	vars.SnapshotInfoschema = temptable.AttachLocalTemporaryTableInfoSchema(e.Ctx(), snapInfo)
+	vars.SnapshotInfoschema = temptable.AttachLocalTemporaryTableInfoSchema(sctx, snapInfo)
 	return nil
 }
diff --git a/pkg/executor/show.go b/pkg/executor/show.go
index 5ae40b7e35a31..fdde71a6c00c8 100644
--- a/pkg/executor/show.go
+++ b/pkg/executor/show.go
@@ -2398,6 +2398,10 @@ func runWithSystemSession(ctx context.Context, sctx sessionctx.Context, fn func(
 		return err
 	}
 	defer b.ReleaseSysSession(ctx, sysCtx)
+
+	if err = loadSnapshotInfoSchemaIfNeeded(sysCtx, sctx.GetSessionVars().SnapshotTS); err != nil {
+		return err
+	}
 	// `fn` may use KV transaction, so initialize the txn here
 	if err = sessiontxn.NewTxn(ctx, sysCtx); err != nil {
 		return err