Fix ssh setup not working in Azure Government

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

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

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

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)