Improved usability of databricks auth login ... --configure-cluster flow by displaying cluster type and runtime version

libs/databrickscfg/uc_clusters.go

@@ -0,0 +1,147 @@
+package databrickscfg
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"regexp"
+	"strings"
+
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/fatih/color"
+	"github.com/manifoldco/promptui"
+	"golang.org/x/mod/semver"
+)
+
+var minUcRuntime = canonicalVersion("v12.0")
+
+var dbrVersionRegex = regexp.MustCompile(`^(\d+\.\d+)\.x-.*`)
+
+func canonicalVersion(v string) string {
+	return semver.Canonical("v" + strings.TrimPrefix(v, "v"))
+}
+
+func GetRuntimeVersion(cluster *compute.ClusterDetails) (string, bool) {
+	match := dbrVersionRegex.FindStringSubmatch(cluster.SparkVersion)
+	if len(match) < 1 {
+		return "", false
+	}
+	return match[1], true
+}
+
+func IsCompatibleWithUC(cluster *compute.ClusterDetails, minVersion string) bool {
+	minVersion = canonicalVersion(minVersion)
+	if semver.Compare(minUcRuntime, minVersion) >= 0 {
+		return false
+	}
+	runtimeVersion, ok := GetRuntimeVersion(cluster)
+	if !ok {
+		return false
+	}
+	clusterRuntime := canonicalVersion(runtimeVersion)
+	if semver.Compare(minVersion, clusterRuntime) >= 0 {
+		return false
+	}
+	switch cluster.DataSecurityMode {
+	case compute.DataSecurityModeUserIsolation, compute.DataSecurityModeSingleUser:
+		return true
+	default:
+		return false
+	}
+}
+
+var ErrNoCompatibleClusters = errors.New("no compatible clusters found")
+
+type compatibleCluster struct {
+	compute.ClusterDetails
+	versionName string
+}
+
+func (v compatibleCluster) Access() string {
+	switch v.DataSecurityMode {
+	case compute.DataSecurityModeUserIsolation:
+		return "Shared"
+	case compute.DataSecurityModeSingleUser:
+		return "Assigned"
+	default:
+		return "Unknown"
+	}
+}
+
+func (v compatibleCluster) Runtime() string {
+	runtime, _, _ := strings.Cut(v.versionName, " (")
+	return runtime
+}
+
+func (v compatibleCluster) State() string {
+	state := v.ClusterDetails.State
+	switch state {
+	case compute.StateRunning, compute.StateResizing:
+		return color.GreenString(state.String())
+	case compute.StateError, compute.StateTerminated, compute.StateTerminating, compute.StateUnknown:
+		return color.RedString(state.String())
+	default:
+		return color.BlueString(state.String())
+	}
+}
+
+func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceClient, minVersion string) (string, error) {
+	all, err := w.Clusters.ListAll(ctx, compute.ListClustersRequest{
+		CanUseClient: "NOTEBOOKS",
+	})
+	if err != nil {
+		return "", fmt.Errorf("list clusters: %w", err)
+	}
+	me, err := w.CurrentUser.Me(ctx)
+	if err != nil {
+		return "", fmt.Errorf("current user: %w", err)
+	}
+	versions := map[string]string{}
+	sv, err := w.Clusters.SparkVersions(ctx)
+	if err != nil {
+		return "", fmt.Errorf("list runtime versions: %w", err)
+	}
+	for _, v := range sv.Versions {
+		versions[v.Key] = v.Name
+	}
+	var compatible []compatibleCluster
+	for _, v := range all {
+		if !IsCompatibleWithUC(&v, minVersion) {
+			continue
+		}
+		if v.SingleUserName != "" && v.SingleUserName != me.UserName {
+			continue
+		}
+		compatible = append(compatible, compatibleCluster{
+			ClusterDetails: v,
+			versionName:    versions[v.SparkVersion],
+		})
+	}
+	if len(compatible) == 0 {
+		return "", ErrNoCompatibleClusters
+	}
+	if len(compatible) == 1 {
+		return compatible[0].ClusterId, nil
+	}
+	i, _, err := cmdio.RunSelect(ctx, &promptui.Select{
+		Label: "Choose compatible cluster",
+		Items: compatible,
+		Searcher: func(input string, idx int) bool {
+			lower := strings.ToLower(compatible[idx].ClusterName)
+			return strings.Contains(lower, input)
+		},
+		StartInSearchMode: true,
+		Templates: &promptui.SelectTemplates{
+			Label:    "{{.ClusterName | faint}}",
+			Active:   `{{.ClusterName | bold}} ({{.State}} {{.Access}} Runtime {{.Runtime}})`,
+			Inactive: `{{.ClusterName}} ({{.State}} {{.Access}} Runtime {{.Runtime}})`,
+			Selected: `{{ "Configured cluster" | faint }}: {{ .ClusterName | bold }} ({{.ClusterId | faint}})`,
+		},
+	})
+	if err != nil {
+		return "", err
+	}
+	return compatible[i].ClusterId, nil
+}

libs/databrickscfg/uc_clusters_test.go

