Merge #147693

147693: upgrade: consolidate RPC clients creation r=cthumuluru-crdb a=cthumuluru-crdb

This commit consolidates migration RPC client creation logic in the upgrade package. It is a continuation of the work done in #147606.

Epic: CRDB-48923
Informs: #147757
Release note: none

Co-authored-by: Chandra Thumuluru <[email protected]>

Author: craig[bot]

pkg/server/serverpb/BUILD.bazel

@@ -102,6 +102,7 @@ go_library(
     name = "serverpb",
     srcs = [
         "admin.go",
+        "rpc_clients.go",
         "status.go",
     ],
     embed = [":serverpb_go_proto"],
@@ -110,6 +111,7 @@ go_library(
     deps = [
         "//pkg/gossip",
         "//pkg/roachpb",
+        "//pkg/rpc/rpcbase",
         "//pkg/util/errorutil",
         "//pkg/util/metric",
         "@com_github_prometheus_client_model//go",

pkg/server/serverpb/rpc_clients.go

@@ -0,0 +1,29 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package serverpb
+
+import (
+	context "context"
+
+	roachpb "github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
+)
+
+// DialMigrationClient establishes a DRPC connection if enabled; otherwise,
+// it falls back to gRPC. The established connection is used to create a
+// MigrationClient.
+func DialMigrationClient(
+	nd rpcbase.NodeDialer, ctx context.Context, nodeID roachpb.NodeID, class rpcbase.ConnectionClass,
+) (MigrationClient, error) {
+	if !rpcbase.TODODRPC {
+		conn, err := nd.Dial(ctx, nodeID, class)
+		if err != nil {
+			return nil, err
+		}
+		return NewMigrationClient(conn), nil
+	}
+	return nil, nil
+}

pkg/upgrade/upgradecluster/cluster.go

@@ -139,11 +139,10 @@ func (c *Cluster) ForEveryNodeOrServer(
 		grp.GoCtx(func(ctx context.Context) error {
 			defer alloc.Release()
 
-			conn, err := c.c.Dialer.Dial(ctx, node.ID, rpcbase.DefaultClass)
+			client, err := serverpb.DialMigrationClient(c.c.Dialer, ctx, node.ID, rpcbase.DefaultClass)
 			if err != nil {
 				return err
 			}
-			client := serverpb.NewMigrationClient(conn)
 			return fn(ctx, client)
 		})
 	}

pkg/upgrade/upgradecluster/tenant_cluster.go

@@ -25,7 +25,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/redact"
-	"google.golang.org/grpc"
 )
 
 // TenantCluster represents the set of sql nodes running in a secondary tenant.
@@ -225,7 +224,7 @@ func (t *TenantCluster) ForEveryNodeOrServer(
 		grp.GoCtx(func(ctx context.Context) error {
 			defer alloc.Release()
 
-			var conn *grpc.ClientConn
+			var client serverpb.MigrationClient
 			retryOpts := retry.Options{
 				InitialBackoff: 1 * time.Millisecond,
 				MaxRetries:     20,
@@ -235,12 +234,11 @@ func (t *TenantCluster) ForEveryNodeOrServer(
 			// test flakes due to network issues.
 			if err := retry.WithMaxAttempts(ctx, retryOpts, retryOpts.MaxRetries+1, func() error {
 				var err error
-				conn, err = t.Dialer.Dial(ctx, roachpb.NodeID(instance.InstanceID), rpcbase.DefaultClass)
+				client, err = serverpb.DialMigrationClient(t.Dialer, ctx, roachpb.NodeID(instance.InstanceID), rpcbase.DefaultClass)
 				return err
 			}); err != nil {
 				return annotateDialError(err)
 			}
-			client := serverpb.NewMigrationClient(conn)
 			return fn(ctx, client)
 		})
 	}