sql: use JobRegistry of ApplicationLayer in tests

yuzefovich · yuzefovich · commit 613f230869fb · 2025-10-23T23:43:44.000-07:00
One of the common modifications we need to make for tests to work with
secondary tenants is using the JobRegistry of the ApplicationLayer (as
opposed to the one belonging to the system tenant). This commit audits
all tests in `pkg/sql` directory to ensure the right registry is used.
Note that tests still might need further modifications.

Release note: None
diff --git a/pkg/sql/catalog/externalcatalog/external_catalog_test.go b/pkg/sql/catalog/externalcatalog/external_catalog_test.go
@@ -117,7 +117,7 @@ func TestExtractIngestExternalCatalog(t *testing.T) {
 		sqlDB.CheckQueryResults(t, "SELECT schema_name,table_name FROM [SHOW TABLES]", [][]string{{"public", "tab1"}, {"public", "tab2_rename"}})
 
 		require.NoError(t, sqltestutils.TestingDescsTxn(ctx, srv, func(ctx context.Context, txn isql.Txn, col *descs.Collection) error {
-			return DropIngestedExternalCatalog(ctx, &execCfg, sqlUser, written, txn, srv.JobRegistry().(*jobs.Registry), col, "test gc")
+			return DropIngestedExternalCatalog(ctx, &execCfg, sqlUser, written, txn, s.JobRegistry().(*jobs.Registry), col, "test gc")
 		}))
 		var res int
 		sqlDB.QueryRow(t, "SELECT count(*) FROM [SHOW TABLES]").Scan(&res)
diff --git a/pkg/sql/catalog/schematelemetry/schema_telemetry_test.go b/pkg/sql/catalog/schematelemetry/schema_telemetry_test.go
@@ -97,8 +97,9 @@ func TestSchemaTelemetryJob(t *testing.T) {
 	defer log.Scope(t).Close(t)
 
 	ctx := context.Background()
-	s, db, _ := serverutils.StartServer(t, makeTestServerArgs())
-	defer s.Stopper().Stop(ctx)
+	srv, db, _ := serverutils.StartServer(t, makeTestServerArgs())
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 	tdb := sqlutils.MakeSQLRunner(db)
 	tdb.Exec(t, `SET CLUSTER SETTING server.eventlog.enabled = true`)
 
diff --git a/pkg/sql/gcjob_test/gc_job_test.go b/pkg/sql/gcjob_test/gc_job_test.go
@@ -91,9 +91,10 @@ func doTestSchemaChangeGCJob(t *testing.T, dropItem DropItem, ttlTime TTLTime) {
 			return nil
 		},
 	}
-	s, db, kvDB := serverutils.StartServer(t, params)
+	srv, db, kvDB := serverutils.StartServer(t, params)
 	ctx := context.Background()
-	defer s.Stopper().Stop(ctx)
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 	// The deferred call to unblock the GC job needs to run before the deferred
 	// call to stop the TestServer. Otherwise, the quiesce step of shutting down
 	// can hang forever waiting for the GC job.
@@ -648,10 +649,11 @@ func TestDropWithDeletedDescriptor(t *testing.T) {
 				},
 			}
 		}
-		s, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{
+		srv, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{
 			Knobs: knobs,
 		})
-		defer s.Stopper().Stop(ctx)
+		defer srv.Stopper().Stop(ctx)
+		s := srv.ApplicationLayer()
 		defer cancel()
 		tdb := sqlutils.MakeSQLRunner(sqlDB)
 
