builtins: add tenant_span_stats generator function

Part of: cockroachdb#94332
Part of: cockroachdb#94330

Added new builtin function `tenant_span_stats` that retrieves span
statistics for the current tenant. `tenant_span_stats` can be called as:
- `crdb_internal.tenant_span_stats()`: returns table span statistics for all of the tenants tables
- `crdb_internal.tenant_span_stats(database_id)`: returns table span statistics for the tenant's tables belonging to the specified database id
- `crdb_internal.tenant_span_stats(database_id, table_id)`: returns the tenant's table span statistics for the provided table id

Release note (sql change): new builtin function `tenants_span_stats`,
retrieves the span statistics for the current tenant.

docs/generated/sql/functions.md

@@ -1264,6 +1264,12 @@ the locality flag on node startup. Returns an error if no region is set.</p>
 </span></td><td>Stable</td></tr>
 <tr><td><a name="crdb_internal.scan"></a><code>crdb_internal.scan(start_key: <a href="bytes.html">bytes</a>, end_key: <a href="bytes.html">bytes</a>) &rarr; tuple{bytes AS key, bytes AS value, string AS ts}</code></td><td><span class="funcdesc"><p>Returns the raw keys and values with their timestamp from the specified span</p>
 </span></td><td>Stable</td></tr>
+<tr><td><a name="crdb_internal.tenant_span_stats"></a><code>crdb_internal.tenant_span_stats() &rarr; tuple{int AS database_id, int AS table_id, int AS range_count, int AS approximate_disk_<a href="bytes.html">bytes</a>, int AS live_<a href="bytes.html">bytes</a>, int AS total_<a href="bytes.html">bytes</a>, float AS live_percentage}</code></td><td><span class="funcdesc"><p>Returns statistics (range count, disk size, live range bytes, total range bytes, live range byte percentage) for all of the tenant’s tables.</p>
+</span></td><td>Stable</td></tr>
+<tr><td><a name="crdb_internal.tenant_span_stats"></a><code>crdb_internal.tenant_span_stats(database_id: <a href="int.html">int</a>) &rarr; tuple{int AS database_id, int AS table_id, int AS range_count, int AS approximate_disk_<a href="bytes.html">bytes</a>, int AS live_<a href="bytes.html">bytes</a>, int AS total_<a href="bytes.html">bytes</a>, float AS live_percentage}</code></td><td><span class="funcdesc"><p>Returns statistics (range count, disk size, live range bytes, total range bytes, live range byte percentage) for tables of the provided database id.</p>
+</span></td><td>Stable</td></tr>
+<tr><td><a name="crdb_internal.tenant_span_stats"></a><code>crdb_internal.tenant_span_stats(database_id: <a href="int.html">int</a>, table_id: <a href="int.html">int</a>) &rarr; tuple{int AS database_id, int AS table_id, int AS range_count, int AS approximate_disk_<a href="bytes.html">bytes</a>, int AS live_<a href="bytes.html">bytes</a>, int AS total_<a href="bytes.html">bytes</a>, float AS live_percentage}</code></td><td><span class="funcdesc"><p>Returns statistics (range count, disk size, live range bytes, total range bytes, live range byte percentage) for the provided table id.</p>
+</span></td><td>Stable</td></tr>
 <tr><td><a name="crdb_internal.testing_callback"></a><code>crdb_internal.testing_callback(name: <a href="string.html">string</a>) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>For internal CRDB testing only. The function calls a callback identified by <code>name</code> registered with the server by the test.</p>
 </span></td><td>Volatile</td></tr>
 <tr><td><a name="crdb_internal.unary_table"></a><code>crdb_internal.unary_table() &rarr; tuple</code></td><td><span class="funcdesc"><p>Produces a virtual table containing a single row with no values.</p>

pkg/sql/faketreeeval/evalctx.go

@@ -476,6 +476,20 @@ func (ep *DummyEvalPlanner) GetRangeDescByID(
 	return
 }
 
+// SpanStats is part of the eval.Planner interface.
+func (ep *DummyEvalPlanner) SpanStats(
+	context.Context, roachpb.RKey, roachpb.RKey,
+) (stats *roachpb.SpanStatsResponse, err error) {
+	return
+}
+
+// GetDetailsForSpanStats is part of the eval.Planner interface.
+func (ep *DummyEvalPlanner) GetDetailsForSpanStats(
+	context.Context, int, int,
+) (it eval.InternalRows, err error) {
+	return
+}
+
 // DummyPrivilegedAccessor implements the tree.PrivilegedAccessor interface by returning errors.
 type DummyPrivilegedAccessor struct{}
 

