From 5fc54c81037708ee13c891d16341dc28c2fd7b9b Mon Sep 17 00:00:00 2001
From: crazycs <crazycs520@gmail.com>
Date: Thu, 26 Nov 2020 14:19:17 +0800
Subject: [PATCH] cherry pick #21242 to release-4.0

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
---
 server/server_test.go | 22 ++++++++++++++++++++++
 server/tidb_test.go   |  4 ++++
 session/session.go    | 10 ++++++++++
 3 files changed, 36 insertions(+)

diff --git a/server/server_test.go b/server/server_test.go
index 98b301c9fdf22..c78f08e238f2c 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -925,6 +925,28 @@ func (cli *testServerClient) runTestConcurrentUpdate(c *C) {
 	})
 }
 
+func (cli *testServerClient) runTestExplainForConn(c *C) {
+	cli.runTestsOnNewDB(c, nil, "explain_for_conn", func(dbt *DBTest) {
+		dbt.mustExec("drop table if exists t")
+		dbt.mustExec("create table t (a int key, b int)")
+		dbt.mustExec("insert t values (1, 1)")
+		rows := dbt.mustQuery("select connection_id();")
+		c.Assert(rows.Next(), IsTrue)
+		var connID int64
+		err := rows.Scan(&connID)
+		c.Assert(err, IsNil)
+		c.Assert(rows.Close(), IsNil)
+		dbt.mustQuery("select * from t where a=1")
+		rows = dbt.mustQuery("explain for connection " + strconv.Itoa(int(connID)))
+		c.Assert(rows.Next(), IsTrue)
+		row := make([]string, 9)
+		err = rows.Scan(&row[0], &row[1], &row[2], &row[3], &row[4], &row[5], &row[6], &row[7], &row[8])
+		c.Assert(err, IsNil)
+		c.Assert(strings.Join(row, ","), Matches, "Point_Get_1,1.00,1,root,table:t,time.*loop.*handle:1.*")
+		c.Assert(rows.Close(), IsNil)
+	})
+}
+
 func (cli *testServerClient) runTestErrorCode(c *C) {
 	cli.runTestsOnNewDB(c, nil, "ErrorCode", func(dbt *DBTest) {
 		dbt.mustExec("create table test (c int PRIMARY KEY);")
diff --git a/server/tidb_test.go b/server/tidb_test.go
index 6277e3c2bd4c2..9b5bdf5d3d0af 100644
--- a/server/tidb_test.go
+++ b/server/tidb_test.go
@@ -146,6 +146,10 @@ func (ts *tidbTestSerialSuite) TestLoadData(c *C) {
 	ts.runTestLoadDataForSlowLog(c, ts.server)
 }
 
+func (ts *tidbTestSerialSuite) TestExplainFor(c *C) {
+	ts.runTestExplainForConn(c)
+}
+
 func (ts *tidbTestSerialSuite) TestStmtCount(c *C) {
 	ts.runTestStmtCount(c)
 }
diff --git a/session/session.go b/session/session.go
index b00a3ddaa21c5..7da705ec98612 100644
--- a/session/session.go
+++ b/session/session.go
@@ -1052,6 +1052,16 @@ func (s *session) SetProcessInfo(sql string, t time.Time, command byte, maxExecu
 		StatsInfo:        plannercore.GetStatsInfo,
 		MaxExecutionTime: maxExecutionTime,
 	}
+	if p == nil {
+		// Store the last valid plan when the current plan is nil.
+		// This is for `explain for connection` statement has the ability to query the last valid plan.
+		oldPi := s.ShowProcess()
+		if oldPi != nil && oldPi.Plan != nil && len(oldPi.PlanExplainRows) > 0 {
+			pi.Plan = oldPi.Plan
+			pi.PlanExplainRows = oldPi.PlanExplainRows
+			pi.RuntimeStatsColl = oldPi.RuntimeStatsColl
+		}
+	}
 	_, pi.Digest = s.sessionVars.StmtCtx.SQLDigest()
 	s.currentPlan = nil
 	if s.sessionVars.User != nil {