From 802cb51a96a0d5602855663f69810febf9871295 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Mon, 6 Nov 2023 21:31:59 +0100 Subject: [PATCH 1/7] Expose selectors for UC clusters and SQL warehouses --- cmd/workspace/clusters/uc.go | 147 ++++++++++++++++++++++++++++++ cmd/workspace/clusters/uc_test.go | 31 +++++++ cmd/workspace/warehouses/ask.go | 51 +++++++++++ 3 files changed, 229 insertions(+) create mode 100644 cmd/workspace/clusters/uc.go create mode 100644 cmd/workspace/clusters/uc_test.go create mode 100644 cmd/workspace/warehouses/ask.go diff --git a/cmd/workspace/clusters/uc.go b/cmd/workspace/clusters/uc.go new file mode 100644 index 0000000000..45abed65ed --- /dev/null +++ b/cmd/workspace/clusters/uc.go @@ -0,0 +1,147 @@ +package clusters + +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 IsUnityCatalogCompatible(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 AskForCompatibleCluster(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 !IsUnityCatalogCompatible(&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 +} diff --git a/cmd/workspace/clusters/uc_test.go b/cmd/workspace/clusters/uc_test.go new file mode 100644 index 0000000000..1a7e36c345 --- /dev/null +++ b/cmd/workspace/clusters/uc_test.go @@ -0,0 +1,31 @@ +package clusters + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/stretchr/testify/assert" +) + +func TestIsCompatible(t *testing.T) { + assert.True(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ + SparkVersion: "13.2.x-aarch64-scala2.12", + DataSecurityMode: compute.DataSecurityModeUserIsolation, + }, "13.0")) + assert.False(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ + SparkVersion: "13.2.x-aarch64-scala2.12", + DataSecurityMode: compute.DataSecurityModeNone, + }, "13.0")) + assert.False(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ + SparkVersion: "9.1.x-photon-scala2.12", + DataSecurityMode: compute.DataSecurityModeNone, + }, "13.0")) + assert.False(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ + SparkVersion: "9.1.x-photon-scala2.12", + DataSecurityMode: compute.DataSecurityModeNone, + }, "10.0")) + assert.False(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ + SparkVersion: "custom-9.1.x-photon-scala2.12", + DataSecurityMode: compute.DataSecurityModeNone, + }, "14.0")) +} diff --git a/cmd/workspace/warehouses/ask.go b/cmd/workspace/warehouses/ask.go new file mode 100644 index 0000000000..72963187a2 --- /dev/null +++ b/cmd/workspace/warehouses/ask.go @@ -0,0 +1,51 @@ +package warehouses + +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 AskForCompatibleWarehouses(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") +} From aef6541df3ba6490ab47b0422de08ad97c8ee796 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Wed, 8 Nov 2023 17:23:20 +0100 Subject: [PATCH 2/7] .. --- cmd/workspace/clusters/uc_test.go | 31 ---- .../databrickscfg/uc_clusters.go | 8 +- libs/databrickscfg/uc_clusters_test.go | 135 ++++++++++++++++++ .../databrickscfg/warehouses.go | 4 +- libs/databrickscfg/warehouses_test.go | 70 +++++++++ 5 files changed, 211 insertions(+), 37 deletions(-) delete mode 100644 cmd/workspace/clusters/uc_test.go rename cmd/workspace/clusters/uc.go => libs/databrickscfg/uc_clusters.go (93%) create mode 100644 libs/databrickscfg/uc_clusters_test.go rename cmd/workspace/warehouses/ask.go => libs/databrickscfg/warehouses.go (88%) create mode 100644 libs/databrickscfg/warehouses_test.go diff --git a/cmd/workspace/clusters/uc_test.go b/cmd/workspace/clusters/uc_test.go deleted file mode 100644 index 1a7e36c345..0000000000 --- a/cmd/workspace/clusters/uc_test.go +++ /dev/null @@ -1,31 +0,0 @@ -package clusters - -import ( - "testing" - - "github.com/databricks/databricks-sdk-go/service/compute" - "github.com/stretchr/testify/assert" -) - -func TestIsCompatible(t *testing.T) { - assert.True(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ - SparkVersion: "13.2.x-aarch64-scala2.12", - DataSecurityMode: compute.DataSecurityModeUserIsolation, - }, "13.0")) - assert.False(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ - SparkVersion: "13.2.x-aarch64-scala2.12", - DataSecurityMode: compute.DataSecurityModeNone, - }, "13.0")) - assert.False(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ - SparkVersion: "9.1.x-photon-scala2.12", - DataSecurityMode: compute.DataSecurityModeNone, - }, "13.0")) - assert.False(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ - SparkVersion: "9.1.x-photon-scala2.12", - DataSecurityMode: compute.DataSecurityModeNone, - }, "10.0")) - assert.False(t, IsUnityCatalogCompatible(&compute.ClusterDetails{ - SparkVersion: "custom-9.1.x-photon-scala2.12", - DataSecurityMode: compute.DataSecurityModeNone, - }, "14.0")) -} diff --git a/cmd/workspace/clusters/uc.go b/libs/databrickscfg/uc_clusters.go similarity index 93% rename from cmd/workspace/clusters/uc.go rename to libs/databrickscfg/uc_clusters.go index 45abed65ed..0bb67ebd23 100644 --- a/cmd/workspace/clusters/uc.go +++ b/libs/databrickscfg/uc_clusters.go @@ -1,4 +1,4 @@ -package clusters +package databrickscfg import ( "context" @@ -31,7 +31,7 @@ func GetRuntimeVersion(cluster *compute.ClusterDetails) (string, bool) { return match[1], true } -func IsUnityCatalogCompatible(cluster *compute.ClusterDetails, minVersion string) bool { +func IsCompatibleWithUC(cluster *compute.ClusterDetails, minVersion string) bool { minVersion = canonicalVersion(minVersion) if semver.Compare(minUcRuntime, minVersion) >= 0 { return false @@ -87,7 +87,7 @@ func (v compatibleCluster) State() string { } } -func AskForCompatibleCluster(ctx context.Context, w *databricks.WorkspaceClient, minVersion string) (string, error) { +func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceClient, minVersion string) (string, error) { all, err := w.Clusters.ListAll(ctx, compute.ListClustersRequest{ CanUseClient: "NOTEBOOKS", }) @@ -108,7 +108,7 @@ func AskForCompatibleCluster(ctx context.Context, w *databricks.WorkspaceClient, } var compatible []compatibleCluster for _, v := range all { - if !IsUnityCatalogCompatible(&v, minVersion) { + if !IsCompatibleWithUC(&v, minVersion) { continue } if v.SingleUserName != "" && v.SingleUserName != me.UserName { diff --git a/libs/databrickscfg/uc_clusters_test.go b/libs/databrickscfg/uc_clusters_test.go new file mode 100644 index 0000000000..3226d81016 --- /dev/null +++ b/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) +} diff --git a/cmd/workspace/warehouses/ask.go b/libs/databrickscfg/warehouses.go similarity index 88% rename from cmd/workspace/warehouses/ask.go rename to libs/databrickscfg/warehouses.go index 72963187a2..4feab0de4a 100644 --- a/cmd/workspace/warehouses/ask.go +++ b/libs/databrickscfg/warehouses.go @@ -1,4 +1,4 @@ -package warehouses +package databrickscfg import ( "context" @@ -13,7 +13,7 @@ import ( var ErrNoCompatibleWarehouses = errors.New("no compatible warehouses") -func AskForCompatibleWarehouses(ctx context.Context, w *databricks.WorkspaceClient, types []sql.EndpointInfoWarehouseType) (string, error) { +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) diff --git a/libs/databrickscfg/warehouses_test.go b/libs/databrickscfg/warehouses_test.go new file mode 100644 index 0000000000..a37d04654a --- /dev/null +++ b/libs/databrickscfg/warehouses_test.go @@ -0,0 +1,70 @@ +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/sql" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestFirstCompatibleWarehouse(t *testing.T) { + cfg, server := qa.HTTPFixtures{ + { + Method: "GET", + Resource: "/api/2.0/sql/warehouses?", + Response: sql.ListWarehousesResponse{ + Warehouses: []sql.EndpointInfo{ + { + Id: "efg-id", + Name: "First PRO Warehouse", + WarehouseType: sql.EndpointInfoWarehouseTypePro, + }, + { + Id: "ghe-id", + Name: "Second UNKNOWN Warehouse", + WarehouseType: sql.EndpointInfoWarehouseTypeTypeUnspecified, + }, + }, + }, + }, + }.Config(t) + defer server.Close() + w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg))) + + ctx := context.Background() + clusterID, err := AskForWarehouse(ctx, w, []sql.EndpointInfoWarehouseType{ + sql.EndpointInfoWarehouseTypePro, + }) + require.NoError(t, err) + assert.Equal(t, "efg-id", clusterID) +} + +func TestNoCompatibleWarehouses(t *testing.T) { + cfg, server := qa.HTTPFixtures{ + { + Method: "GET", + Resource: "/api/2.0/sql/warehouses?", + Response: sql.ListWarehousesResponse{ + Warehouses: []sql.EndpointInfo{ + { + Id: "efg-id", + Name: "...", + WarehouseType: sql.EndpointInfoWarehouseTypeClassic, + }, + }, + }, + }, + }.Config(t) + defer server.Close() + w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg))) + + ctx := context.Background() + _, err := AskForWarehouse(ctx, w, []sql.EndpointInfoWarehouseType{ + sql.EndpointInfoWarehouseTypePro, + }) + assert.Equal(t, ErrNoCompatibleWarehouses, err) +} From ec753640edd9c9726244e292cfba3fc1e74bbe53 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Thu, 9 Nov 2023 13:14:39 +0100 Subject: [PATCH 3/7] .. --- cmd/auth/login.go | 15 ++++----------- libs/databrickscfg/profiles.go | 16 ++++++++++------ libs/databrickscfg/uc_clusters.go | 9 +++++++++ 3 files changed, 23 insertions(+), 17 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 3a3f3a6dcf..67d21c00ca 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -10,7 +10,6 @@ import ( "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/service/compute" "github.com/spf13/cobra" ) @@ -28,6 +27,8 @@ func configureHost(ctx context.Context, persistentAuth *auth.PersistentAuth, arg return nil } +const minimalDbConnectVersion = "13.1" + func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { cmd := &cobra.Command{ Use: "login [HOST]", @@ -95,19 +96,11 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { return err } ctx := cmd.Context() - - promptSpinner := cmdio.Spinner(ctx) - promptSpinner <- "Loading list of clusters to select from" - names, err := w.Clusters.ClusterDetailsClusterNameToClusterIdMap(ctx, compute.ListClustersRequest{}) - close(promptSpinner) - if err != nil { - return fmt.Errorf("failed to load clusters list. Original error: %w", err) - } - clusterId, err := cmdio.Select(ctx, names, "Choose cluster") + clusterID, err := databrickscfg.AskForClusterCompatibleWithUC(ctx, w, minimalDbConnectVersion) if err != nil { return err } - cfg.ClusterID = clusterId + cfg.ClusterID = clusterID } if profileName != "" { diff --git a/libs/databrickscfg/profiles.go b/libs/databrickscfg/profiles.go index 864000d034..677fa600e1 100644 --- a/libs/databrickscfg/profiles.go +++ b/libs/databrickscfg/profiles.go @@ -14,9 +14,11 @@ import ( // It should only be used for prompting and filtering. // Use its name to construct a config.Config. type Profile struct { - Name string - Host string - AccountID string + Name string + Host string + AccountID string + ClusterID string + WarehouseID string } func (p Profile) Cloud() string { @@ -117,9 +119,11 @@ func LoadProfiles(fn ProfileMatchFunction) (file string, profiles Profiles, err continue } profile := Profile{ - Name: v.Name(), - Host: host, - AccountID: all["account_id"], + Name: v.Name(), + Host: host, + AccountID: all["account_id"], + ClusterID: all["cluster_id"], + WarehouseID: all["warehouse_id"], } if fn(profile) { profiles = append(profiles, profile) diff --git a/libs/databrickscfg/uc_clusters.go b/libs/databrickscfg/uc_clusters.go index 0bb67ebd23..74d2f46986 100644 --- a/libs/databrickscfg/uc_clusters.go +++ b/libs/databrickscfg/uc_clusters.go @@ -111,6 +111,15 @@ func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceC if !IsCompatibleWithUC(&v, minVersion) { continue } + switch v.ClusterSource { + case compute.ClusterSourceJob, + compute.ClusterSourceModels, + compute.ClusterSourcePipeline, + compute.ClusterSourcePipelineMaintenance, + compute.ClusterSourceSql: + // only UI and API clusters are usable for DBConnect + continue + } if v.SingleUserName != "" && v.SingleUserName != me.UserName { continue } From 12cc8f0b31619c555576454d609c32289e88a352 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Thu, 9 Nov 2023 13:41:05 +0100 Subject: [PATCH 4/7] added support for snapshots --- libs/databrickscfg/uc_clusters.go | 32 ++++++++++++++++++++------ libs/databrickscfg/uc_clusters_test.go | 12 ++++++++++ 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/libs/databrickscfg/uc_clusters.go b/libs/databrickscfg/uc_clusters.go index 74d2f46986..a86ca5e891 100644 --- a/libs/databrickscfg/uc_clusters.go +++ b/libs/databrickscfg/uc_clusters.go @@ -18,6 +18,7 @@ import ( var minUcRuntime = canonicalVersion("v12.0") var dbrVersionRegex = regexp.MustCompile(`^(\d+\.\d+)\.x-.*`) +var dbrSnapshotVersionRegex = regexp.MustCompile(`^(\d+)\.x-snapshot.*`) func canonicalVersion(v string) string { return semver.Canonical("v" + strings.TrimPrefix(v, "v")) @@ -26,6 +27,11 @@ func canonicalVersion(v string) string { func GetRuntimeVersion(cluster *compute.ClusterDetails) (string, bool) { match := dbrVersionRegex.FindStringSubmatch(cluster.SparkVersion) if len(match) < 1 { + match = dbrSnapshotVersionRegex.FindStringSubmatch(cluster.SparkVersion) + if len(match) > 1 { + // we return 14.0 for 14.x-snapshot for semver.Compare() to work + return fmt.Sprintf("%s.0", match[1]), true + } return "", false } return match[1], true @@ -41,7 +47,7 @@ func IsCompatibleWithUC(cluster *compute.ClusterDetails, minVersion string) bool return false } clusterRuntime := canonicalVersion(runtimeVersion) - if semver.Compare(minVersion, clusterRuntime) >= 0 { + if semver.Compare(minVersion, clusterRuntime) > 0 { return false } switch cluster.DataSecurityMode { @@ -87,21 +93,24 @@ func (v compatibleCluster) State() string { } } -func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceClient, minVersion string) (string, error) { +func loadClustersCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceClient, minVersion string) ([]compatibleCluster, error) { + promptSpinner := cmdio.Spinner(ctx) + promptSpinner <- "Loading list of clusters to select from" + defer close(promptSpinner) all, err := w.Clusters.ListAll(ctx, compute.ListClustersRequest{ CanUseClient: "NOTEBOOKS", }) if err != nil { - return "", fmt.Errorf("list clusters: %w", err) + return nil, fmt.Errorf("list clusters: %w", err) } me, err := w.CurrentUser.Me(ctx) if err != nil { - return "", fmt.Errorf("current user: %w", err) + return nil, 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) + return nil, fmt.Errorf("list runtime versions: %w", err) } for _, v := range sv.Versions { versions[v.Key] = v.Name @@ -117,7 +126,8 @@ func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceC compute.ClusterSourcePipeline, compute.ClusterSourcePipelineMaintenance, compute.ClusterSourceSql: - // only UI and API clusters are usable for DBConnect + // only UI and API clusters are usable for DBConnect. + // `CanUseClient: "NOTEBOOKS"`` didn't seem to have an effect. continue } if v.SingleUserName != "" && v.SingleUserName != me.UserName { @@ -128,6 +138,14 @@ func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceC versionName: versions[v.SparkVersion], }) } + return compatible, nil +} + +func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceClient, minVersion string) (string, error) { + compatible, err := loadClustersCompatibleWithUC(ctx, w, minVersion) + if err != nil { + return "", fmt.Errorf("load: %w", err) + } if len(compatible) == 0 { return "", ErrNoCompatibleClusters } @@ -144,7 +162,7 @@ func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceC StartInSearchMode: true, Templates: &promptui.SelectTemplates{ Label: "{{.ClusterName | faint}}", - Active: `{{.ClusterName | bold}} ({{.State}} {{.Access}} Runtime {{.Runtime}})`, + Active: `{{.ClusterName | bold}} ({{.State}} {{.Access}} Runtime {{.Runtime}}) ({{.ClusterId | faint}})`, Inactive: `{{.ClusterName}} ({{.State}} {{.Access}} Runtime {{.Runtime}})`, Selected: `{{ "Configured cluster" | faint }}: {{ .ClusterName | bold }} ({{.ClusterId | faint}})`, }, diff --git a/libs/databrickscfg/uc_clusters_test.go b/libs/databrickscfg/uc_clusters_test.go index 3226d81016..b508036364 100644 --- a/libs/databrickscfg/uc_clusters_test.go +++ b/libs/databrickscfg/uc_clusters_test.go @@ -1,9 +1,12 @@ package databrickscfg import ( + "bytes" "context" "testing" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/qa" "github.com/databricks/databricks-sdk-go/service/compute" @@ -35,6 +38,13 @@ func TestIsCompatible(t *testing.T) { }, "14.0")) } +func TestIsCompatibleWithSnapshots(t *testing.T) { + assert.True(t, IsCompatibleWithUC(&compute.ClusterDetails{ + SparkVersion: "14.x-snapshot-cpu-ml-scala2.12", + DataSecurityMode: compute.DataSecurityModeUserIsolation, + }, "14.0")) +} + func TestFirstCompatibleCluster(t *testing.T) { cfg, server := qa.HTTPFixtures{ { @@ -84,6 +94,7 @@ func TestFirstCompatibleCluster(t *testing.T) { w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg))) ctx := context.Background() + ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "...")) clusterID, err := AskForClusterCompatibleWithUC(ctx, w, "13.1") require.NoError(t, err) assert.Equal(t, "bcd-id", clusterID) @@ -130,6 +141,7 @@ func TestNoCompatibleClusters(t *testing.T) { w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg))) ctx := context.Background() + ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "...")) _, err := AskForClusterCompatibleWithUC(ctx, w, "13.1") assert.Equal(t, ErrNoCompatibleClusters, err) } From 9c01b2d7e570410bcc0290b5093c218713959101 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Thu, 9 Nov 2023 15:16:06 +0100 Subject: [PATCH 5/7] added cfg pickers --- cmd/auth/login.go | 4 ++- .../{ => cfgpickers}/uc_clusters.go | 28 ++++++++++++++----- .../{ => cfgpickers}/uc_clusters_test.go | 23 ++++++++------- .../{ => cfgpickers}/warehouses.go | 28 ++++++++++++++----- .../{ => cfgpickers}/warehouses_test.go | 10 ++----- 5 files changed, 59 insertions(+), 34 deletions(-) rename libs/databrickscfg/{ => cfgpickers}/uc_clusters.go (83%) rename libs/databrickscfg/{ => cfgpickers}/uc_clusters_test.go (85%) rename libs/databrickscfg/{ => cfgpickers}/warehouses.go (76%) rename libs/databrickscfg/{ => cfgpickers}/warehouses_test.go (86%) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 67d21c00ca..911d739e3b 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -8,6 +8,7 @@ import ( "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" + "github.com/databricks/cli/libs/databrickscfg/cfgpickers" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" "github.com/spf13/cobra" @@ -96,7 +97,8 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { return err } ctx := cmd.Context() - clusterID, err := databrickscfg.AskForClusterCompatibleWithUC(ctx, w, minimalDbConnectVersion) + clusterID, err := cfgpickers.AskForInteractiveCluster(ctx, w, + cfgpickers.WithDatabricksConnect(minimalDbConnectVersion)) if err != nil { return err } diff --git a/libs/databrickscfg/uc_clusters.go b/libs/databrickscfg/cfgpickers/uc_clusters.go similarity index 83% rename from libs/databrickscfg/uc_clusters.go rename to libs/databrickscfg/cfgpickers/uc_clusters.go index a86ca5e891..060ef0fbc5 100644 --- a/libs/databrickscfg/uc_clusters.go +++ b/libs/databrickscfg/cfgpickers/uc_clusters.go @@ -1,4 +1,4 @@ -package databrickscfg +package cfgpickers import ( "context" @@ -24,7 +24,7 @@ func canonicalVersion(v string) string { return semver.Canonical("v" + strings.TrimPrefix(v, "v")) } -func GetRuntimeVersion(cluster *compute.ClusterDetails) (string, bool) { +func GetRuntimeVersion(cluster compute.ClusterDetails) (string, bool) { match := dbrVersionRegex.FindStringSubmatch(cluster.SparkVersion) if len(match) < 1 { match = dbrSnapshotVersionRegex.FindStringSubmatch(cluster.SparkVersion) @@ -37,7 +37,7 @@ func GetRuntimeVersion(cluster *compute.ClusterDetails) (string, bool) { return match[1], true } -func IsCompatibleWithUC(cluster *compute.ClusterDetails, minVersion string) bool { +func IsCompatibleWithUC(cluster compute.ClusterDetails, minVersion string) bool { minVersion = canonicalVersion(minVersion) if semver.Compare(minUcRuntime, minVersion) >= 0 { return false @@ -93,7 +93,15 @@ func (v compatibleCluster) State() string { } } -func loadClustersCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceClient, minVersion string) ([]compatibleCluster, error) { +type clusterFilter func(cluster compute.ClusterDetails) bool + +func WithDatabricksConnect(minVersion string) func(cluster compute.ClusterDetails) bool { + return func(cluster compute.ClusterDetails) bool { + return IsCompatibleWithUC(cluster, minVersion) + } +} + +func loadInteractiveClusters(ctx context.Context, w *databricks.WorkspaceClient, filters []clusterFilter) ([]compatibleCluster, error) { promptSpinner := cmdio.Spinner(ctx) promptSpinner <- "Loading list of clusters to select from" defer close(promptSpinner) @@ -117,7 +125,13 @@ func loadClustersCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceCl } var compatible []compatibleCluster for _, v := range all { - if !IsCompatibleWithUC(&v, minVersion) { + var skip bool + for _, filter := range filters { + if !filter(v) { + skip = true + } + } + if skip { continue } switch v.ClusterSource { @@ -141,8 +155,8 @@ func loadClustersCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceCl return compatible, nil } -func AskForClusterCompatibleWithUC(ctx context.Context, w *databricks.WorkspaceClient, minVersion string) (string, error) { - compatible, err := loadClustersCompatibleWithUC(ctx, w, minVersion) +func AskForInteractiveCluster(ctx context.Context, w *databricks.WorkspaceClient, filters ...clusterFilter) (string, error) { + compatible, err := loadInteractiveClusters(ctx, w, filters) if err != nil { return "", fmt.Errorf("load: %w", err) } diff --git a/libs/databrickscfg/uc_clusters_test.go b/libs/databrickscfg/cfgpickers/uc_clusters_test.go similarity index 85% rename from libs/databrickscfg/uc_clusters_test.go rename to libs/databrickscfg/cfgpickers/uc_clusters_test.go index b508036364..f61e317abd 100644 --- a/libs/databrickscfg/uc_clusters_test.go +++ b/libs/databrickscfg/cfgpickers/uc_clusters_test.go @@ -1,4 +1,4 @@ -package databrickscfg +package cfgpickers import ( "bytes" @@ -11,35 +11,34 @@ import ( "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{ + require.True(t, IsCompatibleWithUC(compute.ClusterDetails{ SparkVersion: "13.2.x-aarch64-scala2.12", DataSecurityMode: compute.DataSecurityModeUserIsolation, }, "13.0")) - assert.False(t, IsCompatibleWithUC(&compute.ClusterDetails{ + require.False(t, IsCompatibleWithUC(compute.ClusterDetails{ SparkVersion: "13.2.x-aarch64-scala2.12", DataSecurityMode: compute.DataSecurityModeNone, }, "13.0")) - assert.False(t, IsCompatibleWithUC(&compute.ClusterDetails{ + require.False(t, IsCompatibleWithUC(compute.ClusterDetails{ SparkVersion: "9.1.x-photon-scala2.12", DataSecurityMode: compute.DataSecurityModeNone, }, "13.0")) - assert.False(t, IsCompatibleWithUC(&compute.ClusterDetails{ + require.False(t, IsCompatibleWithUC(compute.ClusterDetails{ SparkVersion: "9.1.x-photon-scala2.12", DataSecurityMode: compute.DataSecurityModeNone, }, "10.0")) - assert.False(t, IsCompatibleWithUC(&compute.ClusterDetails{ + require.False(t, IsCompatibleWithUC(compute.ClusterDetails{ SparkVersion: "custom-9.1.x-photon-scala2.12", DataSecurityMode: compute.DataSecurityModeNone, }, "14.0")) } func TestIsCompatibleWithSnapshots(t *testing.T) { - assert.True(t, IsCompatibleWithUC(&compute.ClusterDetails{ + require.True(t, IsCompatibleWithUC(compute.ClusterDetails{ SparkVersion: "14.x-snapshot-cpu-ml-scala2.12", DataSecurityMode: compute.DataSecurityModeUserIsolation, }, "14.0")) @@ -95,9 +94,9 @@ func TestFirstCompatibleCluster(t *testing.T) { ctx := context.Background() ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "...")) - clusterID, err := AskForClusterCompatibleWithUC(ctx, w, "13.1") + clusterID, err := AskForInteractiveCluster(ctx, w, WithDatabricksConnect("13.1")) require.NoError(t, err) - assert.Equal(t, "bcd-id", clusterID) + require.Equal(t, "bcd-id", clusterID) } func TestNoCompatibleClusters(t *testing.T) { @@ -142,6 +141,6 @@ func TestNoCompatibleClusters(t *testing.T) { ctx := context.Background() ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "...")) - _, err := AskForClusterCompatibleWithUC(ctx, w, "13.1") - assert.Equal(t, ErrNoCompatibleClusters, err) + _, err := AskForInteractiveCluster(ctx, w, WithDatabricksConnect("13.1")) + require.Equal(t, ErrNoCompatibleClusters, err) } diff --git a/libs/databrickscfg/warehouses.go b/libs/databrickscfg/cfgpickers/warehouses.go similarity index 76% rename from libs/databrickscfg/warehouses.go rename to libs/databrickscfg/cfgpickers/warehouses.go index 4feab0de4a..d0f0e84a96 100644 --- a/libs/databrickscfg/warehouses.go +++ b/libs/databrickscfg/cfgpickers/warehouses.go @@ -1,4 +1,4 @@ -package databrickscfg +package cfgpickers import ( "context" @@ -13,19 +13,33 @@ import ( 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) - } +type warehouseFilter func(sql.EndpointInfo) bool + +func WithWarehouseTypes(types ...sql.EndpointInfoWarehouseType) func(sql.EndpointInfo) bool { allowed := map[sql.EndpointInfoWarehouseType]bool{} for _, v := range types { allowed[v] = true } + return func(ei sql.EndpointInfo) bool { + return allowed[ei.WarehouseType] + } +} + +func AskForWarehouse(ctx context.Context, w *databricks.WorkspaceClient, filters ...warehouseFilter) (string, error) { + all, err := w.Warehouses.ListAll(ctx, sql.ListWarehousesRequest{}) + if err != nil { + return "", fmt.Errorf("list warehouses: %w", err) + } var lastWarehouseID string names := map[string]string{} for _, v := range all { - if !allowed[v.WarehouseType] { + var skip bool + for _, filter := range filters { + if !filter(v) { + skip = true + } + } + if skip { continue } var state string diff --git a/libs/databrickscfg/warehouses_test.go b/libs/databrickscfg/cfgpickers/warehouses_test.go similarity index 86% rename from libs/databrickscfg/warehouses_test.go rename to libs/databrickscfg/cfgpickers/warehouses_test.go index a37d04654a..d6030b4959 100644 --- a/libs/databrickscfg/warehouses_test.go +++ b/libs/databrickscfg/cfgpickers/warehouses_test.go @@ -1,4 +1,4 @@ -package databrickscfg +package cfgpickers import ( "context" @@ -36,9 +36,7 @@ func TestFirstCompatibleWarehouse(t *testing.T) { w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg))) ctx := context.Background() - clusterID, err := AskForWarehouse(ctx, w, []sql.EndpointInfoWarehouseType{ - sql.EndpointInfoWarehouseTypePro, - }) + clusterID, err := AskForWarehouse(ctx, w, WithWarehouseTypes(sql.EndpointInfoWarehouseTypePro)) require.NoError(t, err) assert.Equal(t, "efg-id", clusterID) } @@ -63,8 +61,6 @@ func TestNoCompatibleWarehouses(t *testing.T) { w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg))) ctx := context.Background() - _, err := AskForWarehouse(ctx, w, []sql.EndpointInfoWarehouseType{ - sql.EndpointInfoWarehouseTypePro, - }) + _, err := AskForWarehouse(ctx, w, WithWarehouseTypes(sql.EndpointInfoWarehouseTypePro)) assert.Equal(t, ErrNoCompatibleWarehouses, err) } From 09b142ebe30780c3923de5d6664ba72cb8788a59 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Thu, 9 Nov 2023 15:22:32 +0100 Subject: [PATCH 6/7] .. --- cmd/auth/login.go | 2 +- .../{uc_clusters.go => clusters.go} | 48 ++++++++++--------- .../{uc_clusters_test.go => clusters_test.go} | 4 +- libs/databrickscfg/cfgpickers/warehouses.go | 18 +++---- 4 files changed, 38 insertions(+), 34 deletions(-) rename libs/databrickscfg/cfgpickers/{uc_clusters.go => clusters.go} (86%) rename libs/databrickscfg/cfgpickers/{uc_clusters_test.go => clusters_test.go} (96%) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 911d739e3b..db7bf93c53 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -97,7 +97,7 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { return err } ctx := cmd.Context() - clusterID, err := cfgpickers.AskForInteractiveCluster(ctx, w, + clusterID, err := cfgpickers.AskForCluster(ctx, w, cfgpickers.WithDatabricksConnect(minimalDbConnectVersion)) if err != nil { return err diff --git a/libs/databrickscfg/cfgpickers/uc_clusters.go b/libs/databrickscfg/cfgpickers/clusters.go similarity index 86% rename from libs/databrickscfg/cfgpickers/uc_clusters.go rename to libs/databrickscfg/cfgpickers/clusters.go index 060ef0fbc5..20239b0c55 100644 --- a/libs/databrickscfg/cfgpickers/uc_clusters.go +++ b/libs/databrickscfg/cfgpickers/clusters.go @@ -10,6 +10,7 @@ import ( "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/fatih/color" "github.com/manifoldco/promptui" "golang.org/x/mod/semver" @@ -93,11 +94,27 @@ func (v compatibleCluster) State() string { } } -type clusterFilter func(cluster compute.ClusterDetails) bool +type clusterFilter func(cluster *compute.ClusterDetails, me *iam.User) bool -func WithDatabricksConnect(minVersion string) func(cluster compute.ClusterDetails) bool { - return func(cluster compute.ClusterDetails) bool { - return IsCompatibleWithUC(cluster, minVersion) +func WithDatabricksConnect(minVersion string) func(*compute.ClusterDetails, *iam.User) bool { + return func(cluster *compute.ClusterDetails, me *iam.User) bool { + if !IsCompatibleWithUC(*cluster, minVersion) { + return false + } + switch cluster.ClusterSource { + case compute.ClusterSourceJob, + compute.ClusterSourceModels, + compute.ClusterSourcePipeline, + compute.ClusterSourcePipelineMaintenance, + compute.ClusterSourceSql: + // only UI and API clusters are usable for DBConnect. + // `CanUseClient: "NOTEBOOKS"`` didn't seem to have an effect. + return false + } + if cluster.SingleUserName != "" && cluster.SingleUserName != me.UserName { + return false + } + return true } } @@ -124,38 +141,25 @@ func loadInteractiveClusters(ctx context.Context, w *databricks.WorkspaceClient, versions[v.Key] = v.Name } var compatible []compatibleCluster - for _, v := range all { + for _, cluster := range all { var skip bool for _, filter := range filters { - if !filter(v) { + if !filter(&cluster, me) { skip = true } } if skip { continue } - switch v.ClusterSource { - case compute.ClusterSourceJob, - compute.ClusterSourceModels, - compute.ClusterSourcePipeline, - compute.ClusterSourcePipelineMaintenance, - compute.ClusterSourceSql: - // only UI and API clusters are usable for DBConnect. - // `CanUseClient: "NOTEBOOKS"`` didn't seem to have an effect. - continue - } - if v.SingleUserName != "" && v.SingleUserName != me.UserName { - continue - } compatible = append(compatible, compatibleCluster{ - ClusterDetails: v, - versionName: versions[v.SparkVersion], + ClusterDetails: cluster, + versionName: versions[cluster.SparkVersion], }) } return compatible, nil } -func AskForInteractiveCluster(ctx context.Context, w *databricks.WorkspaceClient, filters ...clusterFilter) (string, error) { +func AskForCluster(ctx context.Context, w *databricks.WorkspaceClient, filters ...clusterFilter) (string, error) { compatible, err := loadInteractiveClusters(ctx, w, filters) if err != nil { return "", fmt.Errorf("load: %w", err) diff --git a/libs/databrickscfg/cfgpickers/uc_clusters_test.go b/libs/databrickscfg/cfgpickers/clusters_test.go similarity index 96% rename from libs/databrickscfg/cfgpickers/uc_clusters_test.go rename to libs/databrickscfg/cfgpickers/clusters_test.go index f61e317abd..362d6904f5 100644 --- a/libs/databrickscfg/cfgpickers/uc_clusters_test.go +++ b/libs/databrickscfg/cfgpickers/clusters_test.go @@ -94,7 +94,7 @@ func TestFirstCompatibleCluster(t *testing.T) { ctx := context.Background() ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "...")) - clusterID, err := AskForInteractiveCluster(ctx, w, WithDatabricksConnect("13.1")) + clusterID, err := AskForCluster(ctx, w, WithDatabricksConnect("13.1")) require.NoError(t, err) require.Equal(t, "bcd-id", clusterID) } @@ -141,6 +141,6 @@ func TestNoCompatibleClusters(t *testing.T) { ctx := context.Background() ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "...")) - _, err := AskForInteractiveCluster(ctx, w, WithDatabricksConnect("13.1")) + _, err := AskForCluster(ctx, w, WithDatabricksConnect("13.1")) require.Equal(t, ErrNoCompatibleClusters, err) } diff --git a/libs/databrickscfg/cfgpickers/warehouses.go b/libs/databrickscfg/cfgpickers/warehouses.go index d0f0e84a96..65b5f8c83c 100644 --- a/libs/databrickscfg/cfgpickers/warehouses.go +++ b/libs/databrickscfg/cfgpickers/warehouses.go @@ -32,10 +32,10 @@ func AskForWarehouse(ctx context.Context, w *databricks.WorkspaceClient, filters } var lastWarehouseID string names := map[string]string{} - for _, v := range all { + for _, warehouse := range all { var skip bool for _, filter := range filters { - if !filter(v) { + if !filter(warehouse) { skip = true } } @@ -43,17 +43,17 @@ func AskForWarehouse(ctx context.Context, w *databricks.WorkspaceClient, filters continue } var state string - switch v.State { + switch warehouse.State { case sql.StateRunning: - state = color.GreenString(v.State.String()) + state = color.GreenString(warehouse.State.String()) case sql.StateStopped, sql.StateDeleted, sql.StateStopping, sql.StateDeleting: - state = color.RedString(v.State.String()) + state = color.RedString(warehouse.State.String()) default: - state = color.BlueString(v.State.String()) + state = color.BlueString(warehouse.State.String()) } - visibleTouser := fmt.Sprintf("%s (%s %s)", v.Name, state, v.WarehouseType) - names[visibleTouser] = v.Id - lastWarehouseID = v.Id + visibleTouser := fmt.Sprintf("%s (%s %s)", warehouse.Name, state, warehouse.WarehouseType) + names[visibleTouser] = warehouse.Id + lastWarehouseID = warehouse.Id } if len(names) == 0 { return "", ErrNoCompatibleWarehouses From 0baccd7e14f2fa70bfb48aa80e956c4e9478c3f0 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Thu, 9 Nov 2023 17:32:06 +0100 Subject: [PATCH 7/7] .. --- libs/databrickscfg/cfgpickers/clusters.go | 4 ++-- libs/databrickscfg/profiles.go | 16 ++++++---------- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/libs/databrickscfg/cfgpickers/clusters.go b/libs/databrickscfg/cfgpickers/clusters.go index 20239b0c55..ac037698e4 100644 --- a/libs/databrickscfg/cfgpickers/clusters.go +++ b/libs/databrickscfg/cfgpickers/clusters.go @@ -30,8 +30,8 @@ func GetRuntimeVersion(cluster compute.ClusterDetails) (string, bool) { if len(match) < 1 { match = dbrSnapshotVersionRegex.FindStringSubmatch(cluster.SparkVersion) if len(match) > 1 { - // we return 14.0 for 14.x-snapshot for semver.Compare() to work - return fmt.Sprintf("%s.0", match[1]), true + // we return 14.999 for 14.x-snapshot for semver.Compare() to work properly + return fmt.Sprintf("%s.999", match[1]), true } return "", false } diff --git a/libs/databrickscfg/profiles.go b/libs/databrickscfg/profiles.go index 677fa600e1..864000d034 100644 --- a/libs/databrickscfg/profiles.go +++ b/libs/databrickscfg/profiles.go @@ -14,11 +14,9 @@ import ( // It should only be used for prompting and filtering. // Use its name to construct a config.Config. type Profile struct { - Name string - Host string - AccountID string - ClusterID string - WarehouseID string + Name string + Host string + AccountID string } func (p Profile) Cloud() string { @@ -119,11 +117,9 @@ func LoadProfiles(fn ProfileMatchFunction) (file string, profiles Profiles, err continue } profile := Profile{ - Name: v.Name(), - Host: host, - AccountID: all["account_id"], - ClusterID: all["cluster_id"], - WarehouseID: all["warehouse_id"], + Name: v.Name(), + Host: host, + AccountID: all["account_id"], } if fn(profile) { profiles = append(profiles, profile)