Merge pull request #154249 from spilchen/backport25.4-153919

release-25.4: sql/inspect: add protected timestamp support to INSPECT

Author: rafiss

pkg/jobs/jobspb/jobs.proto

@@ -1464,6 +1464,14 @@ message InspectDetails {
 
   // AsOf specifies the timestamp at which the inspect checks should be performed.
   util.hlc.Timestamp as_of = 2 [(gogoproto.nullable) = false];
+
+  // ProtectedTimestampRecord is the id of the protected timestamp record that
+  // prevents garbage collection of data at the AsOf timestamp. This is only
+  // set if a fixed timestamp was used for INSPECT (e.g. AsOf timestamp is non-empty).
+  bytes protected_timestamp_record = 3 [
+    (gogoproto.customname) = "ProtectedTimestampRecord",
+    (gogoproto.customtype) = "github.com/cockroachdb/cockroach/pkg/util/uuid.UUID"
+  ];
 }
 
 message UpdateTableMetadataCacheDetails {}

pkg/jobs/jobsprotectedts/jobs_protected_ts_manager.go

@@ -55,6 +55,9 @@ func setProtectedTSOnJob(details jobspb.Details, u *uuid.UUID) jobspb.Details {
 	case jobspb.SchemaChangeDetails:
 		v.ProtectedTimestampRecord = u
 		return v
+	case jobspb.InspectDetails:
+		v.ProtectedTimestampRecord = u
+		return v
 	default:
 		panic(errors.AssertionFailedf("not supported %T", details))
 	}
@@ -68,6 +71,8 @@ func getProtectedTSOnJob(details jobspb.Details) *uuid.UUID {
 		return v.ProtectedTimestampRecord
 	case jobspb.SchemaChangeDetails:
 		return v.ProtectedTimestampRecord
+	case jobspb.InspectDetails:
+		return v.ProtectedTimestampRecord
 	default:
 		panic("not supported")
 	}

pkg/sql/exec_util.go

@@ -2101,6 +2101,9 @@ type InspectTestingKnobs struct {
 	// OnInspectJobStart is called just before the inspect job begins execution.
 	// If it returns an error, the job fails immediately.
 	OnInspectJobStart func() error
+	// OnInspectAfterProtectedTimestamp is called after the protected timestamp
+	// has been created (if applicable). If it returns an error, the job fails.
+	OnInspectAfterProtectedTimestamp func() error
 	// InspectIssueLogger is an override to the default issue logger.
 	InspectIssueLogger interface{}
 }

pkg/sql/inspect/BUILD.bazel

