From 421c8431052f5d8a278e82999da240b93ccb2dc9 Mon Sep 17 00:00:00 2001 From: Bailey Brownie Date: Sun, 18 Jan 2026 18:02:00 -0500 Subject: [PATCH 1/2] Fix ssh setup not working in Azure Government --- experimental/ssh/internal/client/releases.go | 3 +++ experimental/ssh/internal/setup/setup.go | 9 ++++++- experimental/ssh/internal/setup/setup_test.go | 26 +++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/experimental/ssh/internal/client/releases.go b/experimental/ssh/internal/client/releases.go index f147244e9e..4b7b1496ab 100644 --- a/experimental/ssh/internal/client/releases.go +++ b/experimental/ssh/internal/client/releases.go @@ -95,6 +95,9 @@ func getGithubRelease(ctx context.Context, architecture, version, releasesDir st // TODO: download and check databricks_cli__SHA256SUMS fileName := getReleaseName(architecture, version) downloadURL := fmt.Sprintf("https://github.com/databricks/cli/releases/download/v%s/%s", version, fileName) + if strings.Contains(version, "dev") { + downloadURL = fmt.Sprintf("https://github.com/databricks/cli/releases/download/snapshot/%s", fileName) + } cmdio.LogString(ctx, fmt.Sprintf("Downloading %s from %s", fileName, downloadURL)) resp, err := http.Get(downloadURL) diff --git a/experimental/ssh/internal/setup/setup.go b/experimental/ssh/internal/setup/setup.go index b68fa7eef2..38566ea9b6 100644 --- a/experimental/ssh/internal/setup/setup.go +++ b/experimental/ssh/internal/setup/setup.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "regexp" + "slices" "strconv" "strings" "time" @@ -34,12 +35,18 @@ type SetupOptions struct { Profile string } +var validDataSecurityModes = []compute.DataSecurityMode{ + compute.DataSecurityModeSingleUser, + compute.DataSecurityModeLegacySingleUser, + compute.DataSecurityModeLegacySingleUserStandard, +} + func validateClusterAccess(ctx context.Context, client *databricks.WorkspaceClient, clusterID string) error { clusterInfo, err := client.Clusters.Get(ctx, compute.GetClusterRequest{ClusterId: clusterID}) if err != nil { return fmt.Errorf("failed to get cluster information for cluster ID '%s': %w", clusterID, err) } - if clusterInfo.DataSecurityMode != compute.DataSecurityModeSingleUser { + if !slices.Contains(validDataSecurityModes, clusterInfo.DataSecurityMode) { return fmt.Errorf("cluster '%s' does not have dedicated access mode. Current access mode: %s. Please ensure the cluster is configured with dedicated access mode (single user)", clusterID, clusterInfo.DataSecurityMode) } return nil diff --git a/experimental/ssh/internal/setup/setup_test.go b/experimental/ssh/internal/setup/setup_test.go index 7c4cb20925..343d436e18 100644 --- a/experimental/ssh/internal/setup/setup_test.go +++ b/experimental/ssh/internal/setup/setup_test.go @@ -29,6 +29,32 @@ func TestValidateClusterAccess_SingleUser(t *testing.T) { assert.NoError(t, err) } +func TestValidateClusterAccess_LegacySingleUser(t *testing.T) { + ctx := cmdio.MockDiscard(context.Background()) + m := mocks.NewMockWorkspaceClient(t) + clustersAPI := m.GetMockClustersAPI() + + clustersAPI.EXPECT().Get(ctx, compute.GetClusterRequest{ClusterId: "cluster-123"}).Return(&compute.ClusterDetails{ + DataSecurityMode: compute.DataSecurityModeLegacySingleUser, + }, nil) + + err := validateClusterAccess(ctx, m.WorkspaceClient, "cluster-123") + assert.NoError(t, err) +} + +func TestValidateClusterAccess_LegacySingleUserStandard(t *testing.T) { + ctx := cmdio.MockDiscard(context.Background()) + m := mocks.NewMockWorkspaceClient(t) + clustersAPI := m.GetMockClustersAPI() + + clustersAPI.EXPECT().Get(ctx, compute.GetClusterRequest{ClusterId: "cluster-123"}).Return(&compute.ClusterDetails{ + DataSecurityMode: compute.DataSecurityModeLegacySingleUserStandard, + }, nil) + + err := validateClusterAccess(ctx, m.WorkspaceClient, "cluster-123") + assert.NoError(t, err) +} + func TestValidateClusterAccess_InvalidAccessMode(t *testing.T) { ctx := cmdio.MockDiscard(context.Background()) m := mocks.NewMockWorkspaceClient(t) From bfbbfb958f4914374d1b3526811da4da890a68b2 Mon Sep 17 00:00:00 2001 From: Bailey Brownie Date: Sun, 18 Jan 2026 19:16:14 -0500 Subject: [PATCH 2/2] Add fix to changelog --- NEXT_CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index c322ff7f67..497ad9e4f8 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -10,6 +10,9 @@ * Add support for valueFrom property (similar to app.yaml) inside Apps config field in bundle configuration ([#4297](https://github.com/databricks/cli/pull/4297)) * engine/direct: Support bind & unbind. ([#4279](https://github.com/databricks/cli/pull/4279)) +### SSH +* Fix ssh setup not working in Azure Government ([#4308](https://github.com/databricks/cli/pull/4308)) + ### Dependency updates ### API Changes