fix: filter compute discovery to interactive cluster sources

python/databricks_kernel_dialect/cli.py

@@ -8,12 +8,18 @@
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import DatabricksError
+from databricks.sdk.service.compute import ClusterSource, ListClustersFilterBy
 
 PROTOCOL_VERSION = 1
 
+# Notebooks can only attach to interactive compute. Restricting the list at the
+# API keeps discovery fast on shared workspaces, which can accumulate tens of
+# thousands of ephemeral JOB clusters that would otherwise all be fetched.
+_INTERACTIVE_SOURCES = ListClustersFilterBy(cluster_sources=[ClusterSource.UI, ClusterSource.API])
+
 
 class _Clusters(Protocol):
-    def list(self) -> Iterable[object]: ...
+    def list(self, *, filter_by: object = ...) -> Iterable[object]: ...
 
 
 class _Workspace(Protocol):
@@ -96,7 +102,7 @@ def run(
     assert args.command == "compute" and args.action == "list"
     try:
         client = client_factory(args.profile)
-        compute = project_compute(client.clusters.list())
+        compute = project_compute(client.clusters.list(filter_by=_INTERACTIVE_SOURCES))
     except DatabricksError as error:
         if _is_authentication_error(error):
             print(f"Authentication failed for Databricks profile '{args.profile}'.", file=stderr)

python/tests/test_cli.py

@@ -27,8 +27,10 @@ class FakeCluster:
 class FakeClusters:
     def __init__(self, clusters: list[FakeCluster] | Exception) -> None:
         self._clusters = clusters
+        self.filter_by: object = None
 
-    def list(self) -> list[FakeCluster]:
+    def list(self, *, filter_by: object = None) -> list[FakeCluster]:
+        self.filter_by = filter_by
         if isinstance(self._clusters, Exception):
             raise self._clusters
         return self._clusters
@@ -113,6 +115,31 @@ def factory(profile: str) -> FakeWorkspace:
     }
 
 
+def test_cli_requests_only_interactive_cluster_sources() -> None:
+    # The workspace can hold tens of thousands of ephemeral JOB clusters; listing
+    # them all makes discovery unusably slow. The helper must ask the API to return
+    # only the interactive (UI/API) clusters a notebook could actually attach to.
+    from databricks.sdk.service.compute import ClusterSource
+
+    workspace = FakeWorkspace(
+        [FakeCluster("cluster-a", "Compute A", "15.4.x", Value.RUNNING, Value.UI)]
+    )
+
+    exit_code = run(
+        ["compute", "list", "--profile", "dev"],
+        stdout=io.StringIO(),
+        stderr=io.StringIO(),
+        client_factory=lambda _profile: workspace,
+    )
+
+    assert exit_code == 0
+    assert workspace.clusters.filter_by is not None
+    assert set(workspace.clusters.filter_by.cluster_sources) == {
+        ClusterSource.UI,
+        ClusterSource.API,
+    }
+
+
 def test_cli_reports_authentication_and_sdk_failures_on_stderr() -> None:
     cases = [
         (DatabricksError("expired", error_code="UNAUTHENTICATED"), 3, "Authentication failed"),