Extract walk_catalog_schemas scaffold in databricks.py

src/ucode/databricks.py

@@ -2439,7 +2439,7 @@ def resolve_provider_launch_model(model: str | None, provider_models: dict[str,
 
 _UC_LIST_PAGE_SIZE = 200
 _UC_LIST_MAX_PAGES = 50
-_UC_FUNCTION_PROBE_WORKERS = 16
+_SCHEMA_PROBE_WORKERS = 16
 _UC_LIST_HTTP_TIMEOUT = 10
 # Most MCP services live outside `system.ai`, so this workspace-wide walk needs
 # enough time to enumerate them; a slow workspace still degrades to partial
@@ -2513,30 +2513,30 @@ def _paginated_json_items(
     return items, last_reason
 
 
-def list_all_mcp_services(
+def walk_catalog_schemas[T](
     workspace: str,
     token: str,
     *,
-    deadline_seconds: float = _MCP_SERVICES_WALK_DEADLINE_SECONDS,
-    on_progress: Callable[[int, int, int], None] | None = None,
-    on_services: Callable[[list[str]], None] | None = None,
-) -> tuple[list[str], str | None]:
-    """Return sorted unique MCP-service full names across every `<catalog>.<schema>`
-    in the workspace. The mcp-services API is one-schema-per-call, so this walks
-    catalogs -> schemas -> mcp-services in parallel under a wall-clock budget,
-    returning partial results once `deadline_seconds` is exceeded.
-
-    `on_progress`, if given, is called as each schema's listing completes with
-    `(schemas_done, schemas_total, services_found)` so callers can render a live
-    count. `on_services`, if given, is called with each schema's newly-found service
-    names (deduped against everything emitted so far) so callers can stream results
-    into a picker as the walk progresses instead of waiting for the full result. Both
-    are invoked serially from the draining thread (not the workers).
-
-    This walk is the slow, workspace-wide counterpart to `list_mcp_services`
-    (single schema)."""
+    deadline: float,
+    probe: Callable[[str, str], T],
+    collect: Callable[[T, int, int], None],
+    skip_catalogs: frozenset[str] = _UC_FUNCTIONS_SKIP_CATALOGS,
+    max_workers: int = _SCHEMA_PROBE_WORKERS,
+) -> str | None:
+    """Discover every user `<catalog>.<schema>` in the workspace and probe each one in parallel.
+
+    Catalogs and their schemas are listed (skipping `skip_catalogs` and `information_schema`), then
+    each schema is probed concurrently until `deadline` (an absolute `time.monotonic()` value)
+    passes, so a slow workspace returns partial results instead of hanging. The caller supplies two
+    callables and owns whatever they accumulate:
+
+      - `probe(catalog, schema) -> result`: fetch one schema's data (e.g. its MCP services).
+      - `collect(result, done, total)`: handle each probe result as it lands — accumulating,
+        de-duping, streaming — where `done`/`total` are the completed and total schema counts.
+
+    Returns None once the probes run, or a short reason string if there are no catalogs or schemas
+    to probe."""
     hostname = workspace_hostname(workspace)
-    deadline = time.monotonic() + deadline_seconds
 
     catalogs, catalogs_reason = _paginated_json_items(
         f"https://{hostname}/api/2.1/unity-catalog/catalogs",
@@ -2545,23 +2545,20 @@ def list_all_mcp_services(
         timeout=_UC_LIST_HTTP_TIMEOUT,
     )
     if not catalogs:
-        return [], catalogs_reason or "no UC catalogs found"
+        return catalogs_reason or "no UC catalogs found"
 
     catalog_names = [
         c["name"]
         for c in catalogs
-        if isinstance(c.get("name"), str)
-        and c["name"]
-        and c["name"] not in _UC_FUNCTIONS_SKIP_CATALOGS
+        if isinstance(c.get("name"), str) and c["name"] and c["name"] not in skip_catalogs
     ]
     if not catalog_names:
-        return [], "no user UC catalogs found"
+        return "no user UC catalogs found"
     if time.monotonic() > deadline:
-        return [], "deadline exceeded while listing UC catalogs"
+        return "deadline exceeded while listing UC catalogs"
 
-    # Parallel per-catalog schema listing.
-    schema_refs: list[str] = []
-    schema_workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, len(catalog_names)))
+    schema_refs: list[tuple[str, str]] = []
+    schema_workers = max(1, min(max_workers, len(catalog_names)))
     with ThreadPoolExecutor(max_workers=schema_workers) as pool:
         schema_futures = {
             pool.submit(
@@ -2584,40 +2581,76 @@ def collect_schemas(result, catalog):
                     and schema_name
                     and schema_name != "information_schema"
                 ):
-                    schema_refs.append(f"{catalog}.{schema_name}")
+                    schema_refs.append((catalog, schema_name))
 
         _drain_with_deadline(schema_futures, deadline, collect_schemas)
         pool.shutdown(wait=False, cancel_futures=True)
 
     if not schema_refs:
         if time.monotonic() > deadline:
-            return [], "deadline exceeded while listing UC schemas"
-        return [], "no UC schemas found"
+            return "deadline exceeded while listing UC schemas"
+        return "no UC schemas found"
 
-    # Parallel per-schema mcp-services listing.
-    names: set[str] = set()
     schemas_total = len(schema_refs)
     schemas_done = 0
-    probe_workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, schemas_total))
+    probe_workers = max(1, min(max_workers, schemas_total))
     with ThreadPoolExecutor(max_workers=probe_workers) as pool:
-        service_futures = {
-            pool.submit(list_mcp_services, workspace, token, ref): ref for ref in schema_refs
+        probe_futures = {
+            pool.submit(probe, catalog, schema): (catalog, schema)
+            for catalog, schema in schema_refs
         }
 
-        def collect_services(result, _ref):
+        def collect_probe(result, _ref):
             nonlocal schemas_done
-            found, _ = result
-            new = [n for n in found if n not in names]
-            names.update(found)
             schemas_done += 1
-            if on_progress is not None:
-                on_progress(schemas_done, schemas_total, len(names))
-            if on_services is not None and new:
-                on_services(sorted(new))
+            collect(result, schemas_done, schemas_total)
 
-        _drain_with_deadline(service_futures, deadline, collect_services)
+        _drain_with_deadline(probe_futures, deadline, collect_probe)
         pool.shutdown(wait=False, cancel_futures=True)
 
+    return None
+
+
+def list_all_mcp_services(
+    workspace: str,
+    token: str,
+    *,
+    deadline_seconds: float = _MCP_SERVICES_WALK_DEADLINE_SECONDS,
+    on_progress: Callable[[int, int, int], None] | None = None,
+    on_services: Callable[[list[str]], None] | None = None,
+) -> tuple[list[str], str | None]:
+    """Return sorted unique MCP-service full names across every `<catalog>.<schema>`
+    in the workspace. The mcp-services API is one-schema-per-call, so this walks
+    catalogs -> schemas -> mcp-services in parallel under a wall-clock budget,
+    returning partial results once `deadline_seconds` is exceeded.
+
+    `on_progress`, if given, is called as each schema's listing completes with
+    `(schemas_done, schemas_total, services_found)` so callers can render a live
+    count. `on_services`, if given, is called with each schema's newly-found service
+    names (deduped against everything emitted so far) so callers can stream results
+    into a picker as the walk progresses instead of waiting for the full result. Both
+    are invoked serially from the draining thread (not the workers).
+
+    This walk is the slow, workspace-wide counterpart to `list_mcp_services`
+    (single schema)."""
+    deadline = time.monotonic() + deadline_seconds
+    names: set[str] = set()
+
+    def probe(catalog, schema):
+        return list_mcp_services(workspace, token, f"{catalog}.{schema}")
+
+    def collect(result, schemas_done, schemas_total):
+        found, _ = result
+        new = [n for n in found if n not in names]
+        names.update(found)
+        if on_progress is not None:
+            on_progress(schemas_done, schemas_total, len(names))
+        if on_services is not None and new:
+            on_services(sorted(new))
+
+    reason = walk_catalog_schemas(workspace, token, deadline=deadline, probe=probe, collect=collect)
+    if reason is not None:
+        return [], reason
     if not names:
         if time.monotonic() > deadline:
             return [], "deadline exceeded while listing MCP services"

tests/test_databricks.py

@@ -5,6 +5,7 @@
 import json
 import os
 import subprocess
+import time
 from decimal import Decimal
 from urllib.parse import parse_qs
 
@@ -1220,6 +1221,71 @@ def test_http_404_reason_surfaces_for_invalid_parent(self, monkeypatch):
         assert reason and reason.startswith("HTTP 404")
 
 
+class TestWalkCatalogSchemas:
+    """The generic catalogs -> schemas -> per-schema probe scaffold, independent of any probe."""
+
+    def _fake_catalog_http(self, catalogs, schemas_by_catalog):
+        def fake_get(url, token, timeout=30):
+            if "unity-catalog/catalogs" in url:
+                return {"catalogs": [{"name": c} for c in catalogs]}, None
+            if "unity-catalog/schemas" in url:
+                cat = url.split("catalog_name=")[1].split("&")[0]
+                return {"schemas": [{"name": s} for s in schemas_by_catalog.get(cat, [])]}, None
+            return None, "unexpected url"
+
+        return fake_get
+
+    def test_probes_each_user_schema_and_reports_progress(self, monkeypatch):
+        monkeypatch.setattr(
+            db_mod,
+            "_http_get_json",
+            self._fake_catalog_http(
+                catalogs=["mycat", "system"],
+                schemas_by_catalog={"mycat": ["a", "b", "information_schema"]},
+            ),
+        )
+        probed: list[tuple[str, str]] = []
+        collected: list[tuple[str, int, int]] = []
+
+        def probe(catalog, schema):
+            probed.append((catalog, schema))
+            return f"{catalog}.{schema}"
+
+        def collect(result, done, total):
+            collected.append((result, done, total))
+
+        reason = db_mod.walk_catalog_schemas(
+            WS, "token", deadline=time.monotonic() + 30, probe=probe, collect=collect
+        )
+
+        assert reason is None
+        # system is skipped and information_schema is dropped; only user schemas are probed.
+        assert sorted(probed) == [("mycat", "a"), ("mycat", "b")]
+        assert sorted(r for r, _, _ in collected) == ["mycat.a", "mycat.b"]
+        # One collect per probed schema; total is fixed and done climbs to it.
+        assert [total for _, _, total in collected] == [2, 2]
+        assert sorted(done for _, done, _ in collected) == [1, 2]
+
+    def test_returns_reason_when_all_catalogs_skipped(self, monkeypatch):
+        monkeypatch.setattr(
+            db_mod,
+            "_http_get_json",
+            self._fake_catalog_http(catalogs=["system", "samples"], schemas_by_catalog={}),
+        )
+        probed: list[tuple[str, str]] = []
+
+        reason = db_mod.walk_catalog_schemas(
+            WS,
+            "token",
+            deadline=time.monotonic() + 30,
+            probe=lambda catalog, schema: probed.append((catalog, schema)),
+            collect=lambda *args: None,
+        )
+
+        assert reason == "no user UC catalogs found"
+        assert probed == []
+
+
 class TestListAllMcpServices:
     """Workspace-wide walk: catalogs -> schemas -> per-schema mcp-services."""