@@ -0,0 +1,135 @@
+package databrickscfg
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go"
+	"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"
+)
+
+func TestIsCompatible(t *testing.T) {
+	assert.True(t, IsCompatibleWithUC(&compute.ClusterDetails{
+		SparkVersion:     "13.2.x-aarch64-scala2.12",
+		DataSecurityMode: compute.DataSecurityModeUserIsolation,
+	}, "13.0"))
+	assert.False(t, IsCompatibleWithUC(&compute.ClusterDetails{
+		SparkVersion:     "13.2.x-aarch64-scala2.12",
+		DataSecurityMode: compute.DataSecurityModeNone,
+	}, "13.0"))
+	assert.False(t, IsCompatibleWithUC(&compute.ClusterDetails{
+		SparkVersion:     "9.1.x-photon-scala2.12",
+		DataSecurityMode: compute.DataSecurityModeNone,
+	}, "13.0"))
+	assert.False(t, IsCompatibleWithUC(&compute.ClusterDetails{
+		SparkVersion:     "9.1.x-photon-scala2.12",
+		DataSecurityMode: compute.DataSecurityModeNone,
+	}, "10.0"))
+	assert.False(t, IsCompatibleWithUC(&compute.ClusterDetails{
+		SparkVersion:     "custom-9.1.x-photon-scala2.12",
+		DataSecurityMode: compute.DataSecurityModeNone,
+	}, "14.0"))
+}
+
+func TestFirstCompatibleCluster(t *testing.T) {
+	cfg, server := qa.HTTPFixtures{
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/clusters/list?can_use_client=NOTEBOOKS",
+			Response: compute.ListClustersResponse{
+				Clusters: []compute.ClusterDetails{
+					{
+						ClusterId:        "abc-id",
+						ClusterName:      "first shared",
+						DataSecurityMode: compute.DataSecurityModeUserIsolation,
+						SparkVersion:     "12.2.x-whatever",
+						State:            compute.StateRunning,
+					},
+					{
+						ClusterId:        "bcd-id",
+						ClusterName:      "second personal",
+						DataSecurityMode: compute.DataSecurityModeSingleUser,
+						SparkVersion:     "14.5.x-whatever",
+						State:            compute.StateRunning,
+						SingleUserName:   "serge",
+					},
+				},
+			},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/preview/scim/v2/Me",
+			Response: iam.User{
+				UserName: "serge",
+			},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/clusters/spark-versions",
+			Response: compute.GetSparkVersionsResponse{
+				Versions: []compute.SparkVersion{
+					{
+						Key:  "14.5.x-whatever",
+						Name: "14.5 (Awesome)",
+					},
+				},
+			},
+		},
+	}.Config(t)
+	defer server.Close()
+	w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg)))
+
+	ctx := context.Background()
+	clusterID, err := AskForClusterCompatibleWithUC(ctx, w, "13.1")
+	require.NoError(t, err)
+	assert.Equal(t, "bcd-id", clusterID)
+}
+
+func TestNoCompatibleClusters(t *testing.T) {
+	cfg, server := qa.HTTPFixtures{
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/clusters/list?can_use_client=NOTEBOOKS",
+			Response: compute.ListClustersResponse{
+				Clusters: []compute.ClusterDetails{
+					{
+						ClusterId:        "abc-id",
+						ClusterName:      "first shared",
+						DataSecurityMode: compute.DataSecurityModeUserIsolation,
+						SparkVersion:     "12.2.x-whatever",
+						State:            compute.StateRunning,
+					},
+				},
+			},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/preview/scim/v2/Me",
+			Response: iam.User{
+				UserName: "serge",
+			},
+		},
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/clusters/spark-versions",
+			Response: compute.GetSparkVersionsResponse{
+				Versions: []compute.SparkVersion{
+					{
+						Key:  "14.5.x-whatever",
+						Name: "14.5 (Awesome)",
+					},
+				},
+			},
+		},
+	}.Config(t)
+	defer server.Close()
+	w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg)))
+
+	ctx := context.Background()
+	_, err := AskForClusterCompatibleWithUC(ctx, w, "13.1")
+	assert.Equal(t, ErrNoCompatibleClusters, err)
+}

libs/databrickscfg/warehouses.go

@@ -0,0 +1,51 @@
+package databrickscfg
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/sql"
+	"github.com/fatih/color"
+)
+
+var ErrNoCompatibleWarehouses = errors.New("no compatible warehouses")
+
+func AskForWarehouse(ctx context.Context, w *databricks.WorkspaceClient, types []sql.EndpointInfoWarehouseType) (string, error) {
+	all, err := w.Warehouses.ListAll(ctx, sql.ListWarehousesRequest{})
+	if err != nil {
+		return "", fmt.Errorf("list warehouses: %w", err)
+	}
+	allowed := map[sql.EndpointInfoWarehouseType]bool{}
+	for _, v := range types {
+		allowed[v] = true
+	}
+	var lastWarehouseID string
+	names := map[string]string{}
+	for _, v := range all {
+		if !allowed[v.WarehouseType] {
+			continue
+		}
+		var state string
+		switch v.State {
+		case sql.StateRunning:
+			state = color.GreenString(v.State.String())
+		case sql.StateStopped, sql.StateDeleted, sql.StateStopping, sql.StateDeleting:
+			state = color.RedString(v.State.String())
+		default:
+			state = color.BlueString(v.State.String())
+		}
+		visibleTouser := fmt.Sprintf("%s (%s %s)", v.Name, state, v.WarehouseType)
+		names[visibleTouser] = v.Id
+		lastWarehouseID = v.Id
+	}
+	if len(names) == 0 {
+		return "", ErrNoCompatibleWarehouses
+	}
+	if len(names) == 1 {
+		return lastWarehouseID, nil
+	}
+	return cmdio.Select(ctx, names, "Choose SQL Warehouse")
+}