planner: add more test cases for plan cache (#56380)

ref #54057

pkg/planner/core/casetest/instanceplancache/BUILD.bazel

@@ -9,6 +9,10 @@ go_test(
         "others_test.go",
     ],
     flaky = True,
-    shard_count = 9,
-    deps = ["//pkg/testkit"],
+    shard_count = 12,
+    deps = [
+        "//pkg/parser/auth",
+        "//pkg/testkit",
+        "@com_github_stretchr_testify//require",
+    ],
 )

pkg/planner/core/casetest/instanceplancache/others_test.go

@@ -19,7 +19,9 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pingcap/tidb/pkg/parser/auth"
 	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/stretchr/testify/require"
 )
 
 func TestInstancePlanCacheVars(t *testing.T) {
@@ -231,3 +233,100 @@ func TestInstancePlanCacheSchemaChange(t *testing.T) {
 	tk.MustExec(`execute st`)
 	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
 }
+
+func TestInstancePlanCachePrivilegeChanges(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec(`use test`)
+	tk.MustExec(`set global tidb_enable_instance_plan_cache=1`)
+
+	tk.MustExec(`create table t (a int, primary key(a))`)
+	tk.MustExec(`CREATE USER 'u1'`)
+	tk.MustExec(`grant select on test.t to 'u1'`)
+
+	u1 := testkit.NewTestKit(t, store)
+	require.NoError(t, u1.Session().Auth(&auth.UserIdentity{Username: "u1", Hostname: "%"}, nil, nil, nil))
+
+	u1.MustExec(`prepare st from 'select a from test.t where a<1'`)
+	u1.MustExec(`execute st`)
+	u1.MustExec(`execute st`)
+	u1.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+
+	tk.MustExec(`revoke select on test.t from 'u1'`)
+	u1.MustExecToErr(`execute st`) // no privilege
+
+	tk.MustExec(`grant select on test.t to 'u1'`)
+	u1.MustExec(`execute st`)
+	u1.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) // hit the cache again
+}
+
+func TestInstancePlanCacheDifferentUsers(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec(`use test`)
+	tk.MustExec(`set global tidb_enable_instance_plan_cache=1`)
+
+	tk.MustExec(`create table t (a int, primary key(a))`)
+	tk.MustExec(`CREATE USER 'u1'`)
+	tk.MustExec(`grant select on test.t to 'u1'`)
+	tk.MustExec(`CREATE USER 'u1'@'localhost'`)
+	tk.MustExec(`grant select on test.t to 'u1'@'localhost'`)
+	tk.MustExec(`CREATE USER 'u2'`)
+	tk.MustExec(`grant select on test.t to 'u2'`)
+
+	u1 := testkit.NewTestKit(t, store)
+	require.NoError(t, u1.Session().Auth(&auth.UserIdentity{Username: "u1", Hostname: "%"}, nil, nil, nil))
+	u1Local := testkit.NewTestKit(t, store)
+	require.NoError(t, u1Local.Session().Auth(&auth.UserIdentity{Username: "u1", Hostname: "localhost"}, nil, nil, nil))
+	u2 := testkit.NewTestKit(t, store)
+	require.NoError(t, u2.Session().Auth(&auth.UserIdentity{Username: "u2", Hostname: "%"}, nil, nil, nil))
+	u1Dup := testkit.NewTestKit(t, store)
+	require.NoError(t, u1Dup.Session().Auth(&auth.UserIdentity{Username: "u1", Hostname: "%"}, nil, nil, nil))
+
+	u1.MustExec(`prepare st from 'select a from test.t where a=1'`)
+	u1.MustExec(`execute st`)
+	u1.MustExec(`execute st`)
+	u1.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+
+	u1Local.MustExec(`prepare st from 'select a from test.t where a=1'`)
+	u1Local.MustExec(`execute st`)
+	u1Local.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) // can't hit, different localhost
+	u1Local.MustExec(`execute st`)
+	u1Local.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+
+	u2.MustExec(`prepare st from 'select a from test.t where a=1'`)
+	u2.MustExec(`execute st`)
+	u2.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) // can't hit, different user
+	u2.MustExec(`execute st`)
+	u2.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+
+	u1Dup.MustExec(`prepare st from 'select a from test.t where a=1'`)
+	u1Dup.MustExec(`execute st`)
+	u1Dup.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) // can hit, same user and localhost
+	u1Dup.MustExec(`execute st`)
+	u1Dup.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+}
+
+func TestInstancePlanCachePartitioning(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec(`set global tidb_enable_instance_plan_cache=1`)
+	tk.MustExec(`set @@tidb_partition_prune_mode='dynamic'`)
+
+	tk.MustExec(`create table t (a int, b varchar(255)) partition by hash(a) partitions 3`)
+	tk.MustExec(`insert into t values (1,"a"),(2,"b"),(3,"c"),(4,"d"),(5,"e"),(6,"f")`)
+	tk.MustExec(`prepare stmt from 'select a,b from t where a = ?;'`)
+	tk.MustExec(`set @a=1`)
+	tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("1 a"))
+	// Same partition works, due to pruning is not affected
+	tk.MustExec(`set @a=4`)
+	tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d"))
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+
+	tk.MustExec(`set @@tidb_partition_prune_mode='static'`)
+	tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d"))
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0"))
+	tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d"))
+	tk.MustQuery(`show warnings`).Check(testkit.Rows("Warning 1105 skip prepared plan-cache: Static partition pruning mode"))
+}

pkg/planner/core/plan_cache_utils.go

@@ -294,11 +294,24 @@ func NewPlanCacheKey(sctx sessionctx.Context, stmt *PlanCacheStmt) (key, binding
 	}
 	_, connCollation := vars.GetCharsetInfo()
 
+	// not allow to share the same plan among different users for safety.
+	var userName, hostName string
+	if sctx.GetSessionVars().User != nil { // might be nil if in test
+		userName = sctx.GetSessionVars().User.AuthUsername
+		hostName = sctx.GetSessionVars().User.AuthHostname
+	}
+
+	// the user might switch the prune mode dynamically
+	pruneMode := sctx.GetSessionVars().PartitionPruneMode.Load()
+
 	hash := make([]byte, 0, len(stmt.StmtText)*2) // TODO: a Pool for this
+	hash = append(hash, hack.Slice(userName)...)
+	hash = append(hash, hack.Slice(hostName)...)
 	hash = append(hash, hack.Slice(stmtDB)...)
 	hash = append(hash, hack.Slice(stmt.StmtText)...)
 	hash = codec.EncodeInt(hash, stmt.SchemaVersion)
 	hash = hashInt64Uint64Map(hash, stmt.RelateVersion)
+	hash = append(hash, pruneMode...)
 	// Only be set in rc or for update read and leave it default otherwise.
 	// In Rc or ForUpdateRead, we should check whether the information schema has been changed when using plan cache.
 	// If it changed, we should rebuild the plan. lastUpdatedSchemaVersion help us to decide whether we should rebuild