From 855cb23de2aeb1d8f67ac6d4fa2d54360086eefe Mon Sep 17 00:00:00 2001 From: Xiang Shen Date: Fri, 11 Sep 2026 20:42:56 +0000 Subject: [PATCH 1/3] Extract walk_catalog_schemas scaffold in databricks.py Factor the catalogs -> schemas -> parallel-per-schema-probe walk out of list_all_mcp_services into a generic walk_catalog_schemas(probe, collect), and rewrite list_all_mcp_services on top of it. The scaffold owns catalog/schema enumeration, worker pools, and the wall-clock deadline drain; the caller owns the probe, accumulation, dedup, progress, and streaming. No behavior change; the upcoming skills walk (download picker) becomes the second caller. Co-authored-by: Isaac --- src/ucode/databricks.py | 123 +++++++++++++++++++++++---------------- tests/test_databricks.py | 66 +++++++++++++++++++++ 2 files changed, 139 insertions(+), 50 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index e31749a4..8ca5b2e0 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2513,30 +2513,24 @@ 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 `.` - 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, +) -> str | None: + """Walk every user `.` in the workspace, probing each in parallel. + + Phase 1 lists catalogs (minus `skip_catalogs`) and their schemas in parallel, dropping + `information_schema`. Phase 2 runs `probe(catalog, schema)` on each surviving pair in parallel, + draining under `deadline` (an absolute `time.monotonic()` value) so a slow workspace degrades to + partial results. Each completed probe's value is handed to `collect(result, done, total)`, which + owns accumulation, progress, and any streaming. Returns a short reason when phase 1 finds + nothing, else None.""" 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,22 +2539,19 @@ 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_refs: list[tuple[str, str]] = [] schema_workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, len(catalog_names))) with ThreadPoolExecutor(max_workers=schema_workers) as pool: schema_futures = { @@ -2578,46 +2569,78 @@ def list_all_mcp_services( def collect_schemas(result, catalog): schemas, _ = result for schema in schemas: - schema_name = schema.get("name") - if ( - isinstance(schema_name, str) - and schema_name - and schema_name != "information_schema" - ): - schema_refs.append(f"{catalog}.{schema_name}") + name = schema.get("name") + if isinstance(name, str) and name and name != "information_schema": + schema_refs.append((catalog, 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)) 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 `.` + 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" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 0aacc636..9a727989 100644 --- a/tests/test_databricks.py +++ b/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.""" From 9f1dd550299bd69b4228673dd74cd1675500ddad Mon Sep 17 00:00:00 2001 From: Xiang Shen Date: Fri, 11 Sep 2026 22:11:02 +0000 Subject: [PATCH 2/3] Address review: clarify walk_catalog_schemas docstring, keep schema_name Rewrite the docstring around the high-level behavior and spell out the probe/collect callable shapes; revert the incidental schema_name -> name rename. Co-authored-by: Isaac --- src/ucode/databricks.py | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 8ca5b2e0..54d54690 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2522,14 +2522,19 @@ def walk_catalog_schemas[T]( collect: Callable[[T, int, int], None], skip_catalogs: frozenset[str] = _UC_FUNCTIONS_SKIP_CATALOGS, ) -> str | None: - """Walk every user `.` in the workspace, probing each in parallel. - - Phase 1 lists catalogs (minus `skip_catalogs`) and their schemas in parallel, dropping - `information_schema`. Phase 2 runs `probe(catalog, schema)` on each surviving pair in parallel, - draining under `deadline` (an absolute `time.monotonic()` value) so a slow workspace degrades to - partial results. Each completed probe's value is handed to `collect(result, done, total)`, which - owns accumulation, progress, and any streaming. Returns a short reason when phase 1 finds - nothing, else None.""" + """Discover every user `.` 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) catalogs, catalogs_reason = _paginated_json_items( @@ -2569,9 +2574,13 @@ def walk_catalog_schemas[T]( def collect_schemas(result, catalog): schemas, _ = result for schema in schemas: - name = schema.get("name") - if isinstance(name, str) and name and name != "information_schema": - schema_refs.append((catalog, name)) + schema_name = schema.get("name") + if ( + isinstance(schema_name, str) + and schema_name + and schema_name != "information_schema" + ): + schema_refs.append((catalog, schema_name)) _drain_with_deadline(schema_futures, deadline, collect_schemas) pool.shutdown(wait=False, cancel_futures=True) From 8dfc114521f2471d8dc557e28e994b761086abff Mon Sep 17 00:00:00 2001 From: Xiang Shen Date: Fri, 11 Sep 2026 22:16:47 +0000 Subject: [PATCH 3/3] Address review: make walk_catalog_schemas worker count a param Rename _UC_FUNCTION_PROBE_WORKERS -> _SCHEMA_PROBE_WORKERS (a generic per-schema concurrency cap, not UC-function-specific) and expose it as the max_workers param, mirroring skip_catalogs. Co-authored-by: Isaac --- src/ucode/databricks.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 54d54690..aa998c53 100644 --- a/src/ucode/databricks.py +++ b/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 @@ -2521,6 +2521,7 @@ def walk_catalog_schemas[T]( 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 `.` in the workspace and probe each one in parallel. @@ -2557,7 +2558,7 @@ def walk_catalog_schemas[T]( return "deadline exceeded while listing UC catalogs" schema_refs: list[tuple[str, str]] = [] - schema_workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, len(catalog_names))) + schema_workers = max(1, min(max_workers, len(catalog_names))) with ThreadPoolExecutor(max_workers=schema_workers) as pool: schema_futures = { pool.submit( @@ -2592,7 +2593,7 @@ def collect_schemas(result, catalog): 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: probe_futures = { pool.submit(probe, catalog, schema): (catalog, schema)