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"),