From fcc9b0cff2c5a61cb5f2dc2199aec053c7ec371d Mon Sep 17 00:00:00 2001
From: yisaer <disxiaofei@163.com>
Date: Tue, 29 Nov 2022 13:19:49 +0800
Subject: [PATCH] add unique testcase

---
 planner/core/find_best_task.go   | 89 +++++++++++++++++++++-----------
 planner/core/integration_test.go | 12 ++++-
 2 files changed, 70 insertions(+), 31 deletions(-)

diff --git a/planner/core/find_best_task.go b/planner/core/find_best_task.go
index 7046532b9c0dd..896077af2445a 100644
--- a/planner/core/find_best_task.go
+++ b/planner/core/find_best_task.go
@@ -914,8 +914,10 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
 		}
 
 		// if the path is the point get range path with for update lock, we should forbid tiflash as it's store path.
-		if path.StoreType == kv.TiFlash && ds.isPointGetPath(path) && ds.isForUpdateRead {
-			continue
+		if path.StoreType == kv.TiFlash && ds.isForUpdateRead {
+			if ds.isPointGetConditions() {
+				continue
+			}
 		}
 
 		canConvertPointGet := len(path.Ranges) > 0 && path.StoreType == kv.TiKV && ds.isPointGetConvertableSchema()
@@ -1058,33 +1060,6 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
 	return
 }
 
-func (ds *DataSource) isPointGetPath(path *util.AccessPath) bool {
-	if len(path.Ranges) < 1 {
-		return false
-	}
-	for _, ran := range path.Ranges {
-		if !ran.IsPointNonNullable(ds.ctx) {
-			return false
-		}
-	}
-	if path.IsIntHandlePath {
-		return true
-	}
-	if path.Index == nil {
-		return false
-	}
-	if path.Index.HasPrefixIndex() || !path.Index.Unique {
-		return false
-	}
-	idxColsLen := len(path.Index.Columns)
-	for _, ran := range path.Ranges {
-		if len(ran.LowVal) != idxColsLen {
-			return false
-		}
-	}
-	return true
-}
-
 func (ds *DataSource) canConvertToPointGetForPlanCache(path *util.AccessPath) bool {
 	// PointGet might contain some over-optimized assumptions, like `a>=1 and a<=1` --> `a=1`, but
 	// these assumptions may be broken after parameters change.
@@ -1936,6 +1911,62 @@ func (s *LogicalIndexScan) GetPhysicalIndexScan(_ *expression.Schema, stats *pro
 	return is
 }
 
+func (ds *DataSource) isPointGetConditions() bool {
+	t, _ := ds.is.TableByID(ds.physicalTableID)
+	columns := map[string]struct{}{}
+	for _, cond := range ds.allConds {
+		s, ok := cond.(*expression.ScalarFunction)
+		if !ok {
+			return false
+		}
+		if s.FuncName.L != ast.EQ || (s.FuncName.L == ast.In && len(s.GetArgs()) != 2) {
+			return false
+		}
+		arg0 := s.GetArgs()[0]
+		arg1 := s.GetArgs()[1]
+		_, ok1 := arg0.(*expression.Constant)
+		col, ok2 := arg1.(*expression.Column)
+		if ok1 && ok2 {
+			columns[t.Meta().FindColumnNameByID(col.ID)] = struct{}{}
+			continue
+		}
+		col, ok1 = arg0.(*expression.Column)
+		_, ok2 = arg1.(*expression.Constant)
+		if ok1 && ok2 {
+			columns[t.Meta().FindColumnNameByID(col.ID)] = struct{}{}
+			continue
+		}
+	}
+	return ds.findPKOrUniqueIndexMatchColumns(columns)
+}
+
+func (ds *DataSource) findPKOrUniqueIndexMatchColumns(columns map[string]struct{}) bool {
+	for _, idx := range ds.tableInfo.Indices {
+		if !idx.Unique && !idx.Primary {
+			continue
+		}
+		if idx.HasPrefixIndex() {
+			continue
+		}
+		if len(idx.Columns) != len(columns) {
+			continue
+		}
+		flag := true
+		for _, idxCol := range idx.Columns {
+			_, ok := columns[idxCol.Name.String()]
+			if !ok {
+				flag = false
+				break
+			}
+		}
+		if !flag {
+			continue
+		}
+		return true
+	}
+	return false
+}
+
 // convertToTableScan converts the DataSource to table scan.
 func (ds *DataSource) convertToTableScan(prop *property.PhysicalProperty, candidate *candidatePath, _ *physicalOptimizeOp) (task task, err error) {
 	// It will be handled in convertToIndexScan.
diff --git a/planner/core/integration_test.go b/planner/core/integration_test.go
index 7f54f4ed549c8..46b12cd37a1a3 100644
--- a/planner/core/integration_test.go
+++ b/planner/core/integration_test.go
@@ -7558,19 +7558,27 @@ func TestPointGetWithSelectLock(t *testing.T) {
 	store, dom := testkit.CreateMockStoreAndDomain(t)
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")
-	tk.MustExec("create table t(a int, b int, c int, primary key(a, b));")
-	tk.MustExec("insert into t values(1,2,3), (4,5,6);")
+	tk.MustExec("create table t(a int, b int, primary key(a, b));")
+	tk.MustExec("create table t1(c int unique, d int);")
 	tbl, err := dom.InfoSchema().TableByName(model.CIStr{O: "test", L: "test"}, model.CIStr{O: "t", L: "t"})
 	require.NoError(t, err)
 	// Set the hacked TiFlash replica for explain tests.
 	tbl.Meta().TiFlashReplica = &model.TiFlashReplicaInfo{Count: 1, Available: true}
+	tbl1, err := dom.InfoSchema().TableByName(model.CIStr{O: "test", L: "test"}, model.CIStr{O: "t1", L: "t1"})
+	require.NoError(t, err)
+	// Set the hacked TiFlash replica for explain tests.
+	tbl1.Meta().TiFlashReplica = &model.TiFlashReplicaInfo{Count: 1, Available: true}
+
 	tk.MustExec("set @@tidb_enable_tiflash_read_for_write_stmt = on;")
 	tk.MustExec("set @@tidb_isolation_read_engines='tidb,tiflash';")
 	tk.MustExec("begin;")
 	tk.MustGetErrMsg("explain select a, b from t where a = 1 and b = 2 for update;", "[planner:1815]Internal : Can't find a proper physical plan for this query")
+	tk.MustGetErrMsg("explain select c, d from t1 where c = 1 for update;", "[planner:1815]Internal : Can't find a proper physical plan for this query")
 	tk.MustQuery("explain select a, b from t where a = 1 for update;")
+	tk.MustQuery("explain select c, d from t1 where c = 1 and d = 1 for update;")
 	tk.MustExec("set tidb_isolation_read_engines='tidb,tikv,tiflash';")
 	tk.MustQuery("explain select a, b from t where a = 1 and b = 2 for update;")
+	tk.MustQuery("explain select c, d from t1 where c = 1 for update;")
 	tk.MustExec("commit")
 }