*: differentiate types for user variables (#18973)

Co-authored-by: lzmhhh123 <lzmhhh123@gmail.com>

executor/executor_test.go

@@ -135,6 +135,7 @@ var _ = SerialSuites(&tiflashTestSuite{})
 var _ = SerialSuites(&globalIndexSuite{&baseTestSuite{}})
 var _ = SerialSuites(&testSerialSuite{&baseTestSuite{}})
 var _ = SerialSuites(&testCoprCache{})
+var _ = SerialSuites(&testPrepareSuite{})
 
 type testSuite struct{ *baseTestSuite }
 type testSuiteP1 struct{ *baseTestSuite }
@@ -153,6 +154,7 @@ type testCoprCache struct {
 	dom   *domain.Domain
 	cls   cluster.Cluster
 }
+type testPrepareSuite struct{ testData testutil.TestData }
 
 type baseTestSuite struct {
 	cluster cluster.Cluster
@@ -201,6 +203,16 @@ func (s *testSuiteWithData) TearDownSuite(c *C) {
 	c.Assert(s.testData.GenerateOutputIfNeeded(), IsNil)
 }
 
+func (s *testPrepareSuite) SetUpSuite(c *C) {
+	var err error
+	s.testData, err = testutil.LoadTestSuiteData("testdata", "prepare_suite")
+	c.Assert(err, IsNil)
+}
+
+func (s *testPrepareSuite) TearDownSuite(c *C) {
+	c.Assert(s.testData.GenerateOutputIfNeeded(), IsNil)
+}
+
 func (s *baseTestSuite) TearDownSuite(c *C) {
 	s.domain.Close()
 	s.store.Close()

executor/explainfor_test.go

@@ -426,3 +426,64 @@ func (s *testSuite) TestExplainTiFlashSystemTables(c *C) {
 	tk.MustQuery(fmt.Sprintf("desc select * from information_schema.TIFLASH_SEGMENTS where TIFLASH_INSTANCE = '%s' and TIDB_DATABASE = '%s' and TIDB_TABLE = '%s'", tiflashInstance, database, table)).Check(testkit.Rows(
 		fmt.Sprintf("MemTableScan_5 10000.00 root table:TIFLASH_SEGMENTS tiflash_instances:[\"%s\"], tidb_databases:[\"%s\"], tidb_tables:[\"%s\"]", tiflashInstance, database, table)))
 }
+
+func (s *testPrepareSerialSuite) TestPointGetUserVarPlanCache(c *C) {
+	orgEnable := core.PreparedPlanCacheEnabled()
+	defer func() {
+		core.SetPreparedPlanCache(orgEnable)
+	}()
+	core.SetPreparedPlanCache(true)
+
+	tk := testkit.NewTestKitWithInit(c, s.store)
+	tk.Se.Auth(&auth.UserIdentity{Username: "root", Hostname: "localhost", CurrentUser: true, AuthUsername: "root", AuthHostname: "%"}, nil, []byte("012345678901234567890"))
+
+	tk.MustExec("use test")
+	tk.MustExec("set @@tidb_enable_collect_execution_info=0;")
+	tk.MustExec("drop table if exists t1")
+	tk.MustExec("CREATE TABLE t1 (a BIGINT, b VARCHAR(40), PRIMARY KEY (a, b))")
+	tk.MustExec("INSERT INTO t1 VALUES (1,'3'),(2,'4')")
+	tk.MustExec("drop table if exists t2")
+	tk.MustExec("CREATE TABLE t2 (a BIGINT, b VARCHAR(40), UNIQUE KEY idx_a (a))")
+	tk.MustExec("INSERT INTO t2 VALUES (1,'1'),(2,'2')")
+	tk.MustExec("prepare stmt from 'select * from t1, t2 where t1.a = t2.a and t2.a = ?'")
+	tk.MustExec("set @a=1")
+	tk.MustQuery("execute stmt using @a").Check(testkit.Rows(
+		"1 3 1 1",
+	))
+	tkProcess := tk.Se.ShowProcess()
+	ps := []*util.ProcessInfo{tkProcess}
+	tk.Se.SetSessionManager(&mockSessionManager1{PS: ps})
+	// t2 should use PointGet.
+	tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).Check(testkit.Rows(
+		"Projection_7 1.00 root  test.t1.a, test.t1.b, test.t2.a, test.t2.b",
+		"└─IndexMergeJoin_19 1.00 root  inner join, inner:TableReader_14, outer key:test.t2.a, inner key:test.t1.a",
+		"  ├─Selection_41(Build) 0.80 root  not(isnull(test.t2.a))",
+		"  │ └─Point_Get_40 1.00 root table:t2, index:idx_a(a) ",
+		"  └─TableReader_14(Probe) 0.00 root  data:Selection_13",
+		"    └─Selection_13 0.00 cop[tikv]  eq(test.t1.a, 1)",
+		"      └─TableRangeScan_12 1.00 cop[tikv] table:t1 range: decided by [test.t2.a], keep order:true, stats:pseudo",
+	))
+	tk.MustExec("set @a=2")
+	tk.MustQuery("execute stmt using @a").Check(testkit.Rows(
+		"2 4 2 2",
+	))
+	tkProcess = tk.Se.ShowProcess()
+	ps = []*util.ProcessInfo{tkProcess}
+	tk.Se.SetSessionManager(&mockSessionManager1{PS: ps})
+	// t2 should use PointGet, range is changed to [2,2].
+	tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).Check(testkit.Rows(
+		"Projection_7 1.00 root  test.t1.a, test.t1.b, test.t2.a, test.t2.b",
+		"└─IndexMergeJoin_19 1.00 root  inner join, inner:TableReader_14, outer key:test.t2.a, inner key:test.t1.a",
+		"  ├─Selection_41(Build) 0.80 root  not(isnull(test.t2.a))",
+		"  │ └─Point_Get_40 1.00 root table:t2, index:idx_a(a) ",
+		"  └─TableReader_14(Probe) 0.00 root  data:Selection_13",
+		"    └─Selection_13 0.00 cop[tikv]  eq(test.t1.a, 2)",
+		"      └─TableRangeScan_12 1.00 cop[tikv] table:t1 range: decided by [test.t2.a], keep order:true, stats:pseudo",
+	))
+	tk.MustQuery("execute stmt using @a").Check(testkit.Rows(
+		"2 4 2 2",
+	))
+	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows(
+		"1",
+	))
+}