pkg/sql/logictest/testdata/logic_test/tenant_span_stats

@@ -0,0 +1,73 @@
+# LogicTest: 3node-tenant
+
+# Create a second database.
+statement ok
+CREATE DATABASE a
+
+# Create a table for database.
+statement ok
+CREATE TABLE a.b (id INT PRIMARY KEY)
+
+# Create a second table for database.
+statement ok
+CREATE TABLE a.c (id INT PRIMARY KEY)
+
+# SELECT * FROM crdb_internal.tenant_span_stats: span stats for all of the tenant's tables.
+# Assert the schema.
+query IIIIIIR colnames
+SELECT * FROM crdb_internal.tenant_span_stats() LIMIT 0
+----
+database_id table_id range_count approximate_disk_bytes live_bytes total_bytes live_percentage
+
+# SELECT DISTINCT(database_id) FROM crdb_internal.tenant_span_stats:
+# Assert that we are collecting span stats for tables across multiple databases.
+query I colnames
+SELECT DISTINCT(database_id) FROM crdb_internal.tenant_span_stats()
+----
+database_id
+1
+106
+
+# SELECT database_id, table_id FROM crdb_internal.tenant_span_stats(106):
+# Assert that we are collecting span stats scoped to the provided database id.
+query II colnames
+SELECT database_id, table_id FROM crdb_internal.tenant_span_stats(106)
+----
+database_id table_id
+106         108
+106         109
+
+# SELECT database_id, table_id FROM crdb_internal.tenant_span_stats(106, 108):
+# Assert that we are collecting span stats scoped to the provided database/table id combo.
+query II colnames
+SELECT database_id, table_id FROM crdb_internal.tenant_span_stats(106, 108)
+----
+database_id table_id
+106         108
+
+# Assert that we cannot provide an invalid database id.
+query error pq: provided database id must be greater than or equal to 1
+SELECT database_id, table_id FROM crdb_internal.tenant_span_stats(0)
+
+# Assert that we cannot provide an invalid table id.
+query error pq: provided table id must be greater than or equal to 1
+SELECT database_id, table_id FROM crdb_internal.tenant_span_stats(1, -1)
+
+# Assert that we cannot provide an invalid database id with a valid table id.
+query error pq: provided database id must be greater than or equal to 1
+SELECT database_id, table_id FROM crdb_internal.tenant_span_stats(-1, 1)
+
+# SELECT * FROM crdb_internal.tenant_span_stats(1000):
+# Assert that we get empty rows for a database id that does not correspond to a database.
+query IIIIIIR colnames
+SELECT * FROM crdb_internal.tenant_span_stats(1000)
+----
+database_id table_id range_count approximate_disk_bytes live_bytes total_bytes live_percentage
+
+# SELECT * FROM crdb_internal.tenant_span_stats(1, 1000):
+# Assert that we get empty rows for a table id that does not correspond to a table.
+query IIIIIIR colnames
+SELECT * FROM crdb_internal.tenant_span_stats(1, 1000)
+----
+database_id table_id range_count approximate_disk_bytes live_bytes total_bytes live_percentage
+

