diff --git a/.nextchanges/cli/ssh-setup-duplicate-cluster-names.md b/.nextchanges/cli/ssh-setup-duplicate-cluster-names.md new file mode 100644 index 00000000000..8a8ed780be3 --- /dev/null +++ b/.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. diff --git a/experimental/ssh/internal/setup/setup.go b/experimental/ssh/internal/setup/setup.go index 286ab7be7a7..338ac9fe1b5 100644 --- a/experimental/ssh/internal/setup/setup.go +++ b/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 } diff --git a/experimental/ssh/internal/setup/setup_test.go b/experimental/ssh/internal/setup/setup_test.go index 9b3267c6b62..c90f7734e42 100644 --- a/experimental/ssh/internal/setup/setup_test.go +++ b/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()