From 9b7f2b1fa3dc24d0f9f29c0f251b3b501541b627 Mon Sep 17 00:00:00 2001 From: Tom Corbin Date: Mon, 3 Aug 2026 13:35:08 +0100 Subject: [PATCH] fix: filter compute discovery to interactive cluster sources The compute helper listed every cluster in the workspace before the plugin reported readiness. On shared workspaces this can be tens of thousands of ephemeral JOB clusters (observed: 22,590), taking ~8 minutes and blocking DatabricksRunCell the entire time. Restrict the SDK list call to UI/API sources via ListClustersFilterBy, so the API returns only the interactive compute a notebook could attach to. Discovery drops from ~8m to ~4s against the same workspace. --- python/databricks_kernel_dialect/cli.py | 10 +++++++-- python/tests/test_cli.py | 29 ++++++++++++++++++++++++- 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/python/databricks_kernel_dialect/cli.py b/python/databricks_kernel_dialect/cli.py index c65c593..e4d4d05 100644 --- a/python/databricks_kernel_dialect/cli.py +++ b/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) diff --git a/python/tests/test_cli.py b/python/tests/test_cli.py index 3c6036f..dca6e58 100644 --- a/python/tests/test_cli.py +++ b/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"),