@@ -18,6 +18,7 @@ go_library(
     deps = [
         "//pkg/jobs",
         "//pkg/jobs/jobspb",
+        "//pkg/kv/kvserver/protectedts/ptpb",
         "//pkg/roachpb",
         "//pkg/security/username",
         "//pkg/settings",
@@ -72,6 +73,7 @@ go_test(
         "//pkg/jobs/jobspb",
         "//pkg/keys",
         "//pkg/kv",
+        "//pkg/kv/kvserver/protectedts",
         "//pkg/roachpb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
@@ -83,6 +85,7 @@ go_test(
         "//pkg/sql/catalog/desctestutils",
         "//pkg/sql/execinfra",
         "//pkg/sql/execinfrapb",
+        "//pkg/sql/isql",
         "//pkg/sql/rowenc",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondata",

pkg/sql/inspect/index_consistency_check.go

@@ -196,7 +196,6 @@ func (c *indexConsistencyCheck) Start(
 	)
 
 	// Wrap the query with AS OF SYSTEM TIME to ensure it uses the specified timestamp
-	// TODO(148573): use a protected timestamp record for this timestamp.
 	queryWithAsOf := fmt.Sprintf("SELECT * FROM (%s) AS OF SYSTEM TIME %s", checkQuery, c.asOf.AsOfSystemTime())
 
 	// Store the query for error reporting

pkg/sql/inspect/inspect_job.go

@@ -10,9 +10,11 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts/ptpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
@@ -21,6 +23,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/errors"
 )
 
 type inspectResumer struct {
@@ -40,11 +43,24 @@ func (c *inspectResumer) Resume(ctx context.Context, execCtx interface{}) error
 		return err
 	}
 
+	details := c.job.Details().(jobspb.InspectDetails)
+	if len(details.Checks) == 0 {
+		return nil
+	}
+
 	pkSpans, err := c.getPrimaryIndexSpans(ctx, execCfg)
 	if err != nil {
 		return err
 	}
 
+	if err := c.maybeProtectTimestamp(ctx, execCfg, details); err != nil {
+		return err
+	}
+
+	if err := c.maybeRunAfterProtectedTimestampHook(execCfg); err != nil {
+		return err
+	}
+
 	// TODO(149460): add a goroutine that will replan the job on topology changes
 	plan, planCtx, err := c.planInspectProcessors(ctx, jobExecCtx, pkSpans)
 	if err != nil {
@@ -55,13 +71,18 @@ func (c *inspectResumer) Resume(ctx context.Context, execCtx interface{}) error
 		return err
 	}
 
+	c.maybeCleanupProtectedTimestamp(ctx, execCfg)
+
 	return c.markJobComplete(ctx)
 }
 
 // OnFailOrCancel implements the Resumer interface
 func (c *inspectResumer) OnFailOrCancel(
 	ctx context.Context, execCtx interface{}, jobErr error,
 ) error {
+	jobExecCtx := execCtx.(sql.JobExecContext)
+	execCfg := jobExecCtx.ExecCfg()
+	c.maybeCleanupProtectedTimestamp(ctx, execCfg)
 	return nil
 }
 
@@ -77,6 +98,13 @@ func (c *inspectResumer) maybeRunOnJobStartHook(execCfg *sql.ExecutorConfig) err
 	return execCfg.InspectTestingKnobs.OnInspectJobStart()
 }
 
+func (c *inspectResumer) maybeRunAfterProtectedTimestampHook(execCfg *sql.ExecutorConfig) error {
+	if execCfg.InspectTestingKnobs == nil || execCfg.InspectTestingKnobs.OnInspectAfterProtectedTimestamp == nil {
+		return nil
+	}
+	return execCfg.InspectTestingKnobs.OnInspectAfterProtectedTimestamp()
+}
+
 // getPrimaryIndexSpans returns the primary index spans for all tables involved in
 // the INSPECT job's checks.
 func (c *inspectResumer) getPrimaryIndexSpans(
@@ -202,6 +230,50 @@ func (c *inspectResumer) markJobComplete(ctx context.Context) error {
 	)
 }
 
+// maybeProtectTimestamp creates a protected timestamp record for the AsOf
+// timestamp to prevent garbage collection during the inspect operation.
+// If no AsOf timestamp is specified, this function does nothing.
+// The protection target includes all tables involved in the inspect checks.
+// Uses the jobsprotectedts.Manager to store the PTS ID in job details.
+func (c *inspectResumer) maybeProtectTimestamp(
+	ctx context.Context, execCfg *sql.ExecutorConfig, details jobspb.InspectDetails,
+) error {
+	// If we are not running a historical query, nothing to do here.
+	if details.AsOf.IsEmpty() {
+		return nil
+	}
+
+	// Create a target for the specific tables involved in the inspect checks
+	var tableIDSet catalog.DescriptorIDSet
+	for _, check := range details.Checks {
+		tableIDSet.Add(check.TableID)
+	}
+	target := ptpb.MakeSchemaObjectsTarget(tableIDSet.Ordered())
+
+	// Protect will store the PTS ID in job details.
+	_, err := execCfg.ProtectedTimestampManager.Protect(ctx, c.job, target, details.AsOf)
+	if err != nil {
+		return errors.Wrapf(err, "failed to protect timestamp %s for INSPECT job %d", details.AsOf, c.job.ID())
+	}
+
+	log.Dev.Infof(ctx, "protected timestamp created for INSPECT job %d at %s", c.job.ID(), details.AsOf)
+	return nil
+}
+
+// maybeCleanupProtectedTimestamp cleans up any protected timestamp record
+// associated with this job. If no protected timestamp was created, this
+// function does nothing.
+func (c *inspectResumer) maybeCleanupProtectedTimestamp(
+	ctx context.Context, execCfg *sql.ExecutorConfig,
+) {
+	details := c.job.Details().(jobspb.InspectDetails)
+	if details.ProtectedTimestampRecord != nil {
+		if err := execCfg.ProtectedTimestampManager.Unprotect(ctx, c.job); err != nil {
+			log.Dev.Warningf(ctx, "failed to clean up protected timestamp: %v", err)
+		}
+	}
+}
+
 func init() {
 	createResumerFn := func(job *jobs.Job, settings *cluster.Settings) jobs.Resumer {
 		return &inspectResumer{job: job}

pkg/sql/inspect/inspect_job_test.go

@@ -13,7 +13,10 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts"
 	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
@@ -145,3 +148,159 @@ func TestInspectJobImplicitTxnSemantics(t *testing.T) {
 		})
 	}
 }
