spanconfigsqltranslator: introduce a pts table reader

pkg/spanconfig/spanconfigsqltranslator/BUILD.bazel

@@ -2,13 +2,17 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "spanconfigsqltranslator",
-    srcs = ["sqltranslator.go"],
+    srcs = [
+        "protectedts_state_reader.go",
+        "sqltranslator.go",
+    ],
     importpath = "github.com/cockroachdb/cockroach/pkg/spanconfig/spanconfigsqltranslator",
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/config/zonepb",
         "//pkg/keys",
         "//pkg/kv",
+        "//pkg/kv/kvserver/protectedts/ptpb",
         "//pkg/roachpb:with-mocks",
         "//pkg/spanconfig",
         "//pkg/sql",
@@ -23,5 +27,21 @@ go_library(
 
 go_test(
     name = "spanconfigsqltranslator_test",
-    srcs = ["sqltranslator_test.go"],
+    srcs = [
+        "protectedts_state_reader_test.go",
+        "sqltranslator_test.go",
+    ],
+    data = glob(["testdata/**"]),
+    embed = [":spanconfigsqltranslator"],
+    deps = [
+        "//pkg/jobs/jobsprotectedts",
+        "//pkg/kv/kvserver/protectedts/ptpb",
+        "//pkg/roachpb:with-mocks",
+        "//pkg/sql/catalog/descpb",
+        "//pkg/util/hlc",
+        "//pkg/util/leaktest",
+        "//pkg/util/log",
+        "//pkg/util/uuid",
+        "@com_github_stretchr_testify//require",
+    ],
 )

pkg/spanconfig/spanconfigsqltranslator/protectedts_state_reader.go

@@ -0,0 +1,98 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package spanconfigsqltranslator
+
+import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptpb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+)
+
+// protectedTimestampStateReader provides a target specific view of the
+// protected timestamp records stored in the system table.
+type protectedTimestampStateReader struct {
+	schemaObjectProtections map[descpb.ID][]hlc.Timestamp
+	tenantProtections       []tenantProtectedTimestamps
+	clusterProtections      []hlc.Timestamp
+}
+
+// newProtectedTimestampStateReader returns an instance of a
+// protectedTimestampStateReader that can be used to fetch target specific
+// protected timestamp records given the supplied ptpb.State. The ptpb.State is
+// the transactional state of the `system.protected_ts_records` table.
+func newProtectedTimestampStateReader(
+	_ context.Context, ptsState ptpb.State,
+) (*protectedTimestampStateReader, error) {
+	reader := &protectedTimestampStateReader{
+		schemaObjectProtections: make(map[descpb.ID][]hlc.Timestamp),
+		tenantProtections:       make([]tenantProtectedTimestamps, 0),
+		clusterProtections:      make([]hlc.Timestamp, 0),
+	}
+	if err := reader.loadProtectedTimestampRecords(ptsState); err != nil {
+		return nil, err
+	}
+	return reader, nil
+}
+
+// GetProtectedTimestampsForCluster returns all the protected timestamps that
+// apply to the entire cluster's keyspace.
+func (p *protectedTimestampStateReader) GetProtectedTimestampsForCluster() []hlc.Timestamp {
+	return p.clusterProtections
+}
+
+// tenantProtectedTimestamps represents all the protections that apply to a
+// tenant's keyspace.
+type tenantProtectedTimestamps struct {
+	protections []hlc.Timestamp
+	tenantID    roachpb.TenantID
+}
+
+// GetProtectedTimestampsForTenants returns all the protected timestamps that
+// apply to a particular tenant's keyspace. It returns this for all tenants that
+// have protected timestamp records.
+func (p *protectedTimestampStateReader) GetProtectedTimestampsForTenants() []tenantProtectedTimestamps {
+	return p.tenantProtections
+}
+
+// GetProtectedTimestampsForSchemaObject returns all the protected timestamps
+// that apply to the descID's keyspan.
+func (p *protectedTimestampStateReader) GetProtectedTimestampsForSchemaObject(
+	descID descpb.ID,
+) []hlc.Timestamp {
+	return p.schemaObjectProtections[descID]
+}
+
+func (p *protectedTimestampStateReader) loadProtectedTimestampRecords(ptsState ptpb.State) error {
+	tenantProtections := make(map[roachpb.TenantID][]hlc.Timestamp)
+	for _, record := range ptsState.Records {
+		switch t := record.Target.GetUnion().(type) {
+		case *ptpb.Target_Cluster:
+			p.clusterProtections = append(p.clusterProtections, record.Timestamp)
+		case *ptpb.Target_Tenants:
+			for _, tenID := range t.Tenants.IDs {
+				tenantProtections[tenID] = append(tenantProtections[tenID], record.Timestamp)
+			}
+		case *ptpb.Target_SchemaObjects:
+			for _, descID := range t.SchemaObjects.IDs {
+				p.schemaObjectProtections[descID] = append(p.schemaObjectProtections[descID], record.Timestamp)
+			}
+		}
+	}
+
+	for tenID, tenantProtections := range tenantProtections {
+		p.tenantProtections = append(p.tenantProtections,
+			tenantProtectedTimestamps{tenantID: tenID, protections: tenantProtections})
+	}
+	return nil
+}

