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))
@@ -114,32 +124,43 @@ func populatePartitionIDAndNames(tableName *ast.TableName, partitionNames []mode
  if err != nil {
  return 0, nil, err
  }
- pidNames[pid] = partitionName.L
+ pidNames[pid] = genFullPartitionName(table, partitionName.L)
  }
 
  return tbl.Meta().ID, pidNames, nil
 }
 
 // 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: "test.t1 partition (p1)",
+ 3: "test.t1 partition (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/lock_stats_handler.go

@@ -16,17 +16,16 @@ package handle
 
 import (
  "github.com/pingcap/errors"
- "github.com/pingcap/tidb/parser/ast"
  "github.com/pingcap/tidb/statistics/handle/lockstats"
  "github.com/pingcap/tidb/util/sqlexec"
 )
 
 // LockTables add locked tables id to store.
-// - tids: table ids of which will be locked.
-// - pids: partition ids of which will be locked.
+// - tidAndNames: table ids and names of which will be locked.
+// - pidAndNames: partition ids and names of which will be locked.
 // - tables: table names of which will be locked.
 // Return the message of skipped tables and error.
-func (h *Handle) LockTables(tids []int64, pids []int64, tables []*ast.TableName) (string, error) {
+func (h *Handle) LockTables(tidAndNames map[int64]string, pidAndNames map[int64]string) (string, error) {
  se, err := h.pool.Get()
  if err != nil {
  return "", errors.Trace(err)
@@ -35,7 +34,7 @@ func (h *Handle) LockTables(tids []int64, pids []int64, tables []*ast.TableName)
 
  exec := se.(sqlexec.RestrictedSQLExecutor)
 
- return lockstats.AddLockedTables(exec, tids, pids, tables)
+ return lockstats.AddLockedTables(exec, tidAndNames, pidAndNames)
 }
 
 // LockPartitions add locked partitions id to store.
@@ -46,7 +45,7 @@ func (h *Handle) LockTables(tids []int64, pids []int64, tables []*ast.TableName)
 // Return the message of skipped tables and error.
 func (h *Handle) LockPartitions(
  tid int64,
- tableName *ast.TableName,
+ tableName string,
  pidNames map[int64]string,
 ) (string, error) {
  se, err := h.pool.Get()
@@ -61,19 +60,19 @@ func (h *Handle) LockPartitions(
 }
 
 // RemoveLockedTables remove tables from table locked records.
-// - tids: table ids of which will be unlocked.
-// - pids: partition ids of which will be unlocked.
+// - tidAndNames: table ids and names of which will be unlocked.
+// - pidAndNames: partition ids and names of which will be unlocked.
 // - tables: table names of which will be unlocked.
 // Return the message of skipped tables and error.
-func (h *Handle) RemoveLockedTables(tids []int64, pids []int64, tables []*ast.TableName) (string, error) {
+func (h *Handle) RemoveLockedTables(tidAndNames map[int64]string, pidAndNames map[int64]string) (string, error) {
  se, err := h.pool.Get()
  if err != nil {
  return "", errors.Trace(err)
  }
  defer h.pool.Put(se)
 
  exec := se.(sqlexec.RestrictedSQLExecutor)
- return lockstats.RemoveLockedTables(exec, tids, pids, tables)
+ return lockstats.RemoveLockedTables(exec, tidAndNames, pidAndNames)
 }
 
 // RemoveLockedPartitions remove partitions from table locked records.
@@ -82,7 +81,7 @@ func (h *Handle) RemoveLockedTables(tids []int64, pids []int64, tables []*ast.Ta
 // - pidNames: partition ids of which will be unlocked.
 func (h *Handle) RemoveLockedPartitions(
  tid int64,
- tableName *ast.TableName,
+ tableName string,
  pidNames map[int64]string,
 ) (string, error) {
  se, err := h.pool.Get()

statistics/handle/lockstats/BUILD.bazel

@@ -11,7 +11,6 @@ go_library(
  visibility = ["//visibility:public"],
  deps = [
  "//kv",
- "//parser/ast",
  "//parser/terror",
  "//statistics/handle/cache",
  "//util/logutil",
@@ -34,8 +33,6 @@ go_test(
  shard_count = 8,
  deps = [
  "//kv",
- "//parser/ast",
- "//parser/model",
  "//parser/mysql",
  "//types",
  "//util/chunk",