ssh setup: use the shared cluster picker so duplicate cluster names work

.nextchanges/cli/ssh-setup-duplicate-cluster-names.md

@@ -0,0 +1 @@
+`databricks ssh setup` no longer fails to list clusters when a workspace has several clusters sharing the same name. The picker now shows the cluster ID alongside the name so duplicates can be told apart.

experimental/ssh/internal/setup/setup.go

@@ -10,8 +10,8 @@ import (
 	"github.com/databricks/cli/experimental/ssh/internal/keys"
 	"github.com/databricks/cli/experimental/ssh/internal/sshconfig"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/databrickscfg/cfgpickers"
 	"github.com/databricks/databricks-sdk-go"
-	"github.com/databricks/databricks-sdk-go/service/compute"
 )
 
 type SetupOptions struct {
@@ -47,20 +47,13 @@ func generateHostConfig(ctx context.Context, opts SetupOptions, proxyCommand str
 var clusterSelectionPrompt = defaultClusterSelectionPrompt
 
 func defaultClusterSelectionPrompt(ctx context.Context, client *databricks.WorkspaceClient) (string, error) {
-	sp := cmdio.NewSpinner(ctx)
-	sp.Update("Loading clusters.")
-	clusters, err := client.Clusters.ClusterDetailsClusterNameToClusterIdMap(ctx, compute.ListClustersRequest{
-		FilterBy: &compute.ListClustersFilterBy{
-			ClusterSources: []compute.ClusterSource{compute.ClusterSourceApi, compute.ClusterSourceUi},
-		},
-	})
-	sp.Close()
+	// AskForCluster keys the picker by list position instead of by display name,
+	// so workspaces that have several clusters sharing a name still work. It
+	// lists with the same API/UI cluster-source filter this prompt needs, and
+	// shows the cluster ID alongside the name to disambiguate duplicates.
+	id, err := cfgpickers.AskForCluster(ctx, client)
 	if err != nil {
-		return "", fmt.Errorf("failed to load names for Clusters drop-down. Please manually specify cluster argument. Original error: %w", err)
-	}
-	id, err := cmdio.Select(ctx, clusters, "The cluster to connect to")
-	if err != nil {
-		return "", err
+		return "", fmt.Errorf("failed to select a cluster. Please manually specify cluster argument. Original error: %w", err)
 	}
 	return id, nil
 }

experimental/ssh/internal/setup/setup_test.go

@@ -12,7 +12,9 @@ import (
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/qa"
 	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -293,6 +295,81 @@ func TestSetup_PromptsForClusterWhenNotProvided(t *testing.T) {
 	assert.NotContains(t, hostConfigStr, "--cluster= ")
 }
 
+// clusterListFixtures returns the API responses the cluster picker needs to
+// list the given clusters.
+func clusterListFixtures(clusters []compute.ClusterDetails) qa.HTTPFixtures {
+	return qa.HTTPFixtures{
+		{
+			Method:   "GET",
+			Resource: "/api/2.1/clusters/list?filter_by.cluster_sources=API&filter_by.cluster_sources=UI&page_size=100",
+			Response: compute.ListClustersResponse{Clusters: clusters},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/preview/scim/v2/Me?",
+			Response: iam.User{UserName: "someone@example.com"},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.1/clusters/spark-versions",
+			Response: compute.GetSparkVersionsResponse{
+				Versions: []compute.SparkVersion{{Key: "14.5.x-scala2.12", Name: "14.5 (Scala 2.12)"}},
+			},
+		},
+	}
+}
+
+func TestDefaultClusterSelectionPrompt_DuplicateClusterNames(t *testing.T) {
+	cfg, server := clusterListFixtures([]compute.ClusterDetails{
+		{
+			ClusterId:     "cluster-a",
+			ClusterName:   "shared name",
+			ClusterSource: compute.ClusterSourceUi,
+			SparkVersion:  "14.5.x-scala2.12",
+			State:         compute.StateRunning,
+		},
+		{
+			ClusterId:     "cluster-b",
+			ClusterName:   "shared name",
+			ClusterSource: compute.ClusterSourceApi,
+			SparkVersion:  "14.5.x-scala2.12",
+			State:         compute.StateRunning,
+		},
+	}).Config(t)
+	defer server.Close()
+	w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg)))
+
+	ctx := cmdio.MockDiscard(t.Context())
+	_, err := defaultClusterSelectionPrompt(ctx, w)
+
+	// Both clusters share a display name. Listing must still succeed and reach
+	// the picker, which cannot run without a TTY -- the point is that it fails
+	// there rather than while loading the cluster list.
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "Choose compatible cluster")
+	assert.NotContains(t, err.Error(), "duplicate")
+}
+
+func TestDefaultClusterSelectionPrompt_SingleCluster(t *testing.T) {
+	cfg, server := clusterListFixtures([]compute.ClusterDetails{
+		{
+			ClusterId:     "only-cluster",
+			ClusterName:   "the only one",
+			ClusterSource: compute.ClusterSourceUi,
+			SparkVersion:  "14.5.x-scala2.12",
+			State:         compute.StateRunning,
+		},
+	}).Config(t)
+	defer server.Close()
+	w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg)))
+
+	ctx := cmdio.MockDiscard(t.Context())
+	id, err := defaultClusterSelectionPrompt(ctx, w)
+
+	require.NoError(t, err)
+	assert.Equal(t, "only-cluster", id)
+}
+
 func TestSetup_SuccessfulWithExistingConfigFile(t *testing.T) {
 	ctx := cmdio.MockDiscard(t.Context())
 	tmpDir := t.TempDir()