diff --git a/pkg/sql/logictest/testdata/logic_test/ranges_in_span b/pkg/sql/logictest/testdata/logic_test/ranges_in_span new file mode 100644 index 000000000000..45318b5be585 --- /dev/null +++ b/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 diff --git a/pkg/sql/logictest/testdata/logic_test/tenant_ranges_per_table b/pkg/sql/logictest/testdata/logic_test/tenant_ranges_per_table new file mode 100644 index 000000000000..41cc102009b7 --- /dev/null +++ b/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} diff --git a/pkg/sql/logictest/tests/local/generated_test.go b/pkg/sql/logictest/tests/local/generated_test.go index 8569be3004de..379a1a77fee8 100644 --- a/pkg/sql/logictest/tests/local/generated_test.go +++ b/pkg/sql/logictest/tests/local/generated_test.go @@ -1486,6 +1486,13 @@ func TestLogic_rand_ident( runLogicTest(t, "rand_ident") } +func TestLogic_ranges_in_span( + t *testing.T, +) { + defer leaktest.AfterTest(t)() + runLogicTest(t, "ranges_in_span") +} + func TestLogic_reassign_owned_by( t *testing.T, ) { @@ -2088,6 +2095,13 @@ func TestLogic_tenant_builtins( runLogicTest(t, "tenant_builtins") } +func TestLogic_tenant_ranges_per_table( + t *testing.T, +) { + defer leaktest.AfterTest(t)() + runLogicTest(t, "tenant_ranges_per_table") +} + func TestLogic_time( t *testing.T, ) { diff --git a/pkg/sql/sem/builtins/BUILD.bazel b/pkg/sql/sem/builtins/BUILD.bazel index 6eb2ff4cb876..eaa50bdd3812 100644 --- a/pkg/sql/sem/builtins/BUILD.bazel +++ b/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", diff --git a/pkg/sql/sem/builtins/fixed_oids.go b/pkg/sql/sem/builtins/fixed_oids.go index 9c6035efa79f..9804d847c697 100644 --- a/pkg/sql/sem/builtins/fixed_oids.go +++ b/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 diff --git a/pkg/sql/sem/builtins/generator_builtins.go b/pkg/sql/sem/builtins/generator_builtins.go index 53923a9ad800..e6729e1bf5ab 100644 --- a/pkg/sql/sem/builtins/generator_builtins.go +++ b/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 +}