Added listing cluster filtering for cluster lookups

.codegen/lookup.go.tmpl

@@ -116,6 +116,10 @@ func allResolvers() *resolvers {
 	{{range .Services -}}
 	{{- if in $allowlist .KebabName -}}
 	r.{{.Singular.PascalName}} = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
+		fn, ok := lookupOverrides["{{.Singular.PascalName}}"]
+		if ok {
+			return fn(ctx, w, name)
+		}
 		entity, err := w.{{.PascalName}}.GetBy{{range .NamedIdMap.NamePath}}{{.PascalName}}{{end}}(ctx, name)
 		if err != nil {
 			return "", err

bundle/config/mutator/resolve_resource_references_test.go

@@ -2,7 +2,6 @@ package mutator
 
 import (
 	"context"
-	"fmt"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
@@ -44,11 +43,13 @@ func TestResolveClusterReference(t *testing.T) {
 	m := mocks.NewMockWorkspaceClient(t)
 	b.SetWorkpaceClient(m.WorkspaceClient)
 	clusterApi := m.GetMockClustersAPI()
-	clusterApi.EXPECT().GetByClusterName(mock.Anything, clusterRef1).Return(&compute.ClusterDetails{
-		ClusterId: "1234-5678-abcd",
-	}, nil)
-	clusterApi.EXPECT().GetByClusterName(mock.Anything, clusterRef2).Return(&compute.ClusterDetails{
-		ClusterId: "9876-5432-xywz",
+	clusterApi.EXPECT().ListAll(mock.Anything, compute.ListClustersRequest{
+		FilterBy: &compute.ListClustersFilterBy{
+			ClusterSources: []compute.ClusterSource{compute.ClusterSourceApi, compute.ClusterSourceUi},
+		},
+	}).Return([]compute.ClusterDetails{
+		{ClusterId: "1234-5678-abcd", ClusterName: clusterRef1},
+		{ClusterId: "9876-5432-xywz", ClusterName: clusterRef2},
 	}, nil)
 
 	diags := bundle.Apply(context.Background(), b, ResolveResourceReferences())
@@ -78,10 +79,16 @@ func TestResolveNonExistentClusterReference(t *testing.T) {
 	m := mocks.NewMockWorkspaceClient(t)
 	b.SetWorkpaceClient(m.WorkspaceClient)
 	clusterApi := m.GetMockClustersAPI()
-	clusterApi.EXPECT().GetByClusterName(mock.Anything, clusterRef).Return(nil, fmt.Errorf("ClusterDetails named '%s' does not exist", clusterRef))
+	clusterApi.EXPECT().ListAll(mock.Anything, compute.ListClustersRequest{
+		FilterBy: &compute.ListClustersFilterBy{
+			ClusterSources: []compute.ClusterSource{compute.ClusterSourceApi, compute.ClusterSourceUi},
+		},
+	}).Return([]compute.ClusterDetails{
+		{ClusterId: "1234-5678-abcd", ClusterName: "some other cluster"},
+	}, nil)
 
 	diags := bundle.Apply(context.Background(), b, ResolveResourceReferences())
-	require.ErrorContains(t, diags.Error(), "failed to resolve cluster: Random, err: ClusterDetails named 'Random' does not exist")
+	require.ErrorContains(t, diags.Error(), "failed to resolve cluster: Random, err: cluster named 'Random' does not exist")
 }
 
 func TestNoLookupIfVariableIsSet(t *testing.T) {
@@ -158,8 +165,14 @@ func TestResolveVariableReferencesInVariableLookups(t *testing.T) {
 	m := mocks.NewMockWorkspaceClient(t)
 	b.SetWorkpaceClient(m.WorkspaceClient)
 	clusterApi := m.GetMockClustersAPI()
-	clusterApi.EXPECT().GetByClusterName(mock.Anything, "cluster-bar-dev").Return(&compute.ClusterDetails{
-		ClusterId: "1234-5678-abcd",
+
+	clusterApi.EXPECT().ListAll(mock.Anything, compute.ListClustersRequest{
+		FilterBy: &compute.ListClustersFilterBy{
+			ClusterSources: []compute.ClusterSource{compute.ClusterSourceApi, compute.ClusterSourceUi},
+		},
+	}).Return([]compute.ClusterDetails{
+		{ClusterId: "1234-5678-abcd", ClusterName: "cluster-bar-dev"},
+		{ClusterId: "9876-5432-xywz", ClusterName: "some other cluster"},
 	}, nil)
 
 	diags := bundle.Apply(context.Background(), b, bundle.Seq(ResolveVariableReferencesInLookup(), ResolveResourceReferences()))

bundle/config/variable/lookup_overrides.go

@@ -0,0 +1,41 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+)
+
+var lookupOverrides = map[string]resolverFunc{
+	"Cluster": resolveCluster,
+}
+
+// We added a custom resolver for the cluster to add filtering for the cluster source when we list all clusters.
+// Without the filtering listing could take a very long time (5-10 mins) which leads to lookup timeouts.
+func resolveCluster(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
+	result, err := w.Clusters.ListAll(ctx, compute.ListClustersRequest{
+		FilterBy: &compute.ListClustersFilterBy{
+			ClusterSources: []compute.ClusterSource{compute.ClusterSourceApi, compute.ClusterSourceUi},
+		},
+	})
+
+	if err != nil {
+		return "", err
+	}
+
+	tmp := map[string][]compute.ClusterDetails{}
+	for _, v := range result {
+		key := v.ClusterName
+		tmp[key] = append(tmp[key], v)
+	}
+	alternatives, ok := tmp[name]
+	if !ok || len(alternatives) == 0 {
+		return "", fmt.Errorf("cluster named '%s' does not exist", name)
+	}
+	if len(alternatives) > 1 {
+		return "", fmt.Errorf("there are %d instances of clusters named '%s'", len(alternatives), name)
+	}
+	return alternatives[0].ClusterId, nil
+}

bundle/tests/variables_test.go

@@ -124,8 +124,13 @@ func TestVariablesWithTargetLookupOverrides(t *testing.T) {
 	}, nil)
 
 	clustersApi := mockWorkspaceClient.GetMockClustersAPI()
-	clustersApi.EXPECT().GetByClusterName(mock.Anything, "some-test-cluster").Return(&compute.ClusterDetails{
-		ClusterId: "4321",
+	clustersApi.EXPECT().ListAll(mock.Anything, compute.ListClustersRequest{
+		FilterBy: &compute.ListClustersFilterBy{
+			ClusterSources: []compute.ClusterSource{compute.ClusterSourceApi, compute.ClusterSourceUi},
+		},
+	}).Return([]compute.ClusterDetails{
+		{ClusterId: "4321", ClusterName: "some-test-cluster"},
+		{ClusterId: "9876", ClusterName: "some-other-cluster"},
 	}, nil)
 
 	clusterPoliciesApi := mockWorkspaceClient.GetMockClusterPoliciesAPI()