From 7147bfaed2bb68f18e5e56df1f42ef7c271b5312 Mon Sep 17 00:00:00 2001
From: Chao Wang <cclcwangchao@hotmail.com>
Date: Wed, 22 Dec 2021 14:20:44 +0800
Subject: [PATCH] *: Get infoschema from `TxnManager` for some places

---
 executor/adapter.go  | 2 +-
 executor/compiler.go | 5 +++--
 session/session.go   | 8 +++++---
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/executor/adapter.go b/executor/adapter.go
index be83c42d99940..37998a3bfc7c4 100644
--- a/executor/adapter.go
+++ b/executor/adapter.go
@@ -315,7 +315,7 @@ func (a *ExecStmt) RebuildPlan(ctx context.Context) (int64, error) {
 		sessiontxn.AssertTxnManagerInfoSchema(a.Ctx, ret.InfoSchema)
 	})
 
-	a.InfoSchema = ret.InfoSchema
+	a.InfoSchema = sessiontxn.GetTxnManager(a.Ctx).GetTxnInfoSchema()
 	a.SnapshotTS = ret.LastSnapshotTS
 	a.IsStaleness = ret.IsStaleness
 	a.ReplicaReadScope = ret.ReadReplicaScope
diff --git a/executor/compiler.go b/executor/compiler.go
index baf49979572c2..5debec73f1590 100644
--- a/executor/compiler.go
+++ b/executor/compiler.go
@@ -73,7 +73,8 @@ func (c *Compiler) Compile(ctx context.Context, stmtNode ast.StmtNode) (*ExecStm
 		sessiontxn.AssertTxnManagerInfoSchema(c.Ctx, ret.InfoSchema)
 	})
 
-	finalPlan, names, err := planner.Optimize(ctx, c.Ctx, stmtNode, ret.InfoSchema)
+	is := sessiontxn.GetTxnManager(c.Ctx).GetTxnInfoSchema()
+	finalPlan, names, err := planner.Optimize(ctx, c.Ctx, stmtNode, is)
 	if err != nil {
 		return nil, err
 	}
@@ -96,7 +97,7 @@ func (c *Compiler) Compile(ctx context.Context, stmtNode ast.StmtNode) (*ExecStm
 		SnapshotTS:       ret.LastSnapshotTS,
 		IsStaleness:      ret.IsStaleness,
 		ReplicaReadScope: ret.ReadReplicaScope,
-		InfoSchema:       ret.InfoSchema,
+		InfoSchema:       is,
 		Plan:             finalPlan,
 		LowerPriority:    lowerPriority,
 		Text:             stmtNode.Text(),
diff --git a/session/session.go b/session/session.go
index cd1ab05da5498..e781a902afdf7 100644
--- a/session/session.go
+++ b/session/session.go
@@ -2130,7 +2130,9 @@ func (s *session) ExecutePreparedStmt(ctx context.Context, stmtID uint32, args [
 	txnCtxProvider := &sessiontxn.SimpleTxnContextProvider{
 		InfoSchema: is,
 	}
-	if err = sessiontxn.GetTxnManager(s).SetContextProvider(txnCtxProvider); err != nil {
+
+	txnManager := sessiontxn.GetTxnManager(s)
+	if err = txnManager.SetContextProvider(txnCtxProvider); err != nil {
 		return nil, err
 	}
 
@@ -2143,9 +2145,9 @@ func (s *session) ExecutePreparedStmt(ctx context.Context, stmtID uint32, args [
 	defer s.txn.onStmtEnd()
 
 	if ok {
-		return s.cachedPlanExec(ctx, is, snapshotTS, stmtID, preparedStmt, args)
+		return s.cachedPlanExec(ctx, txnManager.GetTxnInfoSchema(), snapshotTS, stmtID, preparedStmt, args)
 	}
-	return s.preparedStmtExec(ctx, is, snapshotTS, stmtID, preparedStmt, args)
+	return s.preparedStmtExec(ctx, txnManager.GetTxnInfoSchema(), snapshotTS, stmtID, preparedStmt, args)
 }
 
 func (s *session) DropPreparedStmt(stmtID uint32) error {