sql: new functions crdb_internal.ranges_in_span and `crdb_internal.…

…tenant_ranges_per_table`

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

This PR introduces two new SRFs:
`crdb_internal.ranges_in_span(start_key, end_key)` - returns the set of ranges encompassing this span in the form of `(range_id, start_key, end_key)`
`crdb_internal.tenant_ranges_per_table()` - returns the tenant's tables, each table row contains the set of ranges encompassing the table, each row takes the form of `(database_name, database_id, table_name, table_id, range_count, range_ids)`

The former SRF is a QOL improvement.
The latter SRF is intended to be used as a method to retrieve a large number of tables' range data quickly, particularly for the database details page (where we surface the range count of each table).

The latter SRF is able to the range ids & range count for ~10,000 tables in ~1.5s.

Release note(sql change): introduce two SRFs `crdb_internal.ranges_in_span` and `crdb_internal.tenant_ranges_per_table`

docs/generated/sql/functions.md

@@ -1260,10 +1260,14 @@ the locality flag on node startup. Returns an error if no region is set.</p>
 <tbody>
 <tr><td><a name="aclexplode"></a><code>aclexplode(aclitems: <a href="string.html">string</a>[]) &rarr; tuple{oid AS grantor, oid AS grantee, string AS privilege_type, bool AS is_grantable}</code></td><td><span class="funcdesc"><p>Produces a virtual table containing aclitem stuff (returns no rows as this feature is unsupported in CockroachDB)</p>
 </span></td><td>Stable</td></tr>
+<tr><td><a name="crdb_internal.ranges_in_span"></a><code>crdb_internal.ranges_in_span(start_key: <a href="bytes.html">bytes</a>, end_key: <a href="bytes.html">bytes</a>) &rarr; tuple{int AS range_id, bytes AS start_key, bytes AS end_key}</code></td><td><span class="funcdesc"><p>Returns ranges (id, start key, end key) within the provided span.</p>
+</span></td><td>Stable</td></tr>
 <tr><td><a name="crdb_internal.scan"></a><code>crdb_internal.scan(span: <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 from the specified span</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_ranges_per_table"></a><code>crdb_internal.tenant_ranges_per_table() &rarr; tuple{string AS database_name, int AS database_id, string AS table_name, int AS table_id, int AS range_count, int[] AS range_ids}</code></td><td><span class="funcdesc"><p>Returns range ids for each of the tenant’s tables.</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/logictest/testdata/logic_test/ranges_in_span

@@ -0,0 +1,18 @@
+# LogicTest: local
+
+subtest select_all_ranges_in_span
+
+# SELECT * FROM crdb_internal.ranges_in_span: all ranges within given span.
+# Assert the schema, using 'system.descriptor' table span as parameter.
+query ITT colnames
+SELECT * FROM crdb_internal.ranges_in_span('\x8b', '\x8c') LIMIT 0
+----
+range_id start_key end_key
+
+# SELECT * FROM crdb_internal.ranges_in_span: all ranges within given span.
+# Assert correct values, using 'system.descriptor' table span as parameter.
+query ITT colnames
+SELECT range_id, crdb_internal.pretty_key(start_key, -1) as start_key, crdb_internal.pretty_key(end_key, -1) as end_key FROM crdb_internal.ranges_in_span('\x8b', '\x8c')
+----
+range_id start_key end_key
+7        /Table/3  /Table/4

pkg/sql/logictest/testdata/logic_test/tenant_ranges_per_table

@@ -0,0 +1,18 @@
+# LogicTest: local
+
+subtest select_tenant_ranges_per_table
+
+# SELECT * FROM crdb_internal.tenant_ranges_per_table: all tenant's ranges per table.
+# Assert the schema.
+query TITIIT colnames
+SELECT * FROM crdb_internal.tenant_ranges_in_span() LIMIT 0
+----
+database_name database_id table_name table_id range_count range_ids
+
+# SELECT * FROM crdb_internal.tenant_ranges_per_table: all tenant's ranges per table.
+# Assert correct values.
+query TITIIT colnames
+SELECT * FROM crdb_internal.tenant_ranges_in_span() LIMIT 1
+----
+database_name database_id table_name table_id range_count range_ids
+system        1           descriptor 3        1           {7}

pkg/sql/sem/builtins/BUILD.bazel

@@ -105,6 +105,7 @@ go_library(
         "//pkg/util/fuzzystrmatch",
         "//pkg/util/hlc",
         "//pkg/util/humanizeutil",
+        "//pkg/util/interval",
         "//pkg/util/intsets",
         "//pkg/util/ipaddr",
         "//pkg/util/json",
@@ -114,6 +115,7 @@ go_library(
         "//pkg/util/randident",
         "//pkg/util/randident/randidentcfg",
         "//pkg/util/randutil",
+        "//pkg/util/rangedesc",
         "//pkg/util/ring",
         "//pkg/util/syncutil",
         "//pkg/util/timeofday",

pkg/sql/sem/builtins/fixed_oids.go

@@ -2044,6 +2044,8 @@ var builtinOidsArray = []string{
 	2068: `crdb_internal.gen_rand_ident(name_pattern: string, count: int, parameters: jsonb) -> string`,
 	2069: `crdb_internal.create_tenant(parameters: jsonb) -> int`,
 	2070: `crdb_internal.num_inverted_index_entries(val: tsvector, version: int) -> int`,
+	2071: `crdb_internal.ranges_in_span(start_key: bytes, end_key: bytes) -> tuple{int AS range_id, bytes AS start_key, bytes AS end_key}`,
+	2072: `crdb_internal.tenant_ranges_per_table() -> tuple{string AS database_name, int AS database_id, string AS table_name, int AS table_id, int AS range_count, int[] AS range_ids}`,
 }
 
 var builtinOidsBySignature map[string]oid.Oid

pkg/sql/sem/builtins/generator_builtins.go

@@ -35,11 +35,13 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/arith"
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
+	"github.com/cockroachdb/cockroach/pkg/util/interval"
 	"github.com/cockroachdb/cockroach/pkg/util/json"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
 	"github.com/cockroachdb/cockroach/pkg/util/randident"
 	"github.com/cockroachdb/cockroach/pkg/util/randident/randidentcfg"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+	"github.com/cockroachdb/cockroach/pkg/util/rangedesc"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
@@ -592,6 +594,27 @@ The last argument is a JSONB object containing the following optional fields:
 			volatility.Volatile,
 		),
 	),
+	"crdb_internal.ranges_in_span": makeBuiltin(genProps(),
+		makeGeneratorOverload(
+			tree.ParamTypes{
+				{Name: "start_key", Typ: types.Bytes},
+				{Name: "end_key", Typ: types.Bytes},
+			},
+			rangesInSpanGeneratorType,
+			makeRangesInSpanGenerator,
+			"Returns ranges (id, start key, end key) within the provided span.",
+			volatility.Stable,
+		),
+	),
+	"crdb_internal.tenant_ranges_per_table": makeBuiltin(genProps(),
+		makeGeneratorOverload(
+			tree.ParamTypes{},
+			tenantRangesPerTableGeneratorType,
+			makeTenantRangesPerTableGenerator,
+			"Returns range ids for each of the tenant's tables.",
+			volatility.Stable,
+		),
+	),
 }
 
 var decodePlanGistGeneratorType = types.String