executor/prepared_test.go

@@ -196,16 +196,16 @@ func (s *testSerialSuite) TestPlanCacheClusterIndex(c *C) {
 
 	// For table scan
 	tk.MustExec(`prepare stmt1 from "select * from t1 where t1.a = ? and t1.b > ?"`)
-	tk.MustExec("set @v1 = 1")
-	tk.MustExec("set @v2 = 0")
+	tk.MustExec("set @v1 = '1'")
+	tk.MustExec("set @v2 = '0'")
 	tk.MustQuery("execute stmt1 using @v1,@v2").Check(testkit.Rows("1 1 111"))
 	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0"))
-	tk.MustExec("set @v1 = 2")
-	tk.MustExec("set @v2 = 1")
+	tk.MustExec("set @v1 = '2'")
+	tk.MustExec("set @v2 = '1'")
 	tk.MustQuery("execute stmt1 using @v1,@v2").Check(testkit.Rows("2 2 222"))
 	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1"))
-	tk.MustExec("set @v1 = 3")
-	tk.MustExec("set @v2 = 2")
+	tk.MustExec("set @v1 = '3'")
+	tk.MustExec("set @v2 = '2'")
 	tk.MustQuery("execute stmt1 using @v1,@v2").Check(testkit.Rows("3 3 333"))
 	tkProcess := tk.Se.ShowProcess()
 	ps := []*util.ProcessInfo{tkProcess}
@@ -215,16 +215,16 @@ func (s *testSerialSuite) TestPlanCacheClusterIndex(c *C) {
 
 	// For point get
 	tk.MustExec(`prepare stmt2 from "select * from t1 where t1.a = ? and t1.b = ?"`)
-	tk.MustExec("set @v1 = 1")
-	tk.MustExec("set @v2 = 1")
+	tk.MustExec("set @v1 = '1'")
+	tk.MustExec("set @v2 = '1'")
 	tk.MustQuery("execute stmt2 using @v1,@v2").Check(testkit.Rows("1 1 111"))
 	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0"))
-	tk.MustExec("set @v1 = 2")
-	tk.MustExec("set @v2 = 2")
+	tk.MustExec("set @v1 = '2'")
+	tk.MustExec("set @v2 = '2'")
 	tk.MustQuery("execute stmt2 using @v1,@v2").Check(testkit.Rows("2 2 222"))
 	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1"))
-	tk.MustExec("set @v1 = 3")
-	tk.MustExec("set @v2 = 3")
+	tk.MustExec("set @v1 = '3'")
+	tk.MustExec("set @v2 = '3'")
 	tk.MustQuery("execute stmt2 using @v1,@v2").Check(testkit.Rows("3 3 333"))
 	tkProcess = tk.Se.ShowProcess()
 	ps = []*util.ProcessInfo{tkProcess}
@@ -297,3 +297,91 @@ func (s *testSerialSuite) TestPlanCacheClusterIndex(c *C) {
 	tk.MustQuery(`execute stmt2 using @v1,@v2,@v3,@v4`).Check(testkit.Rows("2 2 222", "3 3 333"))
 	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
 }
+
+func (s *testPrepareSuite) TestPlanCacheWithDifferentVariableTypes(c *C) {
+	store, dom, err := newStoreWithBootstrap()
+	c.Assert(err, IsNil)
+	tk := testkit.NewTestKit(c, store)
+	defer func() {
+		dom.Close()
+		store.Close()
+	}()
+	orgEnable := plannercore.PreparedPlanCacheEnabled()
+	defer func() {
+		plannercore.SetPreparedPlanCache(orgEnable)
+	}()
+	plannercore.SetPreparedPlanCache(true)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t1, t2")
+	tk.MustExec("set @@tidb_enable_collect_execution_info=0;")
+	tk.MustExec("create table t1(a varchar(20), b int, c float, key(b, a))")
+	tk.MustExec("insert into t1 values('1',1,1.1),('2',2,222),('3',3,333)")
+	tk.MustExec("create table t2(a varchar(20), b int, c float, key(b, a))")
+	tk.MustExec("insert into t2 values('3',3,3.3),('2',2,222),('3',3,333)")
+
+	var input []struct {
+		PrepareStmt string
+		Executes    []struct {
+			Vars []struct {
+				Name  string
+				Value string
+			}
+			ExecuteSQL string
+		}
+	}
+	var output []struct {
+		PrepareStmt string
+		Executes    []struct {
+			SQL  string
+			Vars []struct {
+				Name  string
+				Value string
+			}
+			Plan             []string
+			LastPlanUseCache string
+			Result           []string
+		}
+	}
+	s.testData.GetTestCases(c, &input, &output)
+	for i, tt := range input {
+		tk.MustExec(tt.PrepareStmt)
+		s.testData.OnRecord(func() {
+			output[i].PrepareStmt = tt.PrepareStmt
+			output[i].Executes = make([]struct {
+				SQL  string
+				Vars []struct {
+					Name  string
+					Value string
+				}
+				Plan             []string
+				LastPlanUseCache string
+				Result           []string
+			}, len(tt.Executes))
+		})
+		c.Assert(output[i].PrepareStmt, Equals, tt.PrepareStmt)
+		for j, exec := range tt.Executes {
+			for _, v := range exec.Vars {
+				tk.MustExec(fmt.Sprintf(`set @%s = %s`, v.Name, v.Value))
+			}
+			res := tk.MustQuery(exec.ExecuteSQL)
+			lastPlanUseCache := tk.MustQuery("select @@last_plan_from_cache").Rows()[0][0]
+			tk.MustQuery(exec.ExecuteSQL)
+			tkProcess := tk.Se.ShowProcess()
+			ps := []*util.ProcessInfo{tkProcess}
+			tk.Se.SetSessionManager(&mockSessionManager1{PS: ps})
+			plan := tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID))
+			s.testData.OnRecord(func() {
+				output[i].Executes[j].SQL = exec.ExecuteSQL
+				output[i].Executes[j].Plan = s.testData.ConvertRowsToStrings(plan.Rows())
+				output[i].Executes[j].Vars = exec.Vars
+				output[i].Executes[j].LastPlanUseCache = lastPlanUseCache.(string)
+				output[i].Executes[j].Result = s.testData.ConvertRowsToStrings(res.Rows())
+			})
+			c.Assert(output[i].Executes[j].SQL, Equals, exec.ExecuteSQL)
+			plan.Check(testkit.Rows(output[i].Executes[j].Plan...))
+			c.Assert(output[i].Executes[j].Vars, DeepEquals, exec.Vars)
+			c.Assert(output[i].Executes[j].LastPlanUseCache, Equals, lastPlanUseCache.(string))
+			res.Check(testkit.Rows(output[i].Executes[j].Result...))
+		}
+	}
+}

