Merge #158199

158199: sql/inspect: use historical descriptors for INSPECT AS OF r=spilchen a=spilchen

Fix a bug in `INSPECT ... AS OF` that incorrectly used the latest table descriptor instead of a historical one. This could result in incorrect index ID usage if a schema change occurred after the specified AS OF time.

Since this tests schema changes interleaved with INSPECT, this will also create/start the INSPECT job in a separate transaction. This was needed to avoid getting transaction retry errors.

Informs: #158197
Release note: none

Co-authored-by: Matt Spilchen <matt.spilchen@cockroachlabs.com>

Author: craig[bot]

pkg/sql/importer/import_job.go

@@ -351,7 +351,6 @@ func (r *importResumer) Resume(ctx context.Context, execCtx interface{}) error {
 				ctx,
 				fmt.Sprintf("import-validation-%s", tableName),
 				p.ExecCfg(),
-				nil, /* txn */
 				checks,
 				setPublicTimestamp,
 			)

pkg/sql/inspect/index_consistency_check.go

@@ -447,8 +447,19 @@ func (c *indexConsistencyCheck) Close(context.Context) error {
 // descriptor and index metadata in the indexConsistencyCheck struct.
 func (c *indexConsistencyCheck) loadCatalogInfo(ctx context.Context) error {
 	return c.flowCtx.Cfg.DB.DescsTxn(ctx, func(ctx context.Context, txn descs.Txn) error {
+		if !c.asOf.IsEmpty() {
+			if err := txn.KV().SetFixedTimestamp(ctx, c.asOf); err != nil {
+				return err
+			}
+		}
+
+		byIDGetter := txn.Descriptors().ByIDWithLeased(txn.KV())
+		if !c.asOf.IsEmpty() {
+			byIDGetter = txn.Descriptors().ByIDWithoutLeased(txn.KV())
+		}
+
 		var err error
-		c.tableDesc, err = txn.Descriptors().ByIDWithLeased(txn.KV()).WithoutNonPublic().Get().Table(ctx, c.tableID)
+		c.tableDesc, err = byIDGetter.WithoutNonPublic().Get().Table(ctx, c.tableID)
 		if err != nil {
 			return err
 		}

pkg/sql/inspect/index_consistency_check_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/rowenc"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
@@ -827,7 +828,7 @@ func TestInspectWithoutASOFSchemaChange(t *testing.T) {
 			t.Logf("error polling job status: %v", err)
 			return false
 		}
-		return status != "running" && status != "pending"
+		return status == "failed"
 	}, 60*time.Second, time.Second)
 
 	require.Equal(t, "failed", status, "detached job should fail after schema change")
@@ -838,3 +839,73 @@ func TestInspectWithoutASOFSchemaChange(t *testing.T) {
 	defer rows.Close()
 	require.False(t, rows.Next(), "no inspect errors are recorded when the job aborts before running checks")
 }
+
+// TestInspectASOFAfterPrimaryKeySwapFails runs an INSPECT ASOF statement while
+// doing a schema change after the ASOF time.
+func TestInspectASOFAfterPrimaryKeySwap(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(ctx)
+	r := sqlutils.MakeSQLRunner(db)
+
+	r.ExecMultiple(t,
+		`SET enable_inspect_command = true`,
+		`CREATE DATABASE t`,
+		`CREATE TABLE t.pk_swap (
+			old_pk INT PRIMARY KEY,
+			new_pk INT UNIQUE NOT NULL,
+			payload INT NOT NULL,
+			INDEX payload_idx (payload)
+		)`,
+		`INSERT INTO t.pk_swap SELECT i, i+100, i+200 FROM generate_series(1, 5) AS g(i)`,
+	)
+
+	var asOf string
+	r.QueryRow(t, `SELECT cluster_logical_timestamp()`).Scan(&asOf)
+
+	// Rewrite the table descriptors by swapping the primary key.
+	r.Exec(t, `ALTER TABLE t.pk_swap ALTER PRIMARY KEY USING COLUMNS (new_pk)`)
+
+	// Helper function to wait for the latest INSPECT job to complete successfully.
+	waitForInspectJob := func() {
+		var jobID int64
+		var status string
+		var fractionCompleted float64
+		testutils.SucceedsSoon(t, func() error {
+			row := db.QueryRow(`SELECT job_id, status, fraction_completed FROM [SHOW JOBS] WHERE job_type = 'INSPECT' ORDER BY job_id DESC LIMIT 1`)
+			if err := row.Scan(&jobID, &status, &fractionCompleted); err != nil {
+				return err
+			}
+			if status == "succeeded" || status == "failed" {
+				return nil
+			}
+			return errors.Newf("job is not in the succeeded or failed state: %q", status)
+		})
+		require.Equal(t, "succeeded", status, "INSPECT job should succeed")
+		require.InEpsilon(t, 1.0, fractionCompleted, 0.01, "progress should reach 100%% on successful completion")
+		requireCheckCountsMatch(t, r, jobID)
+	}
+
+	t.Run("non-detached", func(t *testing.T) {
+		r.Exec(t, fmt.Sprintf(`INSPECT TABLE t.pk_swap AS OF SYSTEM TIME '%s' WITH OPTIONS INDEX ALL`, asOf))
+		waitForInspectJob()
+	})
+
+	t.Run("detached with implicit transaction", func(t *testing.T) {
+		r.Exec(t, fmt.Sprintf(`INSPECT TABLE t.pk_swap AS OF SYSTEM TIME '%s' WITH OPTIONS INDEX ALL, DETACHED`, asOf))
+		waitForInspectJob()
+	})
+
+	t.Run("detached with explicit transaction", func(t *testing.T) {
+		_, err := db.Exec("BEGIN AS OF SYSTEM TIME '" + asOf + "'")
+		require.NoError(t, err)
+		_, err = db.Exec(fmt.Sprintf(`INSPECT TABLE t.pk_swap AS OF SYSTEM TIME '%s' WITH OPTIONS INDEX ALL, DETACHED`, asOf))
+		require.NoError(t, err)
+		_, err = db.Exec("COMMIT")
+		// Expecting the commit to fail due to the schema change after ASOF time.
+		require.Contains(t, err.Error(), "TransactionRetryError")
+	})
+}

pkg/sql/inspect/inspect_job_entry.go

@@ -25,12 +25,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 )
 
-// TriggerJob starts an inspect job for the snapshot.
+// TriggerJob starts an inspect job for the snapshot. This always uses a fresh
+// transaction to create and start the job.
 func TriggerJob(
 	ctx context.Context,
 	jobRecordDescription string,
 	execCfg *sql.ExecutorConfig,
-	txn isql.Txn,
 	checks []*jobspb.InspectDetails_Check,
 	asOf hlc.Timestamp,
 ) (*jobs.StartableJob, error) {
@@ -55,35 +55,15 @@ func TriggerJob(
 	}
 
 	var job *jobs.StartableJob
-	if txn == nil {
-		if err := execCfg.InternalDB.Txn(ctx, func(ctx context.Context, txn isql.Txn) (err error) {
-			return execCfg.JobRegistry.CreateStartableJobWithTxn(ctx, &job, record.JobID, txn, record)
-		}); err != nil {
-			if job != nil {
-				if cleanupErr := job.CleanupOnRollback(ctx); cleanupErr != nil {
-					log.Dev.Warningf(ctx, "failed to cleanup StartableJob: %v", cleanupErr)
-				}
+	if err := execCfg.InternalDB.Txn(ctx, func(ctx context.Context, txn isql.Txn) (err error) {
+		return execCfg.JobRegistry.CreateStartableJobWithTxn(ctx, &job, record.JobID, txn, record)
+	}); err != nil {
+		if job != nil {
+			if cleanupErr := job.CleanupOnRollback(ctx); cleanupErr != nil {
+				log.Dev.Warningf(ctx, "failed to cleanup StartableJob: %v", cleanupErr)
 			}
-			return nil, err
-		}
-	} else {
-		if err := func() (err error) {
-			defer func() {
-				if err == nil || job == nil {
-					return
-				}
-				if cleanupErr := job.CleanupOnRollback(ctx); cleanupErr != nil {
-					log.Dev.Errorf(ctx, "failed to cleanup job: %v", cleanupErr)
-				}
-			}()
-			if err := execCfg.JobRegistry.CreateStartableJobWithTxn(ctx, &job, record.JobID, txn, record); err != nil {
-				return err
-			}
-
-			return txn.KV().Commit(ctx)
-		}(); err != nil {
-			return nil, err
 		}
+		return nil, err
 	}
 
 	if err := job.Start(ctx); err != nil {

pkg/sql/inspect/inspect_plan.go

@@ -207,7 +207,7 @@ func inspectPlanHook(
 		return nil, nil, false, err
 	}
 
-	if detached {
+	if detached && !p.ExtendedEvalContext().TxnImplicit {
 		fn := func(ctx context.Context, resultsCh chan<- tree.Datums) error {
 			jobID := p.ExecCfg().JobRegistry.MakeJobID()
 			jr := makeInspectJobRecord(tree.AsString(stmt), jobID, run.checks, run.asOfTimestamp)
@@ -216,12 +216,7 @@ func inspectPlanHook(
 			); err != nil {
 				return err
 			}
-			var notice pgnotice.Notice
-			if p.ExtendedEvalContext().TxnImplicit {
-				notice = pgnotice.Newf("INSPECT job %d running in the background", jobID)
-			} else {
-				notice = pgnotice.Newf("INSPECT job %d queued; will start after the current transaction commits", jobID)
-			}
+			notice := pgnotice.Newf("INSPECT job %d queued; will start after the current transaction commits", jobID)
 			if err := p.SendClientNotice(ctx, notice, true /* immediateFlush */); err != nil {
 				return err
 			}
@@ -232,30 +227,37 @@ func inspectPlanHook(
 	}
 
 	fn := func(ctx context.Context, resultsCh chan<- tree.Datums) error {
-		// We create the job record in the planner's transaction to ensure that
-		// the job record creation happens transactionally.
-		plannerTxn := p.InternalSQLTxn()
-
-		sj, err := TriggerJob(ctx, tree.AsString(stmt), p.ExecCfg(), plannerTxn, run.checks, run.asOfTimestamp)
+		sj, err := TriggerJob(ctx, tree.AsString(stmt), p.ExecCfg(), run.checks, run.asOfTimestamp)
 		if err != nil {
 			return err
 		}
 
-		if err = p.SendClientNotice(
-			ctx,
-			pgnotice.Newf("waiting for INSPECT job to complete: %s\nIf the statement is canceled, the job will continue in the background.", sj.ID()),
-			true, /* immediateFlush */
-		); err != nil {
+		var notice pgnotice.Notice
+		if detached {
+			notice = pgnotice.Newf("INSPECT job %d running in the background", sj.ID())
+		} else {
+			notice = pgnotice.Newf("waiting for INSPECT job to complete: %s\nIf the statement is canceled, the job will continue in the background.", sj.ID())
+		}
+		if err = p.SendClientNotice(ctx, notice, true /* immediateFlush */); err != nil {
 			return err
 		}
 
-		if err := sj.AwaitCompletion(ctx); err != nil {
+		if detached {
+			resultsCh <- tree.Datums{tree.NewDInt(tree.DInt(sj.ID()))}
+			return nil
+		}
+
+		if err = sj.AwaitCompletion(ctx); err != nil {
 			return err
 		}
 		return sj.ReportExecutionResults(ctx, resultsCh)
 	}
 
-	return fn, jobs.InspectJobExecutionResultHeader, false, nil
+	header := jobs.InspectJobExecutionResultHeader
+	if detached {
+		header = jobs.DetachedJobExecutionResultHeader
+	}
+	return fn, header, false, nil
 }
 
 func makeInspectJobRecord(