testutils: add testserver APIs useful for DistSQL

Release note: None

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/span_resolver_test.go

@@ -90,7 +90,7 @@ func TestSpanResolverUsesCaches(t *testing.T) {
 		s3.DistSenderI().(*kvcoord.DistSender),
 		s3.GossipI().(*gossip.Gossip),
 		s3.NodeID(),
-		*s3.Locality(),
+		s3.Locality(),
 		s3.Clock(),
 		nil, // rpcCtx
 		replicaoracle.BinPackingChoice)
@@ -202,7 +202,7 @@ func TestSpanResolver(t *testing.T) {
 		s.DistSenderI().(*kvcoord.DistSender),
 		s.GossipI().(*gossip.Gossip),
 		s.NodeID(),
-		*s.Locality(),
+		s.Locality(),
 		s.Clock(),
 		nil, // rpcCtx
 		replicaoracle.BinPackingChoice)
@@ -301,7 +301,7 @@ func TestMixedDirections(t *testing.T) {
 		s.DistSenderI().(*kvcoord.DistSender),
 		s.GossipI().(*gossip.Gossip),
 		s.NodeID(),
-		*s.Locality(),
+		s.Locality(),
 		s.Clock(),
 		nil, // rpcCtx
 		replicaoracle.BinPackingChoice)

pkg/testutils/serverutils/api.go

@@ -404,6 +404,16 @@ type ApplicationLayerInterface interface {
 	// server. The concrete return value is of type
 	// privchecker.SQLPrivilegeChecker (interface).
 	PrivilegeChecker() interface{}
+
+	// NodeDescStoreI returns the node descriptor lookup interface.
+	// The concrete return type is compatible with interface kvcoord.NodeDescStore.
+	NodeDescStoreI() interface{}
+
+	// Locality returns the locality used by the server.
+	Locality() roachpb.Locality
+
+	// DistSQLPlanningNodeID returns the NodeID to use by the DistSQL span resolver.
+	DistSQLPlanningNodeID() roachpb.NodeID
 }
 
 // TenantControlInterface defines the API of a test server that can
@@ -607,14 +617,6 @@ type StorageLayerInterface interface {
 	// TsDB returns the ts.DB instance used by the TestServer.
 	TsDB() interface{}
 
-	// Locality returns a pointer to the locality used by the server.
-	//
-	// TODO(test-eng): investigate if this should really be a pointer.
-	//
-	// TODO(test-eng): Investigate if this method should be on
-	// ApplicationLayerInterface instead.
-	Locality() *roachpb.Locality
-
 	// DefaultSystemZoneConfig returns the internal system zone config
 	// for the server.
 	// Note: most tests should instead use the .DefaultZoneConfig() method