executor/set.go

@@ -33,7 +33,6 @@ import (
 	"github.com/pingcap/tidb/util/gcutil"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/stmtsummary"
-	"github.com/pingcap/tidb/util/stringutil"
 	"go.uber.org/zap"
 )
 
@@ -91,17 +90,15 @@ func (e *SetExecutor) Next(ctx context.Context, req *chunk.Chunk) error {
 			if err != nil {
 				return err
 			}
-
+			sessionVars.UsersLock.Lock()
 			if value.IsNull() {
 				delete(sessionVars.Users, name)
+				delete(sessionVars.UserVarTypes, name)
 			} else {
-				svalue, err1 := value.ToString()
-				if err1 != nil {
-					return err1
-				}
-
-				sessionVars.SetUserVar(name, stringutil.Copy(svalue), value.Collation())
+				sessionVars.Users[name] = value
+				sessionVars.UserVarTypes[name] = v.Expr.GetType()
 			}
+			sessionVars.UsersLock.Unlock()
 			continue
 		}
 

executor/show_test.go

@@ -1066,9 +1066,9 @@ func (s *testSuite5) TestShowBuiltin(c *C) {
 	res := tk.MustQuery("show builtins;")
 	c.Assert(res, NotNil)
 	rows := res.Rows()
-	c.Assert(268, Equals, len(rows))
+	c.Assert(267, Equals, len(rows))
 	c.Assert("abs", Equals, rows[0][0].(string))
-	c.Assert("yearweek", Equals, rows[267][0].(string))
+	c.Assert("yearweek", Equals, rows[266][0].(string))
 }
 
 func (s *testSuite5) TestShowClusterConfig(c *C) {

executor/testdata/prepare_suite_in.json

@@ -0,0 +1,58 @@
+[
+  {
+    "name": "TestPlanCacheWithDifferentVariableTypes",
+    "cases": [
+      {
+        "PrepareStmt": "prepare stmt from \"select ?, ?\"",
+        "Executes": [
+          {
+            "Vars": [{"Name": "v1", "Value": "1"}, {"Name": "v2", "Value": "2"}],
+            "ExecuteSQL": "execute stmt using @v1, @v2"
+          },
+          {
+            "Vars": [{"Name": "v1", "Value": "1"}, {"Name": "v2", "Value": "\"abc\""}],
+            "ExecuteSQL": "execute stmt using @v1, @v2"
+          },
+          {
+            "Vars": [{"Name": "v1", "Value": "10"}, {"Name": "v2", "Value": "\"cba\""}],
+            "ExecuteSQL": "execute stmt using @v1, @v2"
+          }
+        ]
+      },
+      {
+        "PrepareStmt": "prepare stmt from \"select a from t1 where t1.b = ?\"",
+        "Executes": [
+          {
+            "Vars": [{"Name": "v1", "Value": "3"}],
+            "ExecuteSQL": "execute stmt using @v1"
+          },
+          {
+            "Vars": [{"Name": "v1", "Value": "2"}],
+            "ExecuteSQL": "execute stmt using @v1"
+          },
+          {
+            "Vars": [{"Name": "v1", "Value": "\"abc\""}],
+            "ExecuteSQL": "execute stmt using @v1"
+          }
+        ]
+      },
+      {
+        "PrepareStmt": "prepare stmt from \"select t1.c, t2.c from t1 join t2 on t1.b = t2.b and t1.a = t2.a where t1.b = ?\"",
+        "Executes": [
+          {
+            "Vars": [{"Name": "v1", "Value": "1"}],
+            "ExecuteSQL": "execute stmt using @v1"
+          },
+          {
+            "Vars": [{"Name": "v1", "Value": "2"}],
+            "ExecuteSQL": "execute stmt using @v1"
+          },
+          {
+            "Vars": [{"Name": "v1", "Value": "\"abc\""}],
+            "ExecuteSQL": "execute stmt using @v1"
+          }
+        ]
+      }
+    ]
+  }
+]