diff --git a/pkg/sql/inspect/progress_test.go b/pkg/sql/inspect/progress_test.go
@@ -37,7 +37,7 @@ func setupProgressTestInfra(
 	ctx := context.Background()
 	s := serverutils.StartServerOnly(t, base.TestServerArgs{})
 
-	registry := s.JobRegistry().(*jobs.Registry)
+	registry := s.ApplicationLayer().JobRegistry().(*jobs.Registry)
 
 	// Create a job using INSPECT details
 	record := jobs.Record{
@@ -445,7 +445,7 @@ func TestInspectProgressTracker_ProgressFlushConditions(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	ctx, s, _, _, cleanup := setupProgressTestInfra(t)
+	ctx, s, registry, _, cleanup := setupProgressTestInfra(t)
 	defer cleanup()
 
 	const totalChecks = 1000
@@ -573,7 +573,7 @@ func TestInspectProgressTracker_ProgressFlushConditions(t *testing.T) {
 				Username: username.TestUserName(),
 			}
 
-			freshJob, err := s.JobRegistry().(*jobs.Registry).CreateJobWithTxn(ctx, record, s.JobRegistry().(*jobs.Registry).MakeJobID(), nil)
+			freshJob, err := registry.CreateJobWithTxn(ctx, record, registry.MakeJobID(), nil)
 			require.NoError(t, err)
 
 			freshTracker := newInspectProgressTracker(freshJob, &s.ClusterSettings().SV, s.InternalDB().(isql.DB))
diff --git a/pkg/sql/rowexec/backfiller_test.go b/pkg/sql/rowexec/backfiller_test.go
@@ -66,7 +66,7 @@ func TestWriteResumeSpan(t *testing.T) {
 
 	ctx := context.Background()
 
-	server, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{
+	srv, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{
 		Knobs: base.TestingKnobs{
 			// Disable all schema change execution.
 			SQLSchemaChanger: &sql.SchemaChangerTestingKnobs{
@@ -76,7 +76,8 @@ func TestWriteResumeSpan(t *testing.T) {
 			},
 		},
 	})
-	defer server.Stopper().Stop(ctx)
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 
 	sqlRunner := sqlutils.MakeSQLRunner(sqlDB)
 	sqlRunner.Exec(t, `SET create_table_with_schema_locked=false`)
@@ -97,7 +98,7 @@ func TestWriteResumeSpan(t *testing.T) {
 		{Key: roachpb.Key("q"), EndKey: roachpb.Key("r")},
 	}
 
-	registry := server.JobRegistry().(*jobs.Registry)
+	registry := s.JobRegistry().(*jobs.Registry)
 	tableDesc := desctestutils.TestingGetMutableExistingTableDescriptor(
 		kvDB, keys.SystemSQLCodec, "t", "test")
 
@@ -175,7 +176,7 @@ func TestWriteResumeSpan(t *testing.T) {
 		if test.resume.Key != nil {
 			finished.EndKey = test.resume.Key
 		}
-		if err := sqltestutils.TestingDescsTxn(ctx, server, func(ctx context.Context, txn isql.Txn, col *descs.Collection) error {
+		if err := sqltestutils.TestingDescsTxn(ctx, s, func(ctx context.Context, txn isql.Txn, col *descs.Collection) error {
 			return TestingWriteResumeSpan(
 				ctx,
 				txn,
@@ -215,7 +216,7 @@ func TestWriteResumeSpan(t *testing.T) {
 	}
 
 	var got []roachpb.Span
-	if err := sqltestutils.TestingDescsTxn(ctx, server, func(ctx context.Context, txn isql.Txn, col *descs.Collection) (err error) {
+	if err := sqltestutils.TestingDescsTxn(ctx, s, func(ctx context.Context, txn isql.Txn, col *descs.Collection) (err error) {
 		got, _, _, err = rowexec.GetResumeSpans(
 			ctx, registry, txn, keys.SystemSQLCodec, col, tableDesc.ID, mutationID, backfill.IndexMutationFilter)
 		return err
diff --git a/pkg/sql/schema_changer_test.go b/pkg/sql/schema_changer_test.go
@@ -90,8 +90,9 @@ func TestSchemaChangeProcess(t *testing.T) {
 
 	params, _ := createTestServerParamsAllowTenants()
 
-	s, sqlDB, kvDB := serverutils.StartServer(t, params)
-	defer s.Stopper().Stop(context.Background())
+	srv, sqlDB, kvDB := serverutils.StartServer(t, params)
+	defer srv.Stopper().Stop(context.Background())
+	s := srv.ApplicationLayer()
 
 	// The descriptor changes made must have an immediate effect
 	// so disable leases on tables.
@@ -6499,7 +6500,7 @@ func TestRevertingJobsOnDatabasesAndSchemas(t *testing.T) {
 		params.Knobs = base.TestingKnobs{
 			SQLSchemaChanger: &sql.SchemaChangerTestingKnobs{
 				RunBeforeResume: func(jobID jobspb.JobID) error {
-					scJob, err := s.JobRegistry().(*jobs.Registry).LoadJob(ctx, jobID)
+					scJob, err := s.ApplicationLayer().JobRegistry().(*jobs.Registry).LoadJob(ctx, jobID)
 					if err != nil {
 						return err
 					}
@@ -6584,7 +6585,7 @@ func TestRevertingJobsOnDatabasesAndSchemas(t *testing.T) {
 		params.Knobs = base.TestingKnobs{
 			SQLSchemaChanger: &sql.SchemaChangerTestingKnobs{
 				RunBeforeResume: func(jobID jobspb.JobID) error {
-					scJob, err := s.JobRegistry().(*jobs.Registry).LoadJob(ctx, jobID)
+					scJob, err := s.ApplicationLayer().JobRegistry().(*jobs.Registry).LoadJob(ctx, jobID)
 					if err != nil {
 						return err
 					}
@@ -6658,7 +6659,7 @@ func TestCheckConstraintDropAndColumn(t *testing.T) {
 				// `channel` below.
 				lockHeld := true
 				jobControlMu.Lock()
-				scJob, err := s.JobRegistry().(*jobs.Registry).LoadJob(ctx, jobID)
+				scJob, err := s.ApplicationLayer().JobRegistry().(*jobs.Registry).LoadJob(ctx, jobID)
 				if err != nil {
 					jobControlMu.Unlock()
 					return err
@@ -6790,7 +6791,7 @@ func TestJobsWithoutMutationsAreCancelable(t *testing.T) {
 
 	var registry *jobs.Registry
 	var scJobID jobspb.JobID
-	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
+	srv, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
 		Knobs: base.TestingKnobs{SQLSchemaChanger: &sql.SchemaChangerTestingKnobs{
 			RunBeforeResume: func(jobID jobspb.JobID) error {
 				job, err := registry.LoadJob(ctx, jobID)
@@ -6804,7 +6805,8 @@ func TestJobsWithoutMutationsAreCancelable(t *testing.T) {
 			},
 		}},
 	})
-	defer s.Stopper().Stop(ctx)
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 	tdb := sqlutils.MakeSQLRunner(sqlDB)
 	registry = s.JobRegistry().(*jobs.Registry)
 
diff --git a/pkg/sql/sqlstats/persistedsqlstats/scheduled_sql_stats_compaction_test.go b/pkg/sql/sqlstats/persistedsqlstats/scheduled_sql_stats_compaction_test.go
@@ -36,7 +36,7 @@ import (
 )
 
 type testHelper struct {
-	server           serverutils.TestServerInterface
+	server           serverutils.ApplicationLayerInterface
 	sqlDB            *sqlutils.SQLRunner
 	env              *jobstest.JobSchedulerTestEnv
 	cfg              *scheduledjobs.JobExecutionConfig
@@ -90,7 +90,7 @@ func newTestHelper(
 	require.NotNil(t, helper.cfg)
 
 	helper.sqlDB = sqlutils.MakeSQLRunner(db)
-	helper.server = srv
+	helper.server = srv.ApplicationLayer()
 
 	return helper, func() {
 		srv.Stopper().Stop(context.Background())
@@ -137,7 +137,7 @@ func TestScheduledSQLStatsCompaction(t *testing.T) {
 	helper.sqlDB.Exec(t, "SELECT 1; SELECT 1, 1")
 	sqlstatstestutil.WaitForStatementEntriesAtLeast(t, helper.sqlDB, 2)
 
-	helper.server.ApplicationLayer().SQLServer().(*sql.Server).GetSQLStatsProvider().MaybeFlush(ctx, helper.server.AppStopper())
+	helper.server.SQLServer().(*sql.Server).GetSQLStatsProvider().MaybeFlush(ctx, helper.server.AppStopper())
 	helper.sqlDB.Exec(t, "SET CLUSTER SETTING sql.stats.persisted_rows.max = 1")
 
 	stmtStatsCnt, txnStatsCnt := getPersistedStatsEntry(t, helper.sqlDB)
