From a0f1d7ea2675932ce185b99fe0ec2d5cbfde27b4 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Mon, 14 Aug 2023 17:16:59 +0800
Subject: [PATCH] executor: fix wrong peer count and improper handling of
 partitioned table in `TABLE_STORAGE_STATS` table (#42612) (#42865)

close pingcap/tidb#42611, close pingcap/tidb#42619
---
 executor/infoschema_cluster_table_test.go | 38 +++++++++++++++++++
 executor/infoschema_reader.go             | 46 ++++++++++++++---------
 2 files changed, 67 insertions(+), 17 deletions(-)

diff --git a/executor/infoschema_cluster_table_test.go b/executor/infoschema_cluster_table_test.go
index bdf7ac14235e5..7a4a301098a26 100644
--- a/executor/infoschema_cluster_table_test.go
+++ b/executor/infoschema_cluster_table_test.go
@@ -326,3 +326,41 @@ func TestTableStorageStats(t *testing.T) {
 
 	tk.MustQuery("select count(1) from information_schema.TABLE_STORAGE_STATS where TABLE_SCHEMA = 'mysql'").Check(testkit.Rows(strconv.Itoa(result)))
 }
+
+func TestIssue42619(t *testing.T) {
+	s := createInfosSchemaClusterTableSuite(t)
+	mockAddr := s.mockAddr
+	store := &mockStore{
+		s.store.(helper.Storage),
+		mockAddr,
+	}
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+
+	tk.MustExec("create table t (a int, b int, index idx(a))")
+	tk.MustQuery("SELECT TABLE_SCHEMA, TABLE_NAME, PEER_COUNT, REGION_COUNT, EMPTY_REGION_COUNT, TABLE_SIZE, TABLE_KEYS " +
+		"FROM information_schema.TABLE_STORAGE_STATS " +
+		"WHERE TABLE_SCHEMA = 'test' and TABLE_NAME='t'").Check(
+		testkit.Rows("test t 1 1 1 1 1"))
+
+	tk.MustExec(
+		"CREATE TABLE tp (a int(11) DEFAULT NULL,b int(11) DEFAULT NULL,c int(11) DEFAULT NULL," +
+			"KEY ia(a), KEY ib(b), KEY ic (c))" +
+			"PARTITION BY RANGE (`a`)" +
+			"(PARTITION `p0` VALUES LESS THAN (300)," +
+			"PARTITION `p1` VALUES LESS THAN (600)," +
+			"PARTITION `p2` VALUES LESS THAN (900)," +
+			"PARTITION `p3` VALUES LESS THAN (MAXVALUE))")
+	tk.MustQuery("SELECT TABLE_SCHEMA, TABLE_NAME, PEER_COUNT, REGION_COUNT, EMPTY_REGION_COUNT, TABLE_SIZE, TABLE_KEYS " +
+		"FROM information_schema.TABLE_STORAGE_STATS " +
+		"WHERE TABLE_SCHEMA = 'test'").Check(
+		testkit.Rows(
+			"test t 1 1 1 1 1",
+			"test tp 1 1 1 1 1",
+			"test tp 1 1 1 1 1",
+			"test tp 1 1 1 1 1",
+			"test tp 1 1 1 1 1",
+			"test tp 1 1 1 1 1",
+		))
+}
diff --git a/executor/infoschema_reader.go b/executor/infoschema_reader.go
index d0ca869d499d1..786c3aa41dc4a 100644
--- a/executor/infoschema_reader.go
+++ b/executor/infoschema_reader.go
@@ -2080,25 +2080,37 @@ func (e *tableStorageStatsRetriever) setDataForTableStorageStats(ctx sessionctx.
 	rows := make([][]types.Datum, 0, 1024)
 	count := 0
 	for e.curTable < len(e.initialTables) && count < 1024 {
-		table := e.initialTables[e.curTable]
-		tableID := table.ID
-		err := e.helper.GetPDRegionStats(tableID, &e.stats, false)
-		if err != nil {
-			return nil, err
+		tbl := e.initialTables[e.curTable]
+		tblIDs := make([]int64, 0, 1)
+		tblIDs = append(tblIDs, tbl.ID)
+		if partInfo := tbl.GetPartitionInfo(); partInfo != nil {
+			for _, partDef := range partInfo.Definitions {
+				tblIDs = append(tblIDs, partDef.ID)
+			}
 		}
-		peerCount := len(e.stats.StorePeerCount)
 
-		record := types.MakeDatums(
-			table.db,            // TABLE_SCHEMA
-			table.Name.O,        // TABLE_NAME
-			tableID,             // TABLE_ID
-			peerCount,           // TABLE_PEER_COUNT
-			e.stats.Count,       // TABLE_REGION_COUNT
-			e.stats.EmptyCount,  // TABLE_EMPTY_REGION_COUNT
-			e.stats.StorageSize, // TABLE_SIZE
-			e.stats.StorageKeys, // TABLE_KEYS
-		)
-		rows = append(rows, record)
+		for _, tableID := range tblIDs {
+			err := e.helper.GetPDRegionStats(tableID, &e.stats, false)
+			if err != nil {
+				return nil, err
+			}
+			peerCount := 0
+			for _, cnt := range e.stats.StorePeerCount {
+				peerCount += cnt
+			}
+
+			record := types.MakeDatums(
+				tbl.db,              // TABLE_SCHEMA
+				tbl.Name.O,          // TABLE_NAME
+				tableID,             // TABLE_ID
+				peerCount,           // TABLE_PEER_COUNT
+				e.stats.Count,       // TABLE_REGION_COUNT
+				e.stats.EmptyCount,  // TABLE_EMPTY_REGION_COUNT
+				e.stats.StorageSize, // TABLE_SIZE
+				e.stats.StorageKeys, // TABLE_KEYS
+			)
+			rows = append(rows, record)
+		}
 		count++
 		e.curTable++
 	}