@@ -2902,3 +2925,198 @@ func makeIdentGenerator(
 		count: count,
 	}, nil
 }
+
+type rangeDescSpan roachpb.RangeDescriptor
+
+var _ interval.Interface = rangeDescSpan{}
+
+// ID is part of `interval.Interface`.
+func (rds rangeDescSpan) ID() uintptr { return uintptr(rds.RangeID) }
+
+// Range is part of `interval.Interface`.
+func (rds rangeDescSpan) Range() interval.Range {
+	return interval.Range{Start: []byte(rds.StartKey.AsRawKey()), End: []byte(rds.EndKey.AsRawKey())}
+}
+
+func getRangeIteratorWithinSpan(
+	ctx context.Context, db rangedesc.DB, span roachpb.Span,
+) (rangedesc.Iterator, error) {
+	rangeDescIterator, err := rangedesc.NewIteratorFactory(db).NewIterator(ctx, span)
+	if err != nil {
+		return nil, err
+	}
+	return rangeDescIterator, nil
+}
+
+// rangeSpanIterator is a ValueGenerator that iterates over all
+// ranges of a target span.
+type rangeSpanIterator struct {
+	// The span to iterate
+	span roachpb.Span
+
+	currRangeDesc roachpb.RangeDescriptor
+	rangeDescs    []roachpb.RangeDescriptor
+	rangeIter     rangedesc.Iterator
+
+	// A buffer to avoid allocating an array on every call to Values().
+	buf [3]tree.Datum
+}
+
+func newRangeSpanIterator(_ *eval.Context, span roachpb.Span) *rangeSpanIterator {
+	return &rangeSpanIterator{span: span}
+}
+
+// Start implements the tree.ValueGenerator interface.
+func (rs *rangeSpanIterator) Start(ctx context.Context, txn *kv.Txn) error {
+	var err error
+	rs.rangeIter, err = getRangeIteratorWithinSpan(ctx, txn.DB(), rs.span)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// Next implements the tree.ValueGenerator interface.
+func (rs *rangeSpanIterator) Next(_ context.Context) (bool, error) {
+	exists := rs.rangeIter.Valid()
+	if exists {
+		rs.currRangeDesc = rs.rangeIter.CurRangeDescriptor()
+		rs.rangeIter.Next()
+	}
+	return exists, nil
+}
+
+// Values implements the tree.ValueGenerator interface.
+func (rs *rangeSpanIterator) Values() (tree.Datums, error) {
+	rs.buf[0] = tree.NewDInt(tree.DInt(rs.currRangeDesc.RangeID))
+	rs.buf[1] = tree.NewDBytes(tree.DBytes(rs.currRangeDesc.StartKey))
+	rs.buf[2] = tree.NewDBytes(tree.DBytes(rs.currRangeDesc.EndKey))
+	return rs.buf[:], nil
+}
+
+// Close implements the tree.ValueGenerator interface.
+func (rs *rangeSpanIterator) Close(_ context.Context) {}
+
+// ResolvedType implements the tree.ValueGenerator interface.
+func (rs *rangeSpanIterator) ResolvedType() *types.T {
+	return rangesInSpanGeneratorType
+}
+
+var rangesInSpanGeneratorType = types.MakeLabeledTuple(
+	[]*types.T{types.Int, types.Bytes, types.Bytes},
+	[]string{"range_id", "start_key", "end_key"},
+)
+
+func makeRangesInSpanGenerator(
+	_ context.Context, evalCtx *eval.Context, args tree.Datums,
+) (eval.ValueGenerator, error) {
+	startKey := []byte(tree.MustBeDBytes(args[0]))
+	endKey := []byte(tree.MustBeDBytes(args[1]))
+	return newRangeSpanIterator(evalCtx, roachpb.Span{
+		Key:    startKey,
+		EndKey: endKey,
+	}), nil
+}
+
+type tenantRangesPerTableIterator struct {
+	codec             keys.SQLCodec
+	planner           eval.Planner
+	it                eval.InternalRows
+	span              roachpb.Span
+	rangeIntervalTree interval.Tree
+	// A buffer to avoid allocating an array on every call to Values().
+	buf [6]tree.Datum
+}
+
+// Start implements the tree.ValueGenerator interface.
+func (trpti *tenantRangesPerTableIterator) Start(ctx context.Context, txn *kv.Txn) error {
+	rangeIter, err := getRangeIteratorWithinSpan(ctx, txn.DB(), trpti.span)
+	if err != nil {
+		return err
+	}
+	trpti.rangeIntervalTree = interval.NewTree(interval.ExclusiveOverlapper)
+	for rangeIter.Valid() {
+		rangeDesc := rangeIter.CurRangeDescriptor()
+		err := trpti.rangeIntervalTree.Insert(rangeDescSpan(rangeDesc), false)
+		if err != nil {
+			return err
+		}
+		rangeIter.Next()
+	}
+	const query = `SELECT
+						database_name,
+						parent_id as database_id,
+						name,
+						table_id,
+						name
+					FROM crdb_internal.tables`
+
+	it, err := trpti.planner.QueryIteratorEx(
+		ctx,
+		"crdb_internal.tenant_ranges_per_table",
+		sessiondata.NoSessionDataOverride,
+		query,
+	)
+	if err != nil {
+		return err
+	}
+	trpti.it = it
+	return nil
+}
+
+// Next implements the tree.ValueGenerator interface.
+func (trpti *tenantRangesPerTableIterator) Next(ctx context.Context) (bool, error) {
+	if trpti.it == nil {
+		return false, errors.AssertionFailedf("Start must be called before Next")
+	}
+	return trpti.it.Next(ctx)
+}
+
+// Values implements the tree.ValueGenerator interface.
+func (trpti *tenantRangesPerTableIterator) Values() (tree.Datums, error) {
+	rangeIds := tree.NewDArray(types.Int)
+	row := trpti.it.Cur()
+	tableID := tree.MustBeDInt(row[3])
+
+	tableStartKey := trpti.codec.TablePrefix(uint32(tableID))
+	tableRange := interval.Range{
+		Start: interval.Comparable(tableStartKey),
+		End:   interval.Comparable(tableStartKey.PrefixEnd()),
+	}
+	tableOverlappingRanges := trpti.rangeIntervalTree.Get(tableRange)
+	for _, oRange := range tableOverlappingRanges {
+		err := rangeIds.Append(tree.NewDInt(tree.DInt(oRange.ID())))
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	trpti.buf[0] = row[0]                                               // database_name
+	trpti.buf[1] = row[1]                                               // database_id
+	trpti.buf[2] = row[2]                                               // table_name
+	trpti.buf[3] = row[3]                                               // table_id
+	trpti.buf[4] = tree.NewDInt(tree.DInt(len(tableOverlappingRanges))) // range_count
+	trpti.buf[5] = rangeIds                                             // range_ids
+
+	return trpti.buf[:], nil
+}
+
+// ResolvedType implements the tree.ValueGenerator interface.
+func (trpti *tenantRangesPerTableIterator) ResolvedType() *types.T {
+	return tenantRangesPerTableGeneratorType
+}
+
+// Close implements the tree.ValueGenerator interface.
+func (trpti *tenantRangesPerTableIterator) Close(_ context.Context) {}
+
+var tenantRangesPerTableGeneratorType = types.MakeLabeledTuple(
+	[]*types.T{types.String, types.Int, types.String, types.Int, types.Int, types.IntArray},
+	[]string{"database_name", "database_id", "table_name", "table_id", "range_count", "range_ids"},
+)
+
+func makeTenantRangesPerTableGenerator(
+	_ context.Context, evalCtx *eval.Context, _ tree.Datums,
+) (eval.ValueGenerator, error) {
+	orsi := &tenantRangesPerTableIterator{span: evalCtx.Codec.TenantSpan(), planner: evalCtx.Planner, codec: evalCtx.Codec}
+	return orsi, nil
+}