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`

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,206 @@ 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 getRangesWithinSpan(ctx context.Context, db rangedesc.DB, span roachpb.Span) (interval.Tree, interval.TreeIterator, error) {
+	rangeDescIterator, err := rangedesc.NewIteratorFactory(db).NewIterator(ctx, span)
+	if err != nil {
+		return nil, nil, err
+	}
+	rangeIntervalTree := interval.NewTree(interval.ExclusiveOverlapper)
+	for rangeDescIterator.Valid() {
+		rangeDesc := rangeDescIterator.CurRangeDescriptor()
+		err := rangeIntervalTree.Insert(rangeDescSpan(rangeDesc), false)
+		if err != nil {
+			return nil, nil, err
+		}
+		rangeDescIterator.Next()
+	}
+	rangeTreeIterator := rangeIntervalTree.Iterator()
+	return rangeIntervalTree, rangeTreeIterator, nil
+}
+
+// rangeSpanIterator is a ValueGenerator that iterates over all
+// ranges of a target span.
+type rangeSpanIterator struct {
+	// The span to iterate
+	span roachpb.Span
+
+	rangeIntervalTree interval.Tree
+	rangeTreeIterator interval.TreeIterator
+
+	// index maintains the current position of the iterator.
+	index int
+	// 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.rangeIntervalTree, rs.rangeTreeIterator, err = getRangesWithinSpan(ctx, txn.DB(), rs.span)
+	if err != nil {
+		return err
+	}
+	// The user of the generator first calls Next(), then Values(), so the index
+	// managing the iterator's position needs to start at -1 instead of 0.
+	rs.index = -1
+	return nil
+}
+
+// Next implements the tree.ValueGenerator interface.
+func (rs *rangeSpanIterator) Next(_ context.Context) (bool, error) {
+	rs.index++
+	// If we don't have any range descriptors, then we're out of results.
+	if rs.rangeIntervalTree.Len() == 0 {
+		return false, nil
+	}
+	// If index exceeds rs.rangeIntervalTree, then we have no more range descriptors to return.
+	if rs.index >= rs.rangeIntervalTree.Len() {
+		return false, nil
+	}
+	return true, nil
+}
+
+// Values implements the tree.ValueGenerator interface.
+func (rs *rangeSpanIterator) Values() (tree.Datums, error) {
+	rangeInterface, exists := rs.rangeTreeIterator.Next()
+	if !exists {
+		return nil, errors.New("expected range from rangeTreeIterator, none found")
+	}
+	rs.buf[0] = tree.NewDInt(tree.DInt(rangeInterface.ID()))
+	rs.buf[1] = tree.NewDBytes(tree.DBytes(rangeInterface.Range().Start))
+	rs.buf[2] = tree.NewDBytes(tree.DBytes(rangeInterface.Range().End))
+	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 {
+	var err error
+	trpti.rangeIntervalTree, _, err = getRangesWithinSpan(ctx, txn.DB(), trpti.span)
+	if err != nil {
+		return err
+	}
+	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()),
+	}
+	overlappingRanges := trpti.rangeIntervalTree.Get(tableRange)
+	for _, oRange := range overlappingRanges {
+		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(overlappingRanges))) // 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
+}