From 5c561eb9c798d1f192ca99ea4df3f24d2c35a869 Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Thu, 16 Sep 2021 14:08:43 +0800
Subject: [PATCH] planner/core: fix a panic when select for update on join
 partition table with normal table (#26373) (#26563)

---
 planner/core/integration_test.go | 10 ++++++++++
 planner/core/planbuilder.go      |  4 ++++
 2 files changed, 14 insertions(+)

diff --git a/planner/core/integration_test.go b/planner/core/integration_test.go
index 89d798bb301a6..3941d91e1ec83 100644
--- a/planner/core/integration_test.go
+++ b/planner/core/integration_test.go
@@ -3902,6 +3902,16 @@ func (s *testIntegrationSerialSuite) TestMergeContinuousSelections(c *C) {
 	}
 }
 
+func (s *testIntegrationSerialSuite) TestIssue26250(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("create table tp (id int primary key) partition by range (id) (partition p0 values less than (100));")
+	tk.MustExec("create table tn (id int primary key);")
+	tk.MustExec("insert into tp values(1),(2);")
+	tk.MustExec("insert into tn values(1),(2);")
+	tk.MustQuery("select * from tp,tn where tp.id=tn.id and tn.id=1 for update;").Check(testkit.Rows("1 1"))
+}
+
 func (s *testIntegrationSuite) TestIssue26559(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 	tk.MustExec("use test")
diff --git a/planner/core/planbuilder.go b/planner/core/planbuilder.go
index a97ca266d399d..519b7f54c14e1 100644
--- a/planner/core/planbuilder.go
+++ b/planner/core/planbuilder.go
@@ -1187,6 +1187,10 @@ func (b *PlanBuilder) buildSelectLock(src LogicalPlan, lock *ast.SelectLockInfo)
 func addExtraPIDColumnToDataSource(p LogicalPlan, info *extraPIDInfo) error {
 	switch raw := p.(type) {
 	case *DataSource:
+		// Fix issue 26250, do not add extra pid column to normal table.
+		if raw.tableInfo.GetPartitionInfo() == nil {
+			return nil
+		}
 		raw.addExtraPIDColumn(info)
 		return nil
 	default: