statistics: add better API and better logs for lock stats

executor/lockstats/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "lockstats",
@@ -19,3 +19,17 @@ go_library(
         "@com_github_pingcap_errors//:errors",
     ],
 )
+
+go_test(
+    name = "lockstats_test",
+    timeout = "short",
+    srcs = ["lock_stats_executor_test.go"],
+    embed = [":lockstats"],
+    flaky = True,
+    deps = [
+        "//infoschema",
+        "//parser/ast",
+        "//parser/model",
+        "@com_github_stretchr_testify//require",
+    ],
+)

executor/lockstats/lock_stats_executor.go

@@ -16,6 +16,7 @@ package lockstats
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/domain"
@@ -51,12 +52,13 @@ func (e *LockExec) Next(_ context.Context, _ *chunk.Chunk) error {
 	is := do.InfoSchema()
 
 	if e.onlyLockPartitions() {
-		tableName := e.Tables[0]
-		tid, pidNames, err := populatePartitionIDAndNames(tableName, tableName.PartitionNames, is)
+		table := e.Tables[0]
+		tid, pidNames, err := populatePartitionIDAndNames(table, table.PartitionNames, is)
 		if err != nil {
 			return err
 		}
 
+		tableName := fmt.Sprintf("%s.%s", table.Schema.L, table.Name.L)
 		msg, err := h.LockPartitions(tid, tableName, pidNames)
 		if err != nil {
 			return err
@@ -65,12 +67,12 @@ func (e *LockExec) Next(_ context.Context, _ *chunk.Chunk) error {
 			e.Ctx().GetSessionVars().StmtCtx.AppendWarning(errors.New(msg))
 		}
 	} else {
-		tids, pids, err := populateTableAndPartitionIDs(e.Tables, is)
+		tidAndNames, pidAndNames, err := populateTableAndPartitionIDs(e.Tables, is)
 		if err != nil {
 			return err
 		}
 
-		msg, err := h.LockTables(tids, pids, e.Tables)
+		msg, err := h.LockTables(tidAndNames, pidAndNames)
 		if err != nil {
 			return err
 		}
@@ -97,15 +99,23 @@ func (*LockExec) Open(context.Context) error {
 }
 
 // populatePartitionIDAndNames returns the table ID and partition IDs for the given table name and partition names.
-func populatePartitionIDAndNames(tableName *ast.TableName, partitionNames []model.CIStr, is infoschema.InfoSchema) (int64, map[int64]string, error) {
-	tbl, err := is.TableByName(tableName.Schema, tableName.Name)
+func populatePartitionIDAndNames(
+	table *ast.TableName,
+	partitionNames []model.CIStr,
+	is infoschema.InfoSchema,
+) (int64, map[int64]string, error) {
+	if len(partitionNames) == 0 {
+		return 0, nil, errors.New("partition list should not be empty")
+	}
+	tbl, err := is.TableByName(table.Schema, table.Name)
 	if err != nil {
 		return 0, nil, err
 	}
 
 	pi := tbl.Meta().GetPartitionInfo()
 	if pi == nil {
-		return 0, nil, errors.Errorf("table %s is not a partition table", tableName.Name)
+		return 0, nil, errors.Errorf("table %s is not a partition table",
+			fmt.Sprintf("%s.%s", table.Schema.L, table.Name.L))
 	}
 
 	pidNames := make(map[int64]string, len(partitionNames))
@@ -121,25 +131,36 @@ func populatePartitionIDAndNames(tableName *ast.TableName, partitionNames []mode
 }
 
 // populateTableAndPartitionIDs returns table IDs and partition IDs for the given table names.
-func populateTableAndPartitionIDs(tables []*ast.TableName, is infoschema.InfoSchema) ([]int64, []int64, error) {
-	tids := make([]int64, 0, len(tables))
-	pids := make([]int64, 0)
+func populateTableAndPartitionIDs(
+	tables []*ast.TableName,
+	is infoschema.InfoSchema,
+) (map[int64]string, map[int64]string, error) {
+	if len(tables) == 0 {
+		return nil, nil, errors.New("table list should not be empty")
+	}
+
+	tidAndNames := make(map[int64]string, len(tables))
+	pidAndNames := make(map[int64]string, len(tables))
 
 	for _, table := range tables {
 		tbl, err := is.TableByName(table.Schema, table.Name)
 		if err != nil {
 			return nil, nil, err
 		}
-		tids = append(tids, tbl.Meta().ID)
+		tidAndNames[tbl.Meta().ID] = fmt.Sprintf("%s.%s", table.Schema.L, table.Name.L)
 
 		pi := tbl.Meta().GetPartitionInfo()
 		if pi == nil {
 			continue
 		}
 		for _, p := range pi.Definitions {
-			pids = append(pids, p.ID)
+			pidAndNames[p.ID] = genFullPartitionName(table, p.Name.L)
 		}
 	}
 
-	return tids, pids, nil
+	return tidAndNames, pidAndNames, nil
+}
+
+func genFullPartitionName(table *ast.TableName, partitionName string) string {
+	return fmt.Sprintf("%s.%s partition (%s)", table.Schema.L, table.Name.L, partitionName)
 }

executor/lockstats/lock_stats_executor_test.go

@@ -0,0 +1,108 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package lockstats
+
+import (
+	"testing"
+
+	"github.com/pingcap/tidb/infoschema"
+	"github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/parser/model"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPopulatePartitionIDAndNames(t *testing.T) {
+	fakeInfo := infoschema.MockInfoSchema([]*model.TableInfo{
+		tInfo(1, "t1", "p1", "p2"),
+	})
+
+	table := &ast.TableName{
+		Schema: model.NewCIStr("test"),
+		Name:   model.NewCIStr("t1"),
+		PartitionNames: []model.CIStr{
+			model.NewCIStr("p1"),
+			model.NewCIStr("p2"),
+		},
+	}
+
+	gotTID, gotPIDNames, err := populatePartitionIDAndNames(table, table.PartitionNames, fakeInfo)
+	require.NoError(t, err)
+	require.Equal(t, int64(1), gotTID)
+	require.Equal(t, map[int64]string{
+		2: "p1",
+		3: "p2",
+	}, gotPIDNames)
+
+	// Empty partition names.
+	_, _, err = populatePartitionIDAndNames(nil, nil, fakeInfo)
+	require.Error(t, err)
+}
+
+func TestPopulateTableAndPartitionIDs(t *testing.T) {
+	fakeInfo := infoschema.MockInfoSchema([]*model.TableInfo{
+		tInfo(1, "t1", "p1", "p2"),
+		tInfo(4, "t2"),
+	})
+
+	tables := []*ast.TableName{
+		{
+			Schema: model.NewCIStr("test"),
+			Name:   model.NewCIStr("t1"),
+			PartitionNames: []model.CIStr{
+				model.NewCIStr("p1"),
+				model.NewCIStr("p2"),
+			},
+		},
+		{
+			Schema: model.NewCIStr("test"),
+			Name:   model.NewCIStr("t2"),
+		},
+	}
+
+	gotTIDAndNames, gotPIDAndNames, err := populateTableAndPartitionIDs(tables, fakeInfo)
+	require.NoError(t, err)
+	require.Equal(t, map[int64]string{
+		1: "test.t1",
+		4: "test.t2",
+	}, gotTIDAndNames)
+	require.Equal(t, map[int64]string{
+		2: "test.t1 partition (p1)",
+		3: "test.t1 partition (p2)",
+	}, gotPIDAndNames)
+
+	// Empty table list.
+	_, _, err = populateTableAndPartitionIDs(nil, fakeInfo)
+	require.Error(t, err)
+}
+
+func tInfo(id int, tableName string, partitionNames ...string) *model.TableInfo {
+	tbl := &model.TableInfo{
+		ID:   int64(id),
+		Name: model.NewCIStr(tableName),
+	}
+	if len(partitionNames) > 0 {
+		tbl.Partition = &model.PartitionInfo{
+			Enable: true,
+		}
+		for i, partitionName := range partitionNames {
+			tbl.Partition.Definitions = append(tbl.Partition.Definitions, model.PartitionDefinition{
+				ID:   int64(id + 1 + i),
+				Name: model.NewCIStr(partitionName),
+			})
+		}
+	}
+
+	return tbl
+}

executor/lockstats/unlock_stats_executor.go

@@ -16,6 +16,7 @@ package lockstats
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/domain"
@@ -48,11 +49,12 @@ func (e *UnlockExec) Next(context.Context, *chunk.Chunk) error {
 	is := do.InfoSchema()
 
 	if e.onlyUnlockPartitions() {
-		tableName := e.Tables[0]
-		tid, pidNames, err := populatePartitionIDAndNames(tableName, tableName.PartitionNames, is)
+		table := e.Tables[0]
+		tid, pidNames, err := populatePartitionIDAndNames(table, table.PartitionNames, is)
 		if err != nil {
 			return err
 		}
+		tableName := fmt.Sprintf("%s.%s", table.Schema.O, table.Name.O)
 		msg, err := h.RemoveLockedPartitions(tid, tableName, pidNames)
 		if err != nil {
 			return err
@@ -61,11 +63,11 @@ func (e *UnlockExec) Next(context.Context, *chunk.Chunk) error {
 			e.Ctx().GetSessionVars().StmtCtx.AppendWarning(errors.New(msg))
 		}
 	} else {
-		tids, pids, err := populateTableAndPartitionIDs(e.Tables, is)
+		tidAndNames, pidAndNames, err := populateTableAndPartitionIDs(e.Tables, is)
 		if err != nil {
 			return err
 		}
-		msg, err := h.RemoveLockedTables(tids, pids, e.Tables)
+		msg, err := h.RemoveLockedTables(tidAndNames, pidAndNames)
 		if err != nil {
 			return err
 		}

statistics/handle/handletest/lockstats/BUILD.bazel

@@ -9,7 +9,7 @@ go_test(
         "main_test.go",
     ],
     flaky = True,
-    shard_count = 12,
+    shard_count = 14,
     deps = [
         "//config",
         "//domain",

statistics/handle/handletest/lockstats/lock_partition_stats_test.go

@@ -148,7 +148,7 @@ func TestLockAndUnlockPartitionStatsRepeatedly(t *testing.T) {
 	// Lock the partition again and check the warning.
 	tk.MustExec("lock stats t partition p0")
 	tk.MustQuery("show warnings").Check(testkit.Rows(
-		"Warning 1105 skip locking locked table: test.t partition (p0)",
+		"Warning 1105 skip locking locked partition of table test.t: p0",
 	))
 
 	// Unlock the partition.
@@ -160,7 +160,7 @@ func TestLockAndUnlockPartitionStatsRepeatedly(t *testing.T) {
 	// Unlock the partition again and check the warning.
 	tk.MustExec("unlock stats t partition p0")
 	tk.MustQuery("show warnings").Check(testkit.Rows(
-		"Warning 1105 skip unlocking unlocked table: test.t partition (p0)",
+		"Warning 1105 skip unlocking unlocked partition of table test.t: p0",
 	))
 }
 
@@ -260,6 +260,28 @@ func TestUnlockTheWholeTableWouldUnlockLockedPartitionsAndGenerateWarning(t *tes
 	require.Equal(t, 0, num)
 }
 
+func TestSkipLockALotOfPartitions(t *testing.T) {
+	store, _ := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("set @@tidb_analyze_version = 1")
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a int, b varchar(10), index idx_b (b)) partition by range(a) " +
+		"(partition p0 values less than (10), partition p1 values less than (20), " +
+		"partition a values less than (30), " +
+		"partition b values less than (40), " +
+		"partition g values less than (90), " +
+		"partition h values less than (100))")
+
+	tk.MustExec("lock stats t partition p0, p1, a, b, g, h")
+
+	// Skip locking a lot of partitions.
+	tk.MustExec("lock stats t partition p0, p1, a, b, g, h")
+	tk.MustQuery("show warnings").Check(testkit.Rows(
+		"Warning 1105 skip locking locked partitions of table test.t: a, b, g, h, p0, p1",
+	))
+}
+
 func setupTestEnvironmentWithPartitionedTableT(t *testing.T) (kv.Storage, *testkit.TestKit, *model.TableInfo) {
 	store, dom := testkit.CreateMockStoreAndDomain(t)
 	tk := testkit.NewTestKit(t, store)

statistics/handle/handletest/lockstats/lock_table_stats_test.go

@@ -316,6 +316,26 @@ func TestShowStatsLockedTablePrivilege(t *testing.T) {
 	require.Len(t, rows, 1)
 }
 
+func TestSkipLockALotOfTables(t *testing.T) {
+	store, _ := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("set @@tidb_analyze_version = 1")
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t; drop table if exists a; drop table if exists x; drop table if exists y; drop table if exists z")
+	tk.MustExec("create table t(a int, b varchar(10), index idx_b (b)); " +
+		"create table a(a int, b varchar(10), index idx_b (b)); " +
+		"create table x(a int, b varchar(10), index idx_b (b)); " +
+		"create table y(a int, b varchar(10), index idx_b (b)); " +
+		"create table z(a int, b varchar(10), index idx_b (b))")
+	tk.MustExec("lock stats test.t, test.a, test.x, test.y, test.z")
+
+	// Skip locking a lot of tables.
+	tk.MustExec("lock stats test.t, test.a, test.x, test.y, test.z")
+	tk.MustQuery("show warnings").Check(testkit.Rows(
+		"Warning 1105 skip locking locked tables: test.a, test.t, test.x, test.y, test.z",
+	))
+}
+
 func setupTestEnvironmentWithTableT(t *testing.T) (kv.Storage, *testkit.TestKit, *model.TableInfo) {
 	store, dom := testkit.CreateMockStoreAndDomain(t)
 	tk := testkit.NewTestKit(t, store)