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 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)