sql: make TestScatterResponse work with secondary tenants

pkg/ccl/multiregionccl/multiregion_system_table_test.go

@@ -55,7 +55,7 @@ func TestMrSystemDatabase(t *testing.T) {
 	tenantArgs := base.TestTenantArgs{
 		Settings: cs,
 		TenantID: id,
-		Locality: *cluster.Servers[0].Locality(),
+		Locality: cluster.Servers[0].Locality(),
 	}
 	_, tenantSQL := serverutils.StartTenant(t, cluster.Servers[0], tenantArgs)
 

pkg/ccl/testccl/sqlstatsccl/sql_stats_test.go

@@ -120,7 +120,7 @@ func TestSQLStatsRegions(t *testing.T) {
 		_, tenantDb := serverutils.StartTenant(t, server, base.TestTenantArgs{
 			Settings: st,
 			TenantID: roachpb.MustMakeTenantID(11),
-			Locality: *server.Locality(),
+			Locality: server.Locality(),
 		})
 		tenantDbs = append(tenantDbs, tenantDb)
 	}

pkg/kv/kvserver/client_lease_test.go

@@ -1054,15 +1054,17 @@ func TestLeasePreferencesDuringOutage(t *testing.T) {
 
 	srv, err := tc.FindMemberServer(newLeaseHolder.StoreID)
 	require.NoError(t, err)
-	region, ok := srv.Locality().Find("region")
+	loc := srv.Locality()
+	region, ok := loc.Find("region")
 	require.True(t, ok)
 	require.Equal(t, "us", region)
 	require.Equal(t, 3, len(repl.Desc().Replicas().Voters().VoterDescriptors()))
 	// Validate that we upreplicated outside of SF.
 	for _, replDesc := range repl.Desc().Replicas().Voters().VoterDescriptors() {
 		serv, err := tc.FindMemberServer(replDesc.StoreID)
 		require.NoError(t, err)
-		dc, ok := serv.Locality().Find("dc")
+		memberLoc := serv.Locality()
+		dc, ok := memberLoc.Find("dc")
 		require.True(t, ok)
 		require.NotEqual(t, "sf", dc)
 	}

pkg/server/testserver.go

@@ -901,11 +901,27 @@ func (t *testTenant) DistSenderI() interface{} {
 	return t.sql.execCfg.DistSender
 }
 
+// NodeDescStoreI is part of the serverutils.ApplicationLayerInterface.
+func (t *testTenant) NodeDescStoreI() interface{} {
+	return t.sql.execCfg.DistSQLPlanner.NodeDescStore()
+}
+
 // InternalDB is part of the serverutils.ApplicationLayerInterface.
 func (t *testTenant) InternalDB() interface{} {
 	return t.sql.internalDB
 }
 
+// Locality is part of the serverutils.ApplicationLayerInterface.
+func (t *testTenant) Locality() roachpb.Locality {
+	return t.Cfg.Locality
+}
+
+// DistSQLPlanningNodeID is part of the serverutils.ApplicationLayerInterface.
+func (t *testTenant) DistSQLPlanningNodeID() roachpb.NodeID {
+	// See comments on replicaoracle.Config.
+	return 0
+}
+
 // LeaseManager is part of the serverutils.ApplicationLayerInterface.
 func (t *testTenant) LeaseManager() interface{} {
 	return t.sql.leaseMgr
@@ -1704,9 +1720,14 @@ func (ts *testServer) MustGetSQLNetworkCounter(name string) int64 {
 	return mustGetSQLCounterForRegistry(reg, name)
 }
 
-// Locality is part of the serverutils.StorageLayerInterface.
-func (ts *testServer) Locality() *roachpb.Locality {
-	return &ts.cfg.Locality
+// Locality is part of the serverutils.ApplicationLayerInterface.
+func (ts *testServer) Locality() roachpb.Locality {
+	return ts.cfg.Locality
+}
+
+// DistSQLPlanningNodeID is part of the serverutils.ApplicationLayerInterface.
+func (ts *testServer) DistSQLPlanningNodeID() roachpb.NodeID {
+	return ts.NodeID()
 }
 
 // LeaseManager is part of the serverutils.ApplicationLayerInterface.
@@ -1734,6 +1755,11 @@ func (ts *testServer) DistSenderI() interface{} {
 	return ts.distSender
 }
 
+// NodeDescStoreI is part of the serverutils.ApplicationLayerInterface.
+func (ts *testServer) NodeDescStoreI() interface{} {
+	return ts.sqlServer.execCfg.DistSQLPlanner.NodeDescStore()
+}
+
 // MigrationServer is part of the serverutils.ApplicationLayerInterface.
 func (ts *testServer) MigrationServer() interface{} {
 	return ts.topLevelServer.migrationServer

pkg/sql/distsql_physical_planner.go

@@ -4768,3 +4768,7 @@ func (dsp *DistSQLPlanner) createPlanForInsert(
 		execinfrapb.Ordering{})
 	return plan, nil
 }
+
+func (dsp *DistSQLPlanner) NodeDescStore() kvcoord.NodeDescStore {
+	return dsp.nodeDescs
+}

pkg/sql/physicalplan/BUILD.bazel

@@ -54,7 +54,6 @@ go_test(
     shard_count = 16,
     deps = [
         "//pkg/base",
-        "//pkg/gossip",
         "//pkg/keys",
         "//pkg/kv",
         "//pkg/kv/kvclient/kvcoord",

pkg/sql/physicalplan/aggregator_funcs_test.go

@@ -18,7 +18,6 @@ import (
 
 	"github.com/cockroachdb/apd/v3"
 	"github.com/cockroachdb/cockroach/pkg/base"
-	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
@@ -147,17 +146,17 @@ func checkDistAggregationInfo(
 			Spans: make([]roachpb.Span, 1),
 		}
 		if err := rowenc.InitIndexFetchSpec(
-			&tr.FetchSpec, keys.SystemSQLCodec, tableDesc, tableDesc.GetPrimaryIndex(), columnIDs,
+			&tr.FetchSpec, srv.Codec(), tableDesc, tableDesc.GetPrimaryIndex(), columnIDs,
 		); err != nil {
 			t.Fatal(err)
 		}
 
 		var err error
-		tr.Spans[0].Key, err = randgen.TestingMakePrimaryIndexKey(tableDesc, startPK)
+		tr.Spans[0].Key, err = randgen.TestingMakePrimaryIndexKeyForTenant(tableDesc, srv.Codec(), startPK)
 		if err != nil {
 			t.Fatal(err)
 		}
-		tr.Spans[0].EndKey, err = randgen.TestingMakePrimaryIndexKey(tableDesc, endPK)
+		tr.Spans[0].EndKey, err = randgen.TestingMakePrimaryIndexKeyForTenant(tableDesc, srv.Codec(), endPK)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -449,8 +448,9 @@ func TestSingleArgumentDistAggregateFunctions(t *testing.T) {
 	defer log.Scope(t).Close(t)
 	const numRows = 100
 
-	tc := serverutils.StartCluster(t, 1, base.TestClusterArgs{})
-	defer tc.Stopper().Stop(context.Background())
+	srv, db, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
+	defer srv.Stopper().Stop(context.Background())
+	ts := srv.ApplicationLayer()
 
 	// Create a table with a few columns:
 	//  - k - primary key with values from 0 to number of rows
@@ -466,7 +466,7 @@ func TestSingleArgumentDistAggregateFunctions(t *testing.T) {
 	//  - random ten bytes
 	rng, _ := randutil.NewTestRand()
 	sqlutils.CreateTable(
-		t, tc.ServerConn(0), "t",
+		t, db, "t",
 		"k INT PRIMARY KEY, int1 INT, int2 INT, int3 INT, bool1 BOOL, bool2 BOOL, dec1 DECIMAL, dec2 DECIMAL, float1 FLOAT, float2 FLOAT, b BYTES",
 		numRows,
 		func(row int) []tree.Datum {
@@ -489,8 +489,7 @@ func TestSingleArgumentDistAggregateFunctions(t *testing.T) {
 		},
 	)
 
-	kvDB := tc.Server(0).DB()
-	desc := desctestutils.TestingGetPublicTableDescriptor(kvDB, keys.SystemSQLCodec, "test", "t")
+	desc := desctestutils.TestingGetPublicTableDescriptor(kvDB, ts.Codec(), "test", "t")
 
 	for fn, info := range physicalplan.DistAggregationTable {
 		if fn == execinfrapb.AnyNotNull {
@@ -528,7 +527,8 @@ func TestSingleArgumentDistAggregateFunctions(t *testing.T) {
 				name := fmt.Sprintf("%s/%s/%d", fn, col.GetName(), numRows)
 				t.Run(name, func(t *testing.T) {
 					checkDistAggregationInfo(
-						context.Background(), t, tc.Server(0), desc, []int{col.Ordinal()},
+						// TODO(#76378): pass ts, not srv, here.
+						context.Background(), t, srv, desc, []int{col.Ordinal()},
 						numRows, fn, info,
 					)
 				})
@@ -550,8 +550,11 @@ func TestTwoArgumentRegressionAggregateFunctions(t *testing.T) {
 	defer log.Scope(t).Close(t)
 	const numRows = 100
 
-	tc := serverutils.StartCluster(t, 1, base.TestClusterArgs{})
-	defer tc.Stopper().Stop(context.Background())
+	srv, db, kvDB := serverutils.StartServer(t, base.TestServerArgs{
+		DefaultTestTenant: base.TestIsForStuffThatShouldWorkWithSecondaryTenantsButDoesntYet(108763),
+	})
+	defer srv.Stopper().Stop(context.Background())
+	ts := srv.ApplicationLayer()
 
 	// Create a table with a few columns:
 	//  - k - primary key with values from 0 to number of rows
@@ -563,7 +566,7 @@ func TestTwoArgumentRegressionAggregateFunctions(t *testing.T) {
 	//  - random decimals (with some NULLs)
 	rng, _ := randutil.NewTestRand()
 	sqlutils.CreateTable(
-		t, tc.ServerConn(0), "t",
+		t, db, "t",
 		"k INT PRIMARY KEY, int1 INT, dec1 DECIMAL, float1 FLOAT, int2 INT, dec2 DECIMAL, float2 FLOAT",
 		numRows,
 		func(row int) []tree.Datum {
@@ -579,8 +582,7 @@ func TestTwoArgumentRegressionAggregateFunctions(t *testing.T) {
 		},
 	)
 
-	kvDB := tc.Server(0).DB()
-	desc := desctestutils.TestingGetTableDescriptor(kvDB, keys.SystemSQLCodec, "test", "public", "t")
+	desc := desctestutils.TestingGetTableDescriptor(kvDB, ts.Codec(), "test", "public", "t")
 
 	for fn, info := range physicalplan.DistAggregationTable {
 		if !isTwoArgumentFunction(fn) {
@@ -598,7 +600,8 @@ func TestTwoArgumentRegressionAggregateFunctions(t *testing.T) {
 					name := fmt.Sprintf("%s/%s-%s/%d", fn, cols[i].GetName(), cols[j].GetName(), numRows)
 					t.Run(name, func(t *testing.T) {
 						checkDistAggregationInfo(
-							context.Background(), t, tc.Server(0), desc, []int{i, j}, numRows,
+							// TODO(#76378): pass ts, not srv, here.
+							context.Background(), t, srv, desc, []int{i, j}, numRows,
 							fn, info,
 						)
 					})

pkg/sql/physicalplan/fake_span_resolver_test.go

@@ -34,7 +34,11 @@ func TestFakeSpanResolver(t *testing.T) {
 	defer log.Scope(t).Close(t)
 	ctx := context.Background()
 
-	tc := serverutils.StartCluster(t, 3, base.TestClusterArgs{})
+	tc := serverutils.StartCluster(t, 3, base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{
+			DefaultTestTenant: base.TestIsForStuffThatShouldWorkWithSecondaryTenantsButDoesntYet(108763),
+		},
+	})
 	defer tc.Stopper().Stop(ctx)
 
 	sqlutils.CreateTable(