pkg/sql/planner.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/keyvisualizer"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/repstream"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/spanconfig"
@@ -834,3 +835,45 @@ func (p *planner) GetReplicationStreamManager(
 func (p *planner) GetStreamIngestManager(ctx context.Context) (eval.StreamIngestManager, error) {
 	return repstream.GetStreamIngestManager(ctx, p.EvalContext(), p.InternalSQLTxn())
 }
+
+// SpanStats returns a stats for the given span of keys.
+func (p *planner) SpanStats(
+	ctx context.Context, startKey roachpb.RKey, endKey roachpb.RKey,
+) (*roachpb.SpanStatsResponse, error) {
+	req := &roachpb.SpanStatsRequest{
+		NodeID:   "0",
+		StartKey: startKey,
+		EndKey:   endKey,
+	}
+	return p.ExecCfg().TenantStatusServer.SpanStats(ctx, req)
+}
+
+// GetDetailsForSpanStats ensures that the given database and table id exist.
+// No rows will be returned for database/table ids that do not correspond to an actual
+// database/table.
+func (p *planner) GetDetailsForSpanStats(
+	ctx context.Context, dbId int, tableId int,
+) (eval.InternalRows, error) {
+	query := `SELECT parent_id, table_id from crdb_internal.tables`
+	var args []interface{}
+
+	if tableId != 0 {
+		query += ` where parent_id = $1 and table_id = $2`
+		args = append(args, dbId, tableId)
+	} else if dbId != 0 {
+		query += ` where parent_id = $1`
+		args = append(args, dbId)
+	} else {
+		query += ` where parent_id != $1`
+		args = append(args, dbId)
+	}
+
+	it, err := p.QueryIteratorEx(
+		ctx,
+		"crdb_internal.database_span_stats",
+		sessiondata.NoSessionDataOverride,
+		query,
+		args...,
+	)
+	return it, err
+}

pkg/sql/sem/builtins/fixed_oids.go

@@ -2045,6 +2045,9 @@ var builtinOidsArray = []string{
 	2069: `crdb_internal.create_tenant(parameters: jsonb) -> int`,
 	2070: `crdb_internal.num_inverted_index_entries(val: tsvector, version: int) -> int`,
 	2072: `crdb_internal.upsert_dropped_relation_gc_ttl(desc_id: int, gc_ttl: interval) -> bool`,
+	2073: `crdb_internal.tenant_span_stats() -> tuple{int AS database_id, int AS table_id, int AS range_count, int AS approximate_disk_bytes, int AS live_bytes, int AS total_bytes, float AS live_percentage}`,
+	2074: `crdb_internal.tenant_span_stats(database_id: int) -> tuple{int AS database_id, int AS table_id, int AS range_count, int AS approximate_disk_bytes, int AS live_bytes, int AS total_bytes, float AS live_percentage}`,
+	2075: `crdb_internal.tenant_span_stats(database_id: int, table_id: int) -> tuple{int AS database_id, int AS table_id, int AS range_count, int AS approximate_disk_bytes, int AS live_bytes, int AS total_bytes, float AS live_percentage}`,
 }
 
 var builtinOidsBySignature map[string]oid.Oid

pkg/sql/sem/builtins/generator_builtins.go

@@ -576,6 +576,37 @@ The last argument is a JSONB object containing the following optional fields:
 			volatility.Volatile,
 		),
 	),
+	"crdb_internal.tenant_span_stats": makeBuiltin(genProps(),
+		// Tenant overload
+		makeGeneratorOverload(
+			tree.ParamTypes{},
+			tableSpanStatsGeneratorType,
+			makeTableSpanStatsGenerator,
+			"Returns statistics (range count, disk size, live range bytes, total range bytes, live range byte percentage) for all of the tenant's tables.",
+			volatility.Stable,
+		),
+		// Database overload
+		makeGeneratorOverload(
+			tree.ParamTypes{
+				{Name: "database_id", Typ: types.Int},
+			},
+			tableSpanStatsGeneratorType,
+			makeTableSpanStatsGenerator,
+			"Returns statistics (range count, disk size, live range bytes, total range bytes, live range byte percentage) for tables of the provided database id.",
+			volatility.Stable,
+		),
+		// Table overload
+		makeGeneratorOverload(
+			tree.ParamTypes{
+				{Name: "database_id", Typ: types.Int},
+				{Name: "table_id", Typ: types.Int},
+			},
+			tableSpanStatsGeneratorType,
+			makeTableSpanStatsGenerator,
+			"Returns statistics (range count, disk size, live range bytes, total range bytes, live range byte percentage) for the provided table id.",
+			volatility.Stable,
+		),
+	),
 }
 
 var decodePlanGistGeneratorType = types.String
@@ -2895,3 +2926,113 @@ func makeIdentGenerator(
 		count: count,
 	}, nil
 }
