cli,testutils,ts: consolidate RPC client creation

This commit consolidates RPC client creation logic in testserver code
paths in cli, testutils, and ts packages. It is a continuation of the
work done in #147606.

Epic: CRDB-48923
Fixes: #148158
Release note: none

Author: cthumuluru-crdb

pkg/cli/democluster/demo_cluster.go

@@ -795,7 +795,7 @@ func (c *transientCluster) waitForNodeIDReadiness(
 			if err != nil {
 				return err
 			}
-			c.servers[idx].adminClient = serverpb.NewAdminClient(conn)
+			c.servers[idx].adminClient = conn.NewAdminClient()
 
 		}
 		break
@@ -1219,7 +1219,7 @@ func (c *transientCluster) startServerInternal(
 
 	c.servers[serverIdx] = serverEntry{
 		TestServerInterface: s,
-		adminClient:         serverpb.NewAdminClient(conn),
+		adminClient:         conn.NewAdminClient(),
 		nodeID:              nodeID,
 	}
 

pkg/cli/gen.go

@@ -407,7 +407,7 @@ func generateMetricList(ctx context.Context, skipFiltering bool) (map[string]*La
 		if err != nil {
 			return err
 		}
-		client := serverpb.NewAdminClient(conn)
+		client := conn.NewAdminClient()
 
 		resp, err := client.ChartCatalog(ctx, &serverpb.ChartCatalogRequest{})
 		if err != nil {

pkg/server/server_import_ts_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
 	"github.com/cockroachdb/cockroach/pkg/ts"
 	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
@@ -24,7 +25,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
-	"google.golang.org/grpc"
 )
 
 // TestServerWithTimeseriesImport validates the functionality gated behind the
@@ -73,11 +73,11 @@ func TestServerWithTimeseriesImport(t *testing.T) {
 	require.Equal(t, bytesDumped, bytesDumpedAgain)
 }
 
-func dumpTSNonempty(t *testing.T, cc *grpc.ClientConn, dest string) (bytes int64) {
-	ac := serverpb.NewAdminClient(cc)
+func dumpTSNonempty(t *testing.T, cc serverutils.RPCConn, dest string) (bytes int64) {
+	ac := cc.NewAdminClient()
 	names, err := serverpb.GetInternalTimeseriesNamesFromServer(context.Background(), ac)
 	require.NoError(t, err)
-	c, err := tspb.NewTimeSeriesClient(cc).DumpRaw(context.Background(), &tspb.DumpRequest{
+	c, err := cc.NewTimeSeriesClient().DumpRaw(context.Background(), &tspb.DumpRequest{
 		Names: names,
 	})
 	require.NoError(t, err)

pkg/server/testserver.go

@@ -76,7 +76,6 @@ import (
 	"github.com/cockroachdb/logtags"
 	"github.com/cockroachdb/redact"
 	"github.com/gogo/protobuf/proto"
-	"google.golang.org/grpc"
 )
 
 // makeTestConfig returns a config for testing. It overrides the
@@ -2621,7 +2620,7 @@ func (ts *testServer) NewClientRPCContext(
 // RPCClientConn is part of the serverutils.ApplicationLayerInterface.
 func (ts *testServer) RPCClientConn(
 	test serverutils.TestFataler, user username.SQLUsername,
-) *grpc.ClientConn {
+) serverutils.RPCConn {
 	conn, err := ts.RPCClientConnE(user)
 	if err != nil {
 		test.Fatal(err)
@@ -2630,22 +2629,26 @@ func (ts *testServer) RPCClientConn(
 }
 
 // RPCClientConnE is part of the serverutils.ApplicationLayerInterface.
-func (ts *testServer) RPCClientConnE(user username.SQLUsername) (*grpc.ClientConn, error) {
+func (ts *testServer) RPCClientConnE(user username.SQLUsername) (serverutils.RPCConn, error) {
 	ctx := context.Background()
 	rpcCtx := ts.NewClientRPCContext(ctx, user)
-	return rpcCtx.GRPCDialNode(ts.AdvRPCAddr(), ts.NodeID(), ts.Locality(), rpcbase.DefaultClass).Connect(ctx)
+	conn, err := rpcCtx.GRPCDialNode(ts.AdvRPCAddr(), ts.NodeID(), ts.Locality(), rpcbase.DefaultClass).Connect(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return serverutils.FromGRPCConn(conn), nil
 }
 
 // GetAdminClient is part of the serverutils.ApplicationLayerInterface.
 func (ts *testServer) GetAdminClient(test serverutils.TestFataler) serverpb.AdminClient {
 	conn := ts.RPCClientConn(test, username.RootUserName())
-	return serverpb.NewAdminClient(conn)
+	return conn.NewAdminClient()
 }
 
 // GetStatusClient is part of the serverutils.ApplicationLayerInterface.
 func (ts *testServer) GetStatusClient(test serverutils.TestFataler) serverpb.StatusClient {
 	conn := ts.RPCClientConn(test, username.RootUserName())
-	return serverpb.NewStatusClient(conn)
+	return conn.NewStatusClient()
 }
 
 // NewClientRPCContext is part of the serverutils.ApplicationLayerInterface.
@@ -2662,7 +2665,7 @@ func (t *testTenant) NewClientRPCContext(
 // RPCClientConn is part of the serverutils.ApplicationLayerInterface.
 func (t *testTenant) RPCClientConn(
 	test serverutils.TestFataler, user username.SQLUsername,
-) *grpc.ClientConn {
+) serverutils.RPCConn {
 	conn, err := t.RPCClientConnE(user)
 	if err != nil {
 		test.Fatal(err)
@@ -2671,22 +2674,29 @@ func (t *testTenant) RPCClientConn(
 }
 
 // RPCClientConnE is part of the serverutils.ApplicationLayerInterface.
-func (t *testTenant) RPCClientConnE(user username.SQLUsername) (*grpc.ClientConn, error) {
+func (t *testTenant) RPCClientConnE(user username.SQLUsername) (serverutils.RPCConn, error) {
 	ctx := context.Background()
 	rpcCtx := t.NewClientRPCContext(ctx, user)
-	return rpcCtx.GRPCDialPod(t.AdvRPCAddr(), t.SQLInstanceID(), t.Locality(), rpcbase.DefaultClass).Connect(ctx)
+	if !rpcbase.TODODRPC {
+		conn, err := rpcCtx.GRPCDialPod(t.AdvRPCAddr(), t.SQLInstanceID(), t.Locality(), rpcbase.DefaultClass).Connect(ctx)
+		if err != nil {
+			return nil, err
+		}
+		return serverutils.FromGRPCConn(conn), nil
+	}
+	return nil, nil
 }
 
 // GetAdminClient is part of the serverutils.ApplicationLayerInterface.
 func (t *testTenant) GetAdminClient(test serverutils.TestFataler) serverpb.AdminClient {
 	conn := t.RPCClientConn(test, username.RootUserName())
-	return serverpb.NewAdminClient(conn)
+	return conn.NewAdminClient()
 }
 
 // GetStatusClient is part of the serverutils.ApplicationLayerInterface.
 func (t *testTenant) GetStatusClient(test serverutils.TestFataler) serverpb.StatusClient {
 	conn := t.RPCClientConn(test, username.RootUserName())
-	return serverpb.NewStatusClient(conn)
+	return conn.NewStatusClient()
 }
 
 func newClientRPCContext(

pkg/sql/flowinfra/cluster_test.go

@@ -269,7 +269,7 @@ func TestClusterFlow(t *testing.T) {
 		servers[i] = s
 		conns[i] = s.SQLConn(t)
 		conn := s.RPCClientConn(t, username.RootUserName())
-		clients[i] = execinfrapb.NewDistSQLClient(conn)
+		clients[i] = conn.NewDistSQLClient()
 	}
 
 	runTestClusterFlow(t, servers, conns, clients)
@@ -776,7 +776,7 @@ func BenchmarkInfrastructure(b *testing.B) {
 					for i := 0; i < numNodes; i++ {
 						s := tc.Server(i)
 						conn := s.RPCClientConn(b, username.RootUserName())
-						clients = append(clients, execinfrapb.NewDistSQLClient(conn))
+						clients = append(clients, conn.NewDistSQLClient())
 					}
 
 					b.ResetTimer()

pkg/sql/flowinfra/server_test.go

@@ -119,7 +119,7 @@ func TestServer(t *testing.T) {
 			t.Run(fmt.Sprintf("%d", tc.version), func(t *testing.T) {
 				req := *req
 				req.Version = tc.version
-				distSQLClient := execinfrapb.NewDistSQLClient(conn)
+				distSQLClient := conn.NewDistSQLClient()
 				resp, err := distSQLClient.SetupFlow(ctx, &req)
 				if err == nil && resp.Error != nil {
 					err = resp.Error.ErrorDetail(ctx)

pkg/testutils/serverutils/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
         "api.go",
         "conditional_wrap.go",
         "options.go",
+        "rpc_conn.go",
         "test_cluster_shim.go",
         "test_cluster_utils.go",
         "test_server_shim.go",
@@ -23,6 +24,7 @@ go_library(
         "//pkg/config/zonepb",
         "//pkg/keys",
         "//pkg/kv",
+        "//pkg/kv/kvpb",
         "//pkg/kv/kvprober",
         "//pkg/kv/kvserver/liveness/livenesspb",
         "//pkg/multitenant/tenantcapabilities",
@@ -37,9 +39,11 @@ go_library(
         "//pkg/settings/cluster",
         "//pkg/sql/catalog",
         "//pkg/sql/catalog/descpb",
+        "//pkg/sql/execinfrapb",
         "//pkg/storage",
         "//pkg/testutils/pgurlutils",
         "//pkg/testutils/skip",
+        "//pkg/ts/tspb",
         "//pkg/util/envutil",
         "//pkg/util/hlc",
         "//pkg/util/httputil",

pkg/testutils/serverutils/api.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvprober"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
 	"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
@@ -31,13 +32,14 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/server/status"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
-	"google.golang.org/grpc"
 )
 
 // TestServerInterfaceRaw is the interface of server.testServer.
@@ -146,6 +148,18 @@ func (d DeploymentMode) IsExternal() bool {
 	return d == ExternalProcess
 }
 
+// RPCConn defines a common interface for creating RPC clients. It hides the
+// underlying RPC connection (gRPC or DRPC), making it easy to swap
+// them without changing the caller code.
+type RPCConn interface {
+	NewStatusClient() serverpb.StatusClient
+	NewAdminClient() serverpb.AdminClient
+	NewInitClient() serverpb.InitClient
+	NewTimeSeriesClient() tspb.TimeSeriesClient
+	NewInternalClient() kvpb.InternalClient
+	NewDistSQLClient() execinfrapb.DistSQLClient
+}
+
 // ApplicationLayerInterface defines accessors to the application
 // layer of a test server. Tests written against this interface are
 // effectively agnostic to whether they use a virtual cluster or not.
@@ -275,11 +289,11 @@ type ApplicationLayerInterface interface {
 	NewClientRPCContext(ctx context.Context, userName username.SQLUsername) *rpc.Context
 
 	// RPCClientConn opens a RPC client connection to the server.
-	RPCClientConn(t TestFataler, userName username.SQLUsername) *grpc.ClientConn
+	RPCClientConn(t TestFataler, userName username.SQLUsername) RPCConn
 
 	// RPCClientConnE is like RPCClientConn but it allows the test to check the
 	// error.
-	RPCClientConnE(userName username.SQLUsername) (*grpc.ClientConn, error)
+	RPCClientConnE(userName username.SQLUsername) (RPCConn, error)
 
 	// GetAdminClient creates a serverpb.AdminClient connection to the server.
 	// Shorthand for serverpb.AdminClient(.RPCClientConn(t, "root"))

pkg/testutils/serverutils/rpc_conn.go

@@ -3,12 +3,12 @@
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
-package srvtestutils
+package serverutils
 
 import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
-	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
 	"google.golang.org/grpc"
 )
@@ -20,7 +20,7 @@ type grpcConn struct {
 	conn *grpc.ClientConn
 }
 
-func FromGRPCConn(conn *grpc.ClientConn) serverutils.RPCConn {
+func FromGRPCConn(conn *grpc.ClientConn) RPCConn {
 	return &grpcConn{conn: conn}
 }
 
@@ -43,3 +43,7 @@ func (c *grpcConn) NewTimeSeriesClient() tspb.TimeSeriesClient {
 func (c *grpcConn) NewInternalClient() kvpb.InternalClient {
 	return kvpb.NewInternalClient(c.conn)
 }
+
+func (c *grpcConn) NewDistSQLClient() execinfrapb.DistSQLClient {
+	return execinfrapb.NewDistSQLClient(c.conn)
+}

pkg/ts/server_test.go

@@ -176,7 +176,7 @@ func TestServerQuery(t *testing.T) {
 	}
 
 	conn := s.RPCClientConn(t, username.RootUserName())
-	client := tspb.NewTimeSeriesClient(conn)
+	client := conn.NewTimeSeriesClient()
 	response, err := client.Query(context.Background(), &tspb.TimeSeriesQueryRequest{
 		StartNanos: 500 * 1e9,
 		EndNanos:   526 * 1e9,
@@ -272,7 +272,7 @@ func TestServerQueryStarvation(t *testing.T) {
 	}
 
 	conn := s.RPCClientConn(t, username.RootUserName())
-	client := tspb.NewTimeSeriesClient(conn)
+	client := conn.NewTimeSeriesClient()
 
 	queries := make([]tspb.Query, 0, seriesCount)
 	for i := 0; i < seriesCount; i++ {
@@ -442,7 +442,7 @@ func TestServerQueryTenant(t *testing.T) {
 	}
 
 	conn := s.RPCClientConn(t, username.RootUserName())
-	client := tspb.NewTimeSeriesClient(conn)
+	client := conn.NewTimeSeriesClient()
 	aggregatedResponse, err := client.Query(context.Background(), &tspb.TimeSeriesQueryRequest{
 		StartNanos: 400 * 1e9,
 		EndNanos:   500 * 1e9,
@@ -578,7 +578,7 @@ func TestServerQueryTenant(t *testing.T) {
 	capability := map[tenantcapabilitiespb.ID]string{tenantcapabilitiespb.CanViewTSDBMetrics: "true"}
 	serverutils.WaitForTenantCapabilities(t, s, tenantID, capability, "")
 	tenantConn := tenant.RPCClientConn(t, username.RootUserName())
-	tenantClient := tspb.NewTimeSeriesClient(tenantConn)
+	tenantClient := tenantConn.NewTimeSeriesClient()
 
 	tenantResponse, err := tenantClient.Query(context.Background(), &tspb.TimeSeriesQueryRequest{
 		StartNanos: 400 * 1e9,
@@ -696,7 +696,7 @@ func TestServerQueryMemoryManagement(t *testing.T) {
 	}
 
 	conn := s.RPCClientConn(t, username.RootUserName())
-	client := tspb.NewTimeSeriesClient(conn)
+	client := conn.NewTimeSeriesClient()
 
 	queries := make([]tspb.Query, 0, seriesCount)
 	for i := 0; i < seriesCount; i++ {
@@ -773,7 +773,7 @@ func TestServerDump(t *testing.T) {
 	}
 
 	conn := s.RPCClientConn(t, username.RootUserName())
-	client := tspb.NewTimeSeriesClient(conn)
+	client := conn.NewTimeSeriesClient()
 
 	dumpClient, err := client.Dump(ctx, &tspb.DumpRequest{
 		Names:      names,
@@ -869,7 +869,7 @@ func TestServerDump(t *testing.T) {
 		require.NoError(t, s.DB().Run(ctx, &b))
 
 		conn := s.RPCClientConn(t, username.RootUserName())
-		client := tspb.NewTimeSeriesClient(conn)
+		client := conn.NewTimeSeriesClient()
 
 		dumpClient, err := client.Dump(ctx, &tspb.DumpRequest{
 			Names:      names,
@@ -899,7 +899,7 @@ func BenchmarkServerQuery(b *testing.B) {
 	}
 
 	conn := s.RPCClientConn(b, username.RootUserName())
-	client := tspb.NewTimeSeriesClient(conn)
+	client := conn.NewTimeSeriesClient()
 
 	queries := make([]tspb.Query, 0, seriesCount)
 	for i := 0; i < seriesCount; i++ {