planner: add more test cases for non-prep plan cache (#42715)

ref #36598

planner/core/plan_cache_param.go

@@ -47,22 +47,38 @@ var (
 	}}
 )
 
+// paramReplacer is an ast.Visitor that replaces all values with `?` and collects them.
 type paramReplacer struct {
 	params []*driver.ValueExpr
+
+	// Skip all values in SelectField, e.g.
+	// `select a+1 from t where a<10 and b<23` should be parameterized to
+	// `select a+1 from t where a<? and b<?`, instead of
+	// `select a+? from t where a<? and b<?`.
+	// This is to make the output field names be corresponding to these values.
+	// Use int instead of bool to support nested SelectField.
+	selFieldsCnt int
 }
 
 func (pr *paramReplacer) Enter(in ast.Node) (out ast.Node, skipChildren bool) {
 	switch n := in.(type) {
+	case *ast.SelectField:
+		pr.selFieldsCnt++
 	case *driver.ValueExpr:
-		pr.params = append(pr.params, n)
-		param := ast.NewParamMarkerExpr(len(pr.params) - 1)      // offset is used as order in non-prepared plan cache.
-		param.(*driver.ParamMarkerExpr).Datum = *n.Datum.Clone() // init the ParamMakerExpr's Datum
-		return param, true
+		if pr.selFieldsCnt == 0 { // not in SelectField
+			pr.params = append(pr.params, n)
+			param := ast.NewParamMarkerExpr(len(pr.params) - 1)      // offset is used as order in non-prepared plan cache.
+			param.(*driver.ParamMarkerExpr).Datum = *n.Datum.Clone() // init the ParamMakerExpr's Datum
+			return param, true
+		}
 	}
 	return in, false
 }
 
 func (pr *paramReplacer) Leave(in ast.Node) (out ast.Node, ok bool) {
+	if _, ok := in.(*ast.SelectField); ok {
+		pr.selFieldsCnt--
+	}
 	return in, true
 }
 

planner/core/plan_cache_param_test.go

@@ -56,6 +56,18 @@ func TestParameterize(t *testing.T) {
 			[]interface{}{"a", "bbbbbbbbbbbbbbbbbbbbbbbb"},
 			"SELECT * FROM `t` WHERE `a`=_UTF8MB4'a' AND `b`=_UTF8MB4'bbbbbbbbbbbbbbbbbbbbbbbb'",
 		},
+		{
+			"select 1, 2, 3 from t where a<10",
+			"SELECT 1,2,3 FROM `t` WHERE `a`<?",
+			[]interface{}{int64(10)},
+			"SELECT 1,2,3 FROM `t` WHERE `a`<10",
+		},
+		{
+			"select a+1 from t where a<10",
+			"SELECT `a`+1 FROM `t` WHERE `a`<?",
+			[]interface{}{int64(10)},
+			"SELECT `a`+1 FROM `t` WHERE `a`<10",
+		},
 		// TODO: more test cases
 	}
 

planner/core/plan_cache_test.go

@@ -1440,6 +1440,38 @@ func TestPlanCacheSubquerySPMEffective(t *testing.T) {
 	}
 }
 
+func TestNonPreparedPlanCacheFieldNames(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec(`use test`)
+	tk.MustExec("create table t(a int, index(a))")
+	tk.MustExec("set tidb_enable_non_prepared_plan_cache=1")
+
+	checkFieldName := func(sql, hit string, fields ...string) {
+		rs, err := tk.Exec(sql)
+		require.NoError(t, err)
+		for i, f := range rs.Fields() {
+			require.Equal(t, f.Column.Name.L, fields[i])
+		}
+		require.NoError(t, rs.Close())
+		tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows(hit))
+	}
+
+	checkFieldName(`select a+1 from t where a<10`, `0`, `a+1`)
+	checkFieldName(`select a+1 from t where a<20`, `1`, `a+1`)
+	checkFieldName(`select a+2 from t where a<30`, `0`, `a+2`) // can not hit since field names changed
+	checkFieldName(`select a+2 from t where a<40`, `1`, `a+2`)
+	checkFieldName(`select a,a+1 from t where a<30`, `0`, `a`, `a+1`) // can not hit since field names changed
+	checkFieldName(`select a,a+1 from t where a<40`, `1`, `a`, `a+1`)
+
+	checkFieldName(`select 1 from t where a<10`, `0`, `1`)
+	checkFieldName(`select 1 from t where a<20`, `1`, `1`)
+	checkFieldName(`select 2 from t where a<10`, `0`, `2`)
+	checkFieldName(`select 2 from t where a<20`, `1`, `2`)
+	checkFieldName(`select 1,2 from t where a<10`, `0`, `1`, `2`)
+	checkFieldName(`select 1,2 from t where a<20`, `1`, `1`, `2`)
+}
+
 func TestNonPreparedPlanCacheExplain(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)