From 4870c852d2311840a3d4e7bbba7de7632d23c848 Mon Sep 17 00:00:00 2001
From: tiancaiamao <tiancaiamao@gmail.com>
Date: Wed, 21 Jul 2021 17:23:35 +0800
Subject: [PATCH 1/2] cherry pick #26380 to release-5.1

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
---
 executor/executor.go             |  5 +++--
 executor/partition_table_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/executor/executor.go b/executor/executor.go
index 8dab312b09f1a..e02cb46ffc42e 100644
--- a/executor/executor.go
+++ b/executor/executor.go
@@ -929,8 +929,9 @@ func (e *SelectLockExec) Next(ctx context.Context, req *chunk.Chunk) error {
 				if len(e.partitionedTable) > 0 {
 					// Replace the table ID with partition ID.
 					// The partition ID is returned as an extra column from the table reader.
-					offset := e.tblID2PIDColumnIndex[id]
-					physicalID = row.GetInt64(offset)
+					if offset, ok := e.tblID2PIDColumnIndex[id]; ok {
+						physicalID = row.GetInt64(offset)
+					}
 				}
 
 				for _, col := range cols {
diff --git a/executor/partition_table_test.go b/executor/partition_table_test.go
index 5e9dcc289e04c..b43c4486c6c51 100644
--- a/executor/partition_table_test.go
+++ b/executor/partition_table_test.go
@@ -3055,3 +3055,35 @@ func (s *partitionTableSuite) TestIssue25528(c *C) {
 	tk.MustExec("INSERT INTO issue25528 (`c1`, `c2`, `c3`, `c4`) VALUES (1, 1, 1, 1) , (3, 3, 3, 3) , (2, 2, 2, 2) , (4, 4, 4, 4);")
 	tk.MustQuery("select * from issue25528 where c1 in (3, 4) order by c2 for update;").Check(testkit.Rows("3 3 3 3", "4 4 4 4"))
 }
+
+func (s *partitionTableSuite) TestIssue26251(c *C) {
+	tk1 := testkit.NewTestKit(c, s.store)
+	tk1.MustExec("use test")
+	tk1.MustExec("create table tp (id int primary key) partition by range (id) (partition p0 values less than (100));")
+	tk1.MustExec("create table tn (id int primary key);")
+	tk1.MustExec("insert into tp values(1),(2);")
+	tk1.MustExec("insert into tn values(1),(2);")
+
+	tk2 := testkit.NewTestKit(c, s.store)
+	tk2.MustExec("use test")
+
+	tk1.MustExec("begin pessimistic")
+	tk1.MustQuery("select * from tp,tn where tp.id=tn.id and tn.id<=1 for update;").Check(testkit.Rows("1 1"))
+
+	ch := make(chan struct{}, 1)
+	tk2.MustExec("begin pessimistic")
+	go func() {
+		// This query should block.
+		tk2.MustQuery("select * from tn where id=1 for update;").Check(testkit.Rows("1"))
+		ch <- struct{}{}
+	}()
+
+	select {
+	case <-time.After(100 * time.Millisecond):
+		// Expected, query blocked, not finish within 100ms.
+		tk1.MustExec("rollback")
+	case <-ch:
+		// Unexpected, test fail.
+		c.Fail()
+	}
+}

From dc94b5701f7374ce1046f1d081949691c0e33021 Mon Sep 17 00:00:00 2001
From: tiancaiamao <tiancaiamao@gmail.com>
Date: Thu, 5 Aug 2021 15:06:25 +0800
Subject: [PATCH 2/2] cherry-pick changes from
 https://github.com/pingcap/tidb/pull/26373

---
 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 cfb059eaf4701..e9a802de67427 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: