planner,executor,infoschema: add system table table_storage_stats

executor/builder.go

@@ -1461,6 +1461,16 @@ func (b *executorBuilder) buildMemTable(v *plannercore.PhysicalMemTable) Executo
 					extractor:  v.Extractor.(*plannercore.SlowQueryExtractor),
 				},
 			}
+		case strings.ToLower(infoschema.TableStorageStats):
+			return &MemTableReaderExec{
+				baseExecutor: newBaseExecutor(b.ctx, v.Schema(), v.ExplainID()),
+				table:        v.Table,
+				retriever: &tableStorageStatsRetriever{
+					table:      v.Table,
+					outputCols: v.Columns,
+					extractor:  v.Extractor.(*plannercore.TableStorageStatsExtractor),
+				},
+			}
 		case strings.ToLower(infoschema.TableDDLJobs):
 			return &DDLJobsReaderExec{
 				baseExecutor: newBaseExecutor(b.ctx, v.Schema(), v.ExplainID()),

executor/infoschema_reader.go

@@ -15,7 +15,10 @@ package executor
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"net/http"
+	"net/url"
 	"sort"
 	"strings"
 	"sync"
@@ -29,16 +32,19 @@ import (
 	"github.com/pingcap/tidb/domain/infosync"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/meta/autoid"
+	plannercore "github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/privilege"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/statistics"
 	"github.com/pingcap/tidb/store/helper"
 	"github.com/pingcap/tidb/store/tikv"
 	"github.com/pingcap/tidb/table"
+	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/codec"
 	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/pdapi"
 	"github.com/pingcap/tidb/util/set"
@@ -1306,6 +1312,177 @@ func (e *memtableRetriever) setDataFromTableConstraints(ctx sessionctx.Context,
 	e.rows = rows
 }
 
+//tableStorageStatsRetriever is used to read slow log data.
+type tableStorageStatsRetriever struct {
+	dummyCloser
+	table         *model.TableInfo
+	outputCols    []*model.ColumnInfo
+	retrieved     bool
+	initialized   bool
+	extractor     *plannercore.TableStorageStatsExtractor
+	initialTables []*initialTable
+	curTable      int
+	pdAddress     string
+}
+
+func (e *tableStorageStatsRetriever) retrieve(ctx context.Context, sctx sessionctx.Context) ([][]types.Datum, error) {
+	if e.retrieved {
+		return nil, nil
+	}
+	if !e.initialized {
+		err := e.initialize(sctx)
+		if err != nil {
+			return nil, err
+		}
+	}
+	if len(e.initialTables) == 0 || e.curTable >= len(e.initialTables) {
+		e.retrieved = true
+		return nil, nil
+	}
+
+	rows, err := e.setDataForTableStorageStats(sctx)
+	if err != nil {
+		return nil, err
+	}
+	if len(e.outputCols) == len(e.table.Columns) {
+		return rows, nil
+	}
+	retRows := make([][]types.Datum, len(rows))
+	for i, fullRow := range rows {
+		row := make([]types.Datum, len(e.outputCols))
+		for j, col := range e.outputCols {
+			row[j] = fullRow[col.Offset]
+		}
+		retRows[i] = row
+	}
+	return retRows, nil
+}
+
+type initialTable struct {
+	db string
+	tb *model.TableInfo
+}
+
+func (e *tableStorageStatsRetriever) initialize(sctx sessionctx.Context) error {
+	is := infoschema.GetInfoSchema(sctx)
+	dbs := is.AllSchemas()
+	sort.Sort(infoschema.SchemasSorter(dbs))
+	schemas := e.extractor.TableSchema
+	tables := e.extractor.TableName
+	if len(schemas) != 0 {
+		for _, schema := range dbs {
+			if schemas.Exist(schema.Name.L) {
+				if len(tables) != 0 {
+					for _, table := range schema.Tables {
+						if tables.Exist(table.Name.L) {
+							e.initialTables = append(e.initialTables, &initialTable{db: schema.Name.L, tb: table})
+						}
+					}
+				} else {
+					for _, table := range schema.Tables {
+						e.initialTables = append(e.initialTables, &initialTable{db: schema.Name.L, tb: table})
+					}
+				}
+			}
+
+		}
+		if len(e.initialTables) == 0 && tables == nil {
+			return errors.Errorf("schema or table not exist, please check the schema and table")
+		}
+	} else {
+		if len(tables) != 0 {
+			for _, schema := range dbs {
+				for _, table := range schema.Tables {
+					if tables.Exist(table.Name.L) {
+						e.initialTables = append(e.initialTables, &initialTable{db: schema.Name.L, tb: table})
+					}
+				}
+			}
+		} else {
+			for _, schema := range dbs {
+				for _, table := range schema.Tables {
+					e.initialTables = append(e.initialTables, &initialTable{schema.Name.L, table})
+				}
+			}
+		}
+	}
+
+	// Cache the PD address.
+	tikvStore, ok := sctx.GetStore().(tikv.Storage)
+	if !ok {
+		return errors.New("Information about TiKV region status can be gotten only when the storage is TiKV")
+	}
+	etcd, ok := tikvStore.(tikv.EtcdBackend)
+	if !ok {
+		return errors.New("not implemented")
+	}
+	pdAddrs := etcd.EtcdAddrs()
+	if len(pdAddrs) == 0 {
+		return errors.New("pd unavailable")
+	}
+	e.pdAddress = pdAddrs[0]
+	e.initialized = true
+	return nil
+}
+
+// PdRegionStats is the json response from PD.
+type PdRegionStats struct {
+	Count          int            `json:"count"`
+	EmptyCount     int            `json:"empty_count"`
+	StorageSize    int64          `json:"storage_size"`
+	StorageKeys    int64          `json:"storage_keys"`
+	StorePeerCount map[uint64]int `json:"store_peer_count"`
+}
+
+func (e *tableStorageStatsRetriever) setDataForTableStorageStats(ctx sessionctx.Context) ([][]types.Datum, error) {
+	rows := make([][]types.Datum, 0, 1024)
+
+	count := 0
+	for i := e.curTable; e.curTable < len(e.initialTables) && count < 1024; i++ {
+		table := (e.initialTables)[e.curTable]
+		tableID := table.tb.ID
+		// Include table and index data, because their range located in tableID_i tableID_r
+		startKey := tablecodec.EncodeTablePrefix(tableID)
+		endKey := tablecodec.EncodeTablePrefix(tableID + 1)
+		startKey = codec.EncodeBytes([]byte{}, startKey)
+		endKey = codec.EncodeBytes([]byte{}, endKey)
+
+		statURL := fmt.Sprintf("http://%s/pd/api/v1/stats/region?start_key=%s&end_key=%s",
+			e.pdAddress,
+			url.QueryEscape(string(startKey)),
+			url.QueryEscape(string(endKey)))
+
+		resp, err := http.Get(statURL)
+		if err != nil {
+			return nil, err
+		}
+		var stats PdRegionStats
+		dec := json.NewDecoder(resp.Body)
+		if err := dec.Decode(&stats); err != nil {
+			return nil, err
+		}
+		peerCount := 0
+		for _, v := range stats.StorePeerCount {
+			peerCount = v
+			break
+		}
+		record := types.MakeDatums(
+			table.db,          // TABLE_SCHEMA
+			table.tb.Name.O,   // TABLE_NAME
+			tableID,           // TABLE_ID
+			peerCount,         // TABLE_PEER_COUNT
+			stats.Count,       // TABLE_REGION_COUNT
+			stats.EmptyCount,  // TABLE_EMPTY_REGION_COUNT
+			stats.StorageSize, // TABLE_SIZE
+			stats.StorageKeys, // TABLE_KEYS
+		)
+		rows = append(rows, record)
+		count++
+		e.curTable++
+	}
+	return rows, nil
+}
+
 func (e *memtableRetriever) setDataFromSessionVar(ctx sessionctx.Context) error {
 	var rows [][]types.Datum
 	var err error

executor/infoschema_reader_test.go

@@ -644,6 +644,16 @@ func (s *testInfoschemaClusterTableSuite) setUpMockPDHTTPServer() (*httptest.Ser
 	router.Handle(pdapi.Config, fn.Wrap(mockConfig))
 	// TiDB/TiKV config
 	router.Handle("/config", fn.Wrap(mockConfig))
+	// pd region
+	router.Handle("/pd/api/v1/stats/region", fn.Wrap(func() (*executor.PdRegionStats, error) {
+		return &executor.PdRegionStats{
+			Count:          1,
+			EmptyCount:     1,
+			StorageSize:    1,
+			StorageKeys:    1,
+			StorePeerCount: map[uint64]int{1: 1},
+		}, nil
+	}))
 	return server, mockAddr
 }
 
@@ -748,6 +758,34 @@ func (s *testInfoschemaClusterTableSuite) TestTiDBClusterInfo(c *C) {
 	))
 }
 
+func (s *testInfoschemaClusterTableSuite) TestTableStorageStats(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	err := tk.QueryToErr("select * from information_schema.TABLE_STORAGE_STATS")
+	c.Assert(err, NotNil)
+	mockAddr := s.mockAddr
+	store := &mockStore{
+		s.store.(tikv.Storage),
+		mockAddr,
+	}
+
+	// information_schema.TABLE_STORAGE_STATS
+	tk = testkit.NewTestKit(c, store)
+	tk.MustQuery("select TABLE_NAME from information_schema.TABLE_STORAGE_STATS where TABLE_SCHEMA = 'information_schema' and TABLE_NAME='schemata';").Check(testkit.Rows("SCHEMATA"))
+	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_NAME, TABLE_SIZE from information_schema.TABLE_STORAGE_STATS where TABLE_SCHEMA = 'test' and TABLE_NAME='t';").Check(testkit.Rows("t 1"))
+
+	tk.MustExec("create table t1 (a int, b int, index idx(a))")
+	tk.MustQuery("select TABLE_NAME, sum(TABLE_SIZE) from information_schema.TABLE_STORAGE_STATS where TABLE_SCHEMA = 'test' group by TABLE_NAME;").Sort().Check(testkit.Rows(
+		"t 1",
+		"t1 1",
+	))
+	tk.MustQuery("select TABLE_SCHEMA, sum(TABLE_SIZE) from information_schema.TABLE_STORAGE_STATS where TABLE_SCHEMA = 'test' group by TABLE_SCHEMA;").Check(testkit.Rows(
+		"test 2",
+	))
+}
+
 func (s *testInfoschemaTableSuite) TestSequences(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 	tk.MustExec("CREATE SEQUENCE test.seq maxvalue 10000000")

infoschema/tables.go

@@ -135,6 +135,8 @@ const (
 	TableInspectionSummary = "INSPECTION_SUMMARY"
 	// TableInspectionRules is the string constant of currently implemented inspection and summary rules.
 	TableInspectionRules = "INSPECTION_RULES"
+	// TableStorageStats is a table that contains all tables disk usage
+	TableStorageStats = "TABLE_STORAGE_STATS"
 	// TableDDLJobs is the string constant of DDL job table.
 	TableDDLJobs = "DDL_JOBS"
 	// TableSequences is the string constant of all sequences created by user.
@@ -200,6 +202,7 @@ var tableIDMap = map[string]int64{
 	TableInspectionRules:                    autoid.InformationSchemaDBID + 56,
 	TableDDLJobs:                            autoid.InformationSchemaDBID + 57,
 	TableSequences:                          autoid.InformationSchemaDBID + 58,
+	TableStorageStats:                       autoid.InformationSchemaDBID + 59,
 }
 
 type columnInfo struct {
@@ -988,6 +991,17 @@ var tableMetricSummaryByLabelCols = []columnInfo{
 	{name: "COMMENT", tp: mysql.TypeVarchar, size: 256},
 }
 
+var tableStorageStatsCols = []columnInfo{
+	{name: "TABLE_SCHEMA", tp: mysql.TypeVarchar, size: 64},
+	{name: "TABLE_NAME", tp: mysql.TypeVarchar, size: 64},
+	{name: "TABLE_ID", tp: mysql.TypeLonglong, size: 21},
+	{name: "PEER_COUNT", tp: mysql.TypeLonglong, size: 21},
+	{name: "REGION_COUNT", tp: mysql.TypeLonglong, size: 21},
+	{name: "EMPTY_REGION_COUNT", tp: mysql.TypeLonglong, size: 21},
+	{name: "TABLE_SIZE", tp: mysql.TypeLonglong, size: 64},
+	{name: "TABLE_KEYS", tp: mysql.TypeLonglong, size: 64},
+}
+
 var tableDDLJobsCols = []columnInfo{
 	{name: "JOB_ID", tp: mysql.TypeLonglong, size: 21},
 	{name: "DB_NAME", tp: mysql.TypeVarchar, size: 64},
@@ -1367,6 +1381,7 @@ var tableNameToColumns = map[string][]columnInfo{
 	TableMetricTables:                       tableMetricTablesCols,
 	TableInspectionSummary:                  tableInspectionSummaryCols,
 	TableInspectionRules:                    tableInspectionRulesCols,
+	TableStorageStats:                       tableStorageStatsCols,
 	TableDDLJobs:                            tableDDLJobsCols,
 	TableSequences:                          tableSequencesCols,
 }

planner/core/logical_plan_builder.go

@@ -2848,6 +2848,8 @@ func (b *PlanBuilder) buildMemTable(_ context.Context, dbName model.CIStr, table
 			p.QueryTimeRange = b.timeRangeForSummaryTable()
 		case infoschema.TableSlowQuery:
 			p.Extractor = &SlowQueryExtractor{}
+		case infoschema.TableStorageStats:
+			p.Extractor = &TableStorageStatsExtractor{}
 		}
 	}
 	return p, nil

planner/core/memtable_predicate_extractor.go

@@ -944,6 +944,44 @@ func (e *SlowQueryExtractor) setTimeRange(start, end int64) {
 	e.Enable = true
 }
 
+// TableStorageStatsExtractor is used to extract some predicates of `disk_usage`
+type TableStorageStatsExtractor struct {
+	extractHelper
+	// SkipRequest means the where clause always false, we don't need to request any component
+	SkipRequest bool
+	//usage       set.StringSet
+	// TableSchema represents tableSchema applied to, and we should apply all table disk usage if there is no schema specified
+	// e.g: SELECT * FROM information_schema.disk_usage WHERE table_schema in ('test', 'information_schema')
+	TableSchema set.StringSet
+	// TableName represents tableName applied to, and we should apply all table disk usage if there is no table specified
+	// e.g: SELECT * FROM information_schema.disk_usage WHERE table in ('schemata', 'tables')
+	TableName set.StringSet
+}
+
+// Extract implements the MemTablePredicateExtractor Extract interface
+func (e *TableStorageStatsExtractor) Extract(
+	_ sessionctx.Context,
+	schema *expression.Schema,
+	names []*types.FieldName,
+	predicates []expression.Expression,
+) []expression.Expression {
+	// Extract the `table_schema` columns
+	remained, schemaSkip, tableSchema := e.extractCol(schema, names, predicates, "table_schema", true)
+	// Extract the `table_name` columns
+	remained, tableSkip, tableName := e.extractCol(schema, names, remained, "table_name", true)
+	e.SkipRequest = schemaSkip || tableSkip
+	if e.SkipRequest {
+		return nil
+	}
+	e.TableSchema = tableSchema
+	e.TableName = tableName
+	return remained
+}
+
+func (e *TableStorageStatsExtractor) explainInfo(p *PhysicalMemTable) string {
+	return ""
+}
+
 func (e *SlowQueryExtractor) explainInfo(p *PhysicalMemTable) string {
 	return ""
 }