pkg/spanconfig/spanconfigsqltranslator/protectedts_state_reader_test.go

@@ -0,0 +1,100 @@
+// Copyright 2021 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package spanconfigsqltranslator
+
+import (
+	"context"
+	"sort"
+	"testing"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobsprotectedts"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptpb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+func TestProtectedTimestampStateReader(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	mkRecordAndAddToState := func(state *ptpb.State, ts hlc.Timestamp, target *ptpb.Target) {
+		recordID := uuid.MakeV4()
+		rec := jobsprotectedts.MakeRecord(recordID, int64(1), ts, nil, /* deprecatedSpans */
+			jobsprotectedts.Jobs, target)
+		state.Records = append(state.Records, *rec)
+	}
+
+	protectSchemaObject := func(state *ptpb.State, ts hlc.Timestamp, ids []descpb.ID) {
+		mkRecordAndAddToState(state, ts, ptpb.MakeSchemaObjectsTarget(ids))
+	}
+
+	protectCluster := func(state *ptpb.State, ts hlc.Timestamp) {
+		mkRecordAndAddToState(state, ts, ptpb.MakeClusterTarget())
+	}
+
+	protectTenants := func(state *ptpb.State, ts hlc.Timestamp, ids []roachpb.TenantID) {
+		mkRecordAndAddToState(state, ts, ptpb.MakeTenantsTarget(ids))
+	}
+
+	ts := func(seconds int) hlc.Timestamp {
+		return hlc.Timestamp{WallTime: (time.Duration(seconds) * time.Second).Nanoseconds()}
+	}
+
+	// Create some ptpb.State and then run the ProtectedTimestampStateReader on it
+	// to ensure it outputs the expected records.
+	state := &ptpb.State{}
+	protectSchemaObject(state, ts(1), []descpb.ID{56})
+	protectSchemaObject(state, ts(2), []descpb.ID{56, 57})
+	protectCluster(state, ts(3))
+	protectTenants(state, ts(4), []roachpb.TenantID{roachpb.MakeTenantID(1)})
+	protectTenants(state, ts(5), []roachpb.TenantID{roachpb.MakeTenantID(2)})
+	protectTenants(state, ts(6), []roachpb.TenantID{roachpb.MakeTenantID(2)})
+
+	ptsStateReader, err := newProtectedTimestampStateReader(context.Background(), *state)
+	require.NoError(t, err)
+
+	clusterTimestamps := ptsStateReader.GetProtectedTimestampsForCluster()
+	require.Len(t, clusterTimestamps, 1)
+	require.Equal(t, []hlc.Timestamp{ts(3)}, clusterTimestamps)
+
+	tenantTimestamps := ptsStateReader.GetProtectedTimestampsForTenants()
+	sort.Slice(tenantTimestamps, func(i, j int) bool {
+		return tenantTimestamps[i].tenantID.ToUint64() < tenantTimestamps[j].tenantID.ToUint64()
+	})
+	require.Len(t, tenantTimestamps, 2)
+	require.Equal(t, []tenantProtectedTimestamps{
+		{
+			tenantID:    roachpb.MakeTenantID(1),
+			protections: []hlc.Timestamp{ts(4)},
+		},
+		{
+			tenantID:    roachpb.MakeTenantID(2),
+			protections: []hlc.Timestamp{ts(5), ts(6)},
+		},
+	}, tenantTimestamps)
+
+	tableTimestamps := ptsStateReader.GetProtectedTimestampsForSchemaObject(56)
+	sort.Slice(tableTimestamps, func(i, j int) bool {
+		return tableTimestamps[i].Less(tableTimestamps[j])
+	})
+	require.Len(t, tableTimestamps, 2)
+	require.Equal(t, []hlc.Timestamp{ts(1), ts(2)}, tableTimestamps)
+
+	tableTimestamps2 := ptsStateReader.GetProtectedTimestampsForSchemaObject(57)
+	require.Len(t, tableTimestamps2, 1)
+	require.Equal(t, []hlc.Timestamp{ts(2)}, tableTimestamps2)
+}