+
+// TestInspectJobProtectedTimestamp verifies that INSPECT jobs properly create
+// and clean up protected timestamp records when using AS OF SYSTEM TIME.
+func TestInspectJobProtectedTimestamp(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	for _, tc := range []struct {
+		desc              string
+		forceJobFailure   bool
+		expectedJobStatus string
+		expectError       bool
+	}{
+		{
+			desc:              "job success with cleanup",
+			forceJobFailure:   false,
+			expectedJobStatus: "succeeded",
+			expectError:       false,
+		},
+		{
+			desc:              "job failure with cleanup",
+			forceJobFailure:   true,
+			expectedJobStatus: "failed",
+			expectError:       true,
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			var blockInspectExecution atomic.Bool
+			var protectedTimestampCreated atomic.Bool
+
+			ctx := context.Background()
+			s, db, _ := serverutils.StartServer(t, base.TestServerArgs{
+				Knobs: base.TestingKnobs{
+					Inspect: &sql.InspectTestingKnobs{
+						OnInspectAfterProtectedTimestamp: func() error {
+							protectedTimestampCreated.Store(true)
+							// Block execution until we've verified the protected timestamp
+							for blockInspectExecution.Load() {
+								time.Sleep(10 * time.Millisecond)
+							}
+							if tc.forceJobFailure {
+								return errors.New("forced job failure for testing")
+							}
+							return nil
+						},
+					},
+				},
+			})
+			defer s.Stopper().Stop(ctx)
+
+			runner := sqlutils.MakeSQLRunner(db)
+			runner.Exec(t, `
+				CREATE DATABASE db;
+				SET enable_scrub_job = true;
+				CREATE TABLE db.t (
+					id INT PRIMARY KEY,
+					val INT
+				);
+				CREATE INDEX i1 on db.t (val);
+				INSERT INTO db.t VALUES (1, 2), (2, 3);`)
+
+			// Start blocking inspection execution
+			blockInspectExecution.Store(true)
+
+			// Start INSPECT job with AS OF timestamp in a goroutine
+			errCh := make(chan error, 1)
+			go func() {
+				_, err := db.Exec("EXPERIMENTAL SCRUB TABLE db.t AS OF SYSTEM TIME '-1us'")
+				errCh <- err
+			}()
+
+			// Wait for the protected timestamp hook to be called
+			testutils.SucceedsSoon(t, func() error {
+				if !protectedTimestampCreated.Load() {
+					return errors.New("protected timestamp hook not called yet")
+				}
+				return nil
+			})
+
+			// Get the job ID
+			var jobID int64
+			runner.QueryRow(t, `
+				SELECT id
+				FROM crdb_internal.system_jobs
+				WHERE job_type = 'INSPECT' AND status = 'running'
+				ORDER BY created DESC
+				LIMIT 1
+			`).Scan(&jobID)
+
+			// Load the job and get protected timestamp record
+			execCfg := s.ApplicationLayer().ExecutorConfig().(sql.ExecutorConfig)
+			job, err := execCfg.JobRegistry.LoadJob(ctx, jobspb.JobID(jobID))
+			require.NoError(t, err)
+
+			details := job.Details().(jobspb.InspectDetails)
+			require.NotNil(t, details.ProtectedTimestampRecord, "protected timestamp record should be set")
+			protectedTSID := *details.ProtectedTimestampRecord
+
+			// Check that the protected timestamp record actually exists in the system
+			var recordExists bool
+			require.NoError(t, execCfg.InternalDB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+				pts := execCfg.ProtectedTimestampProvider.WithTxn(txn)
+				_, err := pts.GetRecord(ctx, protectedTSID)
+				if err != nil {
+					if errors.Is(err, protectedts.ErrNotExists) {
+						recordExists = false
+						return nil
+					}
+					return err
+				}
+				recordExists = true
+				return nil
+			}))
+			require.True(t, recordExists, "protected timestamp record should exist in the system")
+
+			// Allow the job to complete
+			blockInspectExecution.Store(false)
+
+			// Wait for job to complete
+			select {
+			case err := <-errCh:
+				if tc.expectError {
+					require.Error(t, err, "INSPECT job should fail due to forced error")
+				} else {
+					require.NoError(t, err, "INSPECT job should complete successfully")
+				}
+			case <-time.After(30 * time.Second):
+				t.Fatal("INSPECT job did not complete within timeout")
+			}
+
+			// Verify job status
+			var status string
+			runner.QueryRow(t, `
+				SELECT status
+				FROM crdb_internal.system_jobs
+				WHERE id = $1
+			`, jobID).Scan(&status)
+			require.Equal(t, tc.expectedJobStatus, status, "job should have expected status")
+
+			// Verify protected timestamp record is cleaned up
+			testutils.SucceedsSoon(t, func() error {
+				return execCfg.InternalDB.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
+					pts := execCfg.ProtectedTimestampProvider.WithTxn(txn)
+					_, err := pts.GetRecord(ctx, protectedTSID)
+					if err != nil {
+						if errors.Is(err, protectedts.ErrNotExists) {
+							return nil // This is what we want
+						}
+						return err
+					}
+					return errors.New("protected timestamp record still exists")
+				})
+			})
+		})
+	}
+}