+
+type tableSpanStatsIterator struct {
+	it                eval.InternalRows
+	codec             keys.SQLCodec
+	p                 eval.Planner
+	currDbId          int
+	currTableId       int
+	currStatsResponse *roachpb.SpanStatsResponse
+	singleTableReq    bool
+}
+
+func newTableSpanStatsIterator(eval *eval.Context, dbId int, tableId int) *tableSpanStatsIterator {
+	return &tableSpanStatsIterator{codec: eval.Codec, p: eval.Planner, currDbId: dbId, currTableId: tableId, singleTableReq: tableId != 0}
+}
+
+// Start implements the tree.ValueGenerator interface.
+func (tssi *tableSpanStatsIterator) Start(ctx context.Context, _ *kv.Txn) error {
+	var err error = nil
+	tssi.it, err = tssi.p.GetDetailsForSpanStats(ctx, tssi.currDbId, tssi.currTableId)
+	return err
+}
+
+// Next implements the tree.ValueGenerator interface.
+func (tssi *tableSpanStatsIterator) Next(ctx context.Context) (bool, error) {
+	if tssi.it == nil {
+		return false, errors.AssertionFailedf("Start must be called before Next")
+	}
+	var next bool
+	var err error
+	next, err = tssi.it.Next(ctx)
+	if err != nil || !next {
+		return false, err
+	}
+	// Pull the current row.
+	row := tssi.it.Cur()
+	tssi.currDbId = int(tree.MustBeDInt(row[0]))
+	tssi.currTableId = int(tree.MustBeDInt(row[1]))
+
+	// Set our current stats response.
+	startKey := roachpb.RKey(tssi.codec.TablePrefix(uint32(tssi.currTableId)))
+	tssi.currStatsResponse, err = tssi.p.SpanStats(ctx, startKey, startKey.PrefixEnd())
+	if err != nil {
+		return false, err
+	}
+	return next, err
+}
+
+// Values implements the tree.ValueGenerator interface.
+func (tssi *tableSpanStatsIterator) Values() (tree.Datums, error) {
+	liveBytes := tssi.currStatsResponse.TotalStats.LiveBytes
+	totalBytes := tssi.currStatsResponse.TotalStats.KeyBytes +
+		tssi.currStatsResponse.TotalStats.ValBytes +
+		tssi.currStatsResponse.TotalStats.RangeKeyBytes +
+		tssi.currStatsResponse.TotalStats.RangeValBytes
+	livePercentage := float64(0)
+	if totalBytes > 0 {
+		livePercentage = float64(liveBytes) / float64(totalBytes)
+	}
+	return []tree.Datum{
+		tree.NewDInt(tree.DInt(tssi.currDbId)),
+		tree.NewDInt(tree.DInt(tssi.currTableId)),
+		tree.NewDInt(tree.DInt(tssi.currStatsResponse.RangeCount)),
+		tree.NewDInt(tree.DInt(tssi.currStatsResponse.ApproximateDiskBytes)),
+		tree.NewDInt(tree.DInt(liveBytes)),
+		tree.NewDInt(tree.DInt(totalBytes)),
+		tree.NewDFloat(tree.DFloat(livePercentage)),
+	}, nil
+}
+
+// Close implements the tree.ValueGenerator interface.
+func (tssi *tableSpanStatsIterator) Close(_ context.Context) {}
+
+// ResolvedType implements the tree.ValueGenerator interface.
+func (tssi *tableSpanStatsIterator) ResolvedType() *types.T {
+	return tableSpanStatsGeneratorType
+}
+
+var tableSpanStatsGeneratorType = types.MakeLabeledTuple(
+	[]*types.T{types.Int, types.Int, types.Int, types.Int, types.Int, types.Int, types.Float},
+	[]string{"database_id", "table_id", "range_count", "approximate_disk_bytes", "live_bytes", "total_bytes", "live_percentage"},
+)
+
+func makeTableSpanStatsGenerator(
+	ctx context.Context, evalCtx *eval.Context, args tree.Datums,
+) (eval.ValueGenerator, error) {
+	// The user must be an admin to use this builtin.
+	isAdmin, err := evalCtx.SessionAccessor.HasAdminRole(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if !isAdmin {
+		return nil, pgerror.Newf(pgcode.InsufficientPrivilege, "user needs the admin role to view range data")
+	}
+	dbId := 0
+	tableId := 0
+	if len(args) > 0 {
+		dbId = int(tree.MustBeDInt(args[0]))
+		if dbId <= 0 {
+			return nil, errors.New("provided database id must be greater than or equal to 1")
+		}
+	}
+	if len(args) > 1 {
+		tableId = int(tree.MustBeDInt(args[1]))
+		if tableId <= 0 {
+			return nil, errors.New("provided table id must be greater than or equal to 1")
+		}
+	}
+
+	return newTableSpanStatsIterator(evalCtx, dbId, tableId), nil
+}

pkg/sql/sem/eval/deps.go

@@ -371,6 +371,10 @@ type Planner interface {
 
 	// GetRangeDescByID gets the RangeDescriptor by the specified RangeID.
 	GetRangeDescByID(context.Context, roachpb.RangeID) (roachpb.RangeDescriptor, error)
+
+	SpanStats(context.Context, roachpb.RKey, roachpb.RKey) (*roachpb.SpanStatsResponse, error)
+
+	GetDetailsForSpanStats(ctx context.Context, dbId int, tableId int) (InternalRows, error)
 }
 
 // InternalRows is an iterator interface that's exposed by the internal