apiutil, roachpb: create utilities to map descriptors to ranges

pkg/BUILD.bazel

@@ -323,6 +323,7 @@ ALL_TESTS = [
     "//pkg/security/username:username_disallowed_imports_test",
     "//pkg/security/username:username_test",
     "//pkg/security:security_test",
+    "//pkg/server/apiutil:apiutil_test",
     "//pkg/server/application_api:application_api_test",
     "//pkg/server/authserver:authserver_test",
     "//pkg/server/debug/goroutineui:goroutineui_test",
@@ -1657,6 +1658,7 @@ GO_TARGETS = [
     "//pkg/security:security_test",
     "//pkg/server/apiconstants:apiconstants",
     "//pkg/server/apiutil:apiutil",
+    "//pkg/server/apiutil:apiutil_test",
     "//pkg/server/application_api:application_api",
     "//pkg/server/application_api:application_api_test",
     "//pkg/server/authserver:authserver",

pkg/roachpb/BUILD.bazel

@@ -60,13 +60,15 @@ go_test(
     srcs = [
         "data_test.go",
         "index_usage_stats_test.go",
+        "key_test.go",
         "main_test.go",
         "merge_spans_test.go",
         "metadata_replicas_test.go",
         "metadata_test.go",
         "span_config_conformance_report_test.go",
         "span_config_test.go",
         "span_group_test.go",
+        "span_test.go",
         "string_test.go",
         "tenant_test.go",
         "version_test.go",

pkg/roachpb/data.go

@@ -218,6 +218,23 @@ func (k Key) Compare(b Key) int {
 	return bytes.Compare(k, b)
 }
 
+// Less says whether key k is less than key b.
+func (k Key) Less(b Key) bool {
+	return k.Compare(b) < 0
+}
+
+// Clamp fixes the key to something within the range a < k < b.
+func (k Key) Clamp(a, b Key) Key {
+	result := k
+	if k.Less(a) {
+		result = a
+	}
+	if b.Less(k) {
+		result = b
+	}
+	return result
+}
+
 // SafeFormat implements the redact.SafeFormatter interface.
 func (k Key) SafeFormat(w redact.SafePrinter, _ rune) {
 	SafeFormatKey(w, nil /* valDirs */, k)
@@ -2360,6 +2377,19 @@ func (s Span) Equal(o Span) bool {
 	return s.Key.Equal(o.Key) && s.EndKey.Equal(o.EndKey)
 }
 
+// ZeroLength returns true if the distance between the start and end key is 0.
+func (s Span) ZeroLength() bool {
+	return s.Key.Equal(s.EndKey)
+}
+
+// Clamp clamps span s's keys within the span defined in bounds.
+func (s Span) Clamp(bounds Span) Span {
+	return Span{
+		s.Key.Clamp(bounds.Key, bounds.EndKey),
+		s.EndKey.Clamp(bounds.Key, bounds.EndKey),
+	}
+}
+
 // Overlaps returns true WLOG for span A and B iff:
 //  1. Both spans contain one key (just the start key) and they are equal; or
 //  2. The span with only one key is contained inside the other span; or

pkg/roachpb/key_test.go

@@ -0,0 +1,108 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package roachpb_test
+
+import (
+	"math"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/encoding"
+)
+
+func TestKeyClampTenants(t *testing.T) {
+	// tp = TablePrefix
+	tp := keys.MakeSQLCodec(roachpb.MustMakeTenantID(3)).TablePrefix
+	lowTp := keys.MakeSQLCodec(roachpb.MustMakeTenantID(1)).TablePrefix
+	highTp := keys.MakeSQLCodec(roachpb.MustMakeTenantID(5)).TablePrefix
+	sysTp := keys.SystemSQLCodec.TablePrefix
+	tests := []struct {
+		name     string
+		k, a, b  roachpb.Key
+		expected roachpb.Key
+	}{
+		{"key within main tenant is unchanged", tp(5), tp(1), tp(10), tp(5)},
+		{"low tenant codec gets clamped to lower bound", lowTp(5), tp(1), tp(10), tp(1)},
+		{"high tenant codec gets clamped to upper bound", highTp(5), tp(1), tp(10), tp(10)},
+		{"system codec occurs below the tenant table boundaries", sysTp(5), tp(1), tp(10), tp(1)},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := tt.k.Clamp(tt.a, tt.b)
+			if !result.Equal(tt.expected) {
+				t.Errorf("Clamp(%v, %v, %v) = %v; want %v", tt.k, tt.a, tt.b, result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestKeyClampTables(t *testing.T) {
+	// tp = TablePrefix
+	tp := keys.MakeSQLCodec(roachpb.MustMakeTenantID(3)).TablePrefix
+	tests := []struct {
+		name     string
+		k, a, b  roachpb.Key
+		expected roachpb.Key
+	}{
+		{"table within prefix is unchanged", tp(5), tp(1), tp(10), tp(5)},
+		{"low table gets clamped to lower bound", tp(0), tp(1), tp(10), tp(1)},
+		{"high table gets clamped to upper bound", tp(11), tp(1), tp(10), tp(10)},
+		{"low table on lower bound is unchanged", tp(1), tp(1), tp(10), tp(1)},
+		{"high table on upper bound is unchanged", tp(10), tp(1), tp(10), tp(10)},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := tt.k.Clamp(tt.a, tt.b)
+			if !result.Equal(tt.expected) {
+				t.Errorf("Clamp(%v, %v, %v) = %v; want %v", tt.k, tt.a, tt.b, result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestKeyClampTenantTablespace(t *testing.T) {
+	timeseriesKeyPrefix := encoding.EncodeVarintAscending(
+		encoding.EncodeBytesAscending(
+			append(roachpb.Key(nil), keys.TimeseriesPrefix...),
+			[]byte("my.fake.metric"),
+		),
+		int64(10),
+	)
+	tsKey := func(source string, timestamp int64) roachpb.Key {
+		return append(encoding.EncodeVarintAscending(timeseriesKeyPrefix, timestamp), source...)
+	}
+
+	tp := keys.MakeSQLCodec(roachpb.MustMakeTenantID(3)).TablePrefix
+	lower := tp(0)
+	upper := tp(math.MaxUint32)
+	tests := []struct {
+		name     string
+		k, a, b  roachpb.Key
+		expected roachpb.Key
+	}{
+		{"KeyMin gets clamped to lower", roachpb.KeyMin, lower, upper, lower},
+		{"KeyMax gets clamped to upper", roachpb.KeyMax, lower, upper, upper},
+		{"Meta1Prefix gets clamped to lower", keys.Meta1Prefix, lower, upper, lower},
+		{"Meta2Prefix gets clamped to lower", keys.Meta2Prefix, lower, upper, lower},
+		{"TableDataMin gets clamped to lower", keys.TableDataMin, lower, upper, lower},
+		// below is an unexpected test case for a tenant codec
+		{"TableDataMax also gets clamped to lower", keys.TableDataMax, lower, upper, lower},
+		{"SystemPrefix gets clamped to lower", keys.SystemPrefix, lower, upper, lower},
+		{"TimeseriesKey gets clamped to lower", tsKey("5", 123), lower, upper, lower},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := tt.k.Clamp(tt.a, tt.b)
+			if !result.Equal(tt.expected) {
+				t.Errorf("Clamp(%v, %v, %v) = %v; want %v", tt.k, tt.a, tt.b, result, tt.expected)
+			}
+		})
+	}
+}

pkg/roachpb/span_test.go

@@ -0,0 +1,75 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package roachpb_test
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+)
+
+func TestSpanZeroLength(t *testing.T) {
+	// create two separate references here.
+	shouldBeEmpty := roachpb.Span{
+		Key:    keys.SystemSQLCodec.TablePrefix(1),
+		EndKey: keys.SystemSQLCodec.TablePrefix(1),
+	}
+	if !shouldBeEmpty.ZeroLength() {
+		t.Fatalf("expected span %s to be empty.", shouldBeEmpty)
+	}
+
+	shouldNotBeEmpty := roachpb.Span{
+		Key:    keys.SystemSQLCodec.TablePrefix(1),
+		EndKey: keys.SystemSQLCodec.TablePrefix(1).Next(),
+	}
+	if shouldNotBeEmpty.ZeroLength() {
+		t.Fatalf("expected span %s to not be empty.", shouldNotBeEmpty)
+	}
+}
+
+func TestSpanClamp(t *testing.T) {
+	tp := keys.SystemSQLCodec.TablePrefix
+	tests := []struct {
+		name   string
+		span   roachpb.Span
+		bounds roachpb.Span
+		want   roachpb.Span
+	}{
+		{
+			name:   "within bounds",
+			span:   roachpb.Span{tp(5), tp(10)},
+			bounds: roachpb.Span{tp(0), tp(15)},
+			want:   roachpb.Span{tp(5), tp(10)},
+		},
+		{
+			name:   "clamp lower bound",
+			span:   roachpb.Span{tp(0), tp(10)},
+			bounds: roachpb.Span{tp(5), tp(15)},
+			want:   roachpb.Span{tp(5), tp(10)},
+		},
+		{
+			name:   "clamp upper bound",
+			span:   roachpb.Span{tp(5), tp(20)},
+			bounds: roachpb.Span{tp(0), tp(15)},
+			want:   roachpb.Span{tp(5), tp(15)},
+		},
+		{
+			name:   "clamp both bounds",
+			span:   roachpb.Span{tp(0), tp(20)},
+			bounds: roachpb.Span{tp(5), tp(15)},
+			want:   roachpb.Span{tp(5), tp(15)},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.span.Clamp(tt.bounds); !got.Equal(tt.want) {
+				t.Errorf("Clamp() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

pkg/server/apiutil/BUILD.bazel

@@ -1,9 +1,36 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "apiutil",
-    srcs = ["apiutil.go"],
+    srcs = [
+        "apiutil.go",
+        "rangeutil.go",
+    ],
     importpath = "github.com/cockroachdb/cockroach/pkg/server/apiutil",
     visibility = ["//visibility:public"],
-    deps = ["//pkg/server/srverrors"],
+    deps = [
+        "//pkg/keys",
+        "//pkg/roachpb",
+        "//pkg/server/srverrors",
+        "//pkg/sql/catalog",
+        "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/descs",
+        "@com_github_cockroachdb_errors//:errors",
+    ],
+)
+
+go_test(
+    name = "apiutil_test",
+    srcs = ["rangeutil_test.go"],
+    deps = [
+        ":apiutil",
+        "//pkg/keys",
+        "//pkg/roachpb",
+        "//pkg/sql/catalog",
+        "//pkg/sql/catalog/dbdesc",
+        "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/tabledesc",
+        "//pkg/sql/sem/catid",
+        "@com_github_stretchr_testify//require",
+    ],
 )