From 85b9768925539e2fddb7c9d3d6c260c2ab0d8fe6 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Tue, 4 Aug 2026 22:52:23 +0000 Subject: [PATCH] databricks: paginate the model-provider-service listing, and get by name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `main.tien_le.openai_all` exists — the UC explore page shows it, and a direct GET returns HTTP 200 — but ucode reported: Model provider service 'main.tien_le.openai_all' was not found. `resolve_provider_service` lists and filters in memory, and the listing made a single unpaginated request. That metastore has 30 services across 2 pages, so ucode saw an arbitrary 16 and discarded the `next_page_token`. Anything on a later page was invisible — to the resolver and to the interactive picker alike. - `list_model_provider_services` now pages, mirroring `list_model_services` right above it (bounded `page_size`, `next_page_token`, a `seen_tokens` guard against a server echoing one back). A mid-pagination failure degrades to partial results rather than an error, since partial data is more useful than none. - It also accepts `parent` to scope the listing to one `catalog.schema`. The metastore-wide default is documented in the proto as an internal scope that is "likely to be deprecated", so callers that know the schema should pass it. - `get_model_provider_service` addresses a service by name, and `resolve_provider_service` falls back to it before concluding "not found" — a named service should never be judged absent on the strength of a listing that might be incomplete. - The entry-parsing moves to `_provider_service_entry`, shared by both paths. Verified against eng-ml-inference.staging: the listing goes from 16 to 30 services, `main.tien_le.openai_all` appears, `parent="main.tien_le"` returns exactly the three services in that schema, and resolving a genuinely absent name still errors. Tests: +8. `test_follows_next_page_token` was checked against the old single-page behavior, where it fails with "Right contains one more item: 'main.s.two'" — the same page-2 loss that caused the report. Co-authored-by: Isaac --- src/ucode/databricks.py | 150 +++++++++++++++++++++++++++++---------- tests/test_databricks.py | 130 +++++++++++++++++++++++++++++++++ 2 files changed, 242 insertions(+), 38 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d9cf166..0a004a1 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1495,7 +1495,15 @@ def _provider_type_tag(provider_type: str | None) -> str: return tag.lower() -def list_model_provider_services(workspace: str, token: str) -> tuple[list[dict], str | None]: +# The listing is paginated; a metastore with more services than one page silently truncated before +# this was honored, making services on later pages look nonexistent. +_PROVIDER_SERVICES_PAGE_SIZE = 100 +_PROVIDER_SERVICES_MAX_PAGES = 50 + + +def list_model_provider_services( + workspace: str, token: str, *, parent: str | None = None +) -> tuple[list[dict], str | None]: """List Unity Catalog Model Provider Services on the workspace. Returns ``(services, reason)`` where each service is @@ -1505,46 +1513,107 @@ def list_model_provider_services(workspace: str, token: str) -> tuple[list[dict] Bedrock model names). ``relayed`` is True for a credential-less Anthropic service (Claude Max/Team/Enterprise subscription relay). A non-None ``reason`` means the listing call itself failed. + + Pages through the endpoint: a metastore with more services than fit on one page used to have the + remainder silently dropped, so a service that plainly existed looked absent. ``parent`` scopes + the listing to one ``catalog.schema`` — the metastore-wide default is documented as an internal, + likely-to-be-deprecated scope, so prefer passing it when the schema is known. """ hostname = workspace_hostname(workspace) - url = f"https://{hostname}/api/2.1/unity-catalog/model-provider-services" - payload, reason = _http_get_json(url, token, timeout=30) - if payload is None: - return [], reason - data = cast(dict, payload) if isinstance(payload, dict) else {} services: list[dict] = [] - for service in data.get("model_provider_services") or []: - if not isinstance(service, dict): - continue - raw_name = service.get("name") - if not isinstance(raw_name, str) or not raw_name: - continue - # The API returns `model-provider-services/..`. - full_name = raw_name.split("/", 1)[1] if "/" in raw_name else raw_name - config = service.get("config") if isinstance(service.get("config"), dict) else {} - targets = [] - for target in config.get("targets") or []: - model_id = target.get("model") if isinstance(target, dict) else None - if isinstance(model_id, str) and model_id: - targets.append(model_id) - # Relayed = credential-less Anthropic (subscription relay). Only whether - # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed - # server-side, so both launch identically. - anthropic_cfg = config.get("anthropic") - relayed = isinstance(anthropic_cfg, dict) and "relayed" in anthropic_cfg - services.append( - { - "name": full_name, - "provider_type": _provider_type_tag(config.get("provider_type")), - "targets": targets, - "allow_all_targets": bool(config.get("allow_all_targets")), - "relayed": relayed, - } + page_token: str | None = None + seen_tokens: set[str] = set() + last_reason: str | None = None + for _ in range(_PROVIDER_SERVICES_MAX_PAGES): + params: dict[str, str] = {"page_size": str(_PROVIDER_SERVICES_PAGE_SIZE)} + if parent: + params["parent"] = f"schemas/{parent}" + if page_token: + params["page_token"] = page_token + url = ( + f"https://{hostname}/api/2.1/unity-catalog/model-provider-services?{urlencode(params)}" ) + payload, reason = _http_get_json(url, token, timeout=30) + if payload is None: + # Surface the failure only if we have nothing yet; a mid-pagination blip still + # returns whatever was collected. + last_reason = reason + break + data = cast(dict, payload) if isinstance(payload, dict) else {} + for service in data.get("model_provider_services") or []: + entry = _provider_service_entry(service) + if entry is not None: + services.append(entry) + page_token = data.get("next_page_token") or None + if not page_token: + last_reason = None + break + if page_token in seen_tokens: + break + seen_tokens.add(page_token) + + if not services and last_reason is not None: + return [], last_reason services.sort(key=lambda s: s["name"]) return services, None +def _provider_service_entry(raw_service: object) -> dict | None: + """Normalize one listing entry, or None when it isn't usable.""" + if not isinstance(raw_service, dict): + return None + # A bare isinstance narrows to dict[Never, Never], which rejects string keys. + service = cast("dict[str, object]", raw_service) + raw_name = service.get("name") + if not isinstance(raw_name, str) or not raw_name: + return None + # The API returns `model-provider-services/..`. + full_name = raw_name.split("/", 1)[1] if "/" in raw_name else raw_name + raw_config = service.get("config") + config = cast("dict[str, object]", raw_config) if isinstance(raw_config, dict) else {} + targets: list[str] = [] + raw_targets = config.get("targets") + for target in raw_targets if isinstance(raw_targets, list) else []: + if not isinstance(target, dict): + continue + model_id = cast("dict[str, object]", target).get("model") + if isinstance(model_id, str) and model_id: + targets.append(model_id) + # Relayed = credential-less Anthropic (subscription relay). Only whether + # it's relayed matters here; the tier (Max vs Team/Enterprise) is governed + # server-side, so both launch identically. + anthropic_cfg = config.get("anthropic") + relayed = isinstance(anthropic_cfg, dict) and "relayed" in anthropic_cfg + raw_type = config.get("provider_type") + return { + "name": full_name, + "provider_type": _provider_type_tag(raw_type if isinstance(raw_type, str) else None), + "targets": targets, + "allow_all_targets": bool(config.get("allow_all_targets")), + "relayed": relayed, + } + + +def get_model_provider_service( + service_name: str, workspace: str, token: str +) -> tuple[dict | None, str | None]: + """Fetch one provider service by its full `catalog.schema.name`, bypassing the listing. + + The listing is paginated and metastore-wide, so any gap in it (a page we failed to fetch, a + server-side filter) makes a service that plainly exists look absent. Addressing it directly + removes that whole class of false negative. + """ + hostname = workspace_hostname(workspace) + url = f"https://{hostname}/api/2.1/unity-catalog/model-provider-services/{service_name}" + payload, reason = _http_get_json(url, token, timeout=30) + if payload is None: + return None, reason + entry = _provider_service_entry(payload) + if entry is None: + return None, "model-provider-service response had an unexpected shape" + return entry, None + + def is_model_provider_feature_unavailable(reason: str | None) -> bool: """True when a model-provider-services API failure means the workspace simply hasn't enabled the feature (HTTP 400 "feature is not available"), @@ -1603,11 +1672,16 @@ def resolve_provider_service( return None, f"Could not list model provider services: {reason}" match = next((s for s in services if s["name"] == service_name), None) if match is None: - usable = [ - s["name"] for s in services if tool_supports_provider_type(tool, s["provider_type"]) - ] - suffix = f" Available for {tool}: {', '.join(usable)}." if usable else "" - return None, f"Model provider service '{service_name}' was not found.{suffix}" + # Don't conclude "not found" from a listing that may be incomplete — a named service can be + # fetched directly. Only when that 404s is it really absent. + match, get_reason = get_model_provider_service(service_name, workspace, token) + if match is None: + usable = [ + s["name"] for s in services if tool_supports_provider_type(tool, s["provider_type"]) + ] + suffix = f" Available for {tool}: {', '.join(usable)}." if usable else "" + detail = f" ({get_reason})" if get_reason and "404" not in get_reason else "" + return None, f"Model provider service '{service_name}' was not found.{detail}{suffix}" provider_type = match["provider_type"] if not tool_supports_provider_type(tool, provider_type): supported = ", ".join(_TOOL_PROVIDER_TYPES.get(tool, ())) or "none" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7e1a73a..5a3a4ca 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -499,6 +499,136 @@ def test_empty_when_no_claude(self): assert db_mod.map_bedrock_claude_models(["amazon.titan-text-express-v1"]) == {} +class TestProviderServicePagination: + """The listing is paginated; ignoring next_page_token hid services on later pages entirely.""" + + @staticmethod + def _page(names, next_token=None): + payload = { + "model_provider_services": [ + { + "name": f"model-provider-services/{n}", + "config": {"provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_OPENAI"}, + } + for n in names + ] + } + if next_token: + payload["next_page_token"] = next_token + return payload + + def test_follows_next_page_token(self, monkeypatch): + pages = [ + self._page(["main.s.one"], next_token="tok2"), + self._page(["main.s.two"]), + ] + seen: list[str] = [] + + def fake_get(url, token, **kwargs): + seen.append(url) + return pages[len(seen) - 1], None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + services, reason = db_mod.list_model_provider_services("https://ws", "tok") + + assert reason is None + assert [s["name"] for s in services] == ["main.s.one", "main.s.two"] + assert "page_token=tok2" in seen[1] + + def test_stops_on_a_repeated_token(self, monkeypatch): + # A server that echoes the same token would otherwise spin forever. + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token, **kw: (self._page(["main.s.one"], next_token="same"), None), + ) + services, reason = db_mod.list_model_provider_services("https://ws", "tok") + assert reason is None + assert len(services) >= 1 + + def test_keeps_earlier_pages_when_a_later_one_fails(self, monkeypatch): + calls = {"n": 0} + + def fake_get(url, token, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + return self._page(["main.s.one"], next_token="tok2"), None + return None, "HTTP 500 Server Error" + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + services, reason = db_mod.list_model_provider_services("https://ws", "tok") + + # A mid-pagination blip should degrade to partial results, not to an error. + assert reason is None + assert [s["name"] for s in services] == ["main.s.one"] + + def test_reports_the_failure_when_nothing_was_collected(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, **kw: (None, "HTTP 403 Forbidden") + ) + services, reason = db_mod.list_model_provider_services("https://ws", "tok") + assert services == [] + assert reason == "HTTP 403 Forbidden" + + def test_parent_scopes_the_listing(self, monkeypatch): + seen: dict = {} + + def fake_get(url, token, **kwargs): + seen["url"] = url + return self._page(["main.tien_le.openai"]), None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + db_mod.list_model_provider_services("https://ws", "tok", parent="main.tien_le") + + assert "parent=schemas%2Fmain.tien_le" in seen["url"] + + def test_page_size_is_always_sent(self, monkeypatch): + seen: dict = {} + + def fake_get(url, token, **kwargs): + seen["url"] = url + return self._page(["main.s.one"]), None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + db_mod.list_model_provider_services("https://ws", "tok") + + assert "page_size=" in seen["url"] + + +class TestGetModelProviderService: + def test_addresses_the_service_directly(self, monkeypatch): + seen: dict = {} + + def fake_get(url, token, **kwargs): + seen["url"] = url + return { + "name": "model-provider-services/main.tien_le.openai_all", + "config": { + "provider_type": "EXTERNAL_MODEL_PROVIDER_TYPE_OPENAI", + "allow_all_targets": True, + }, + }, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + service, reason = db_mod.get_model_provider_service( + "main.tien_le.openai_all", "https://ws", "tok" + ) + + assert reason is None + assert service is not None + assert service["name"] == "main.tien_le.openai_all" + assert service["allow_all_targets"] is True + assert seen["url"].endswith("/model-provider-services/main.tien_le.openai_all") + + def test_missing_service_returns_the_reason(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, **kw: (None, "HTTP 404 Not Found") + ) + service, reason = db_mod.get_model_provider_service("main.a.b", "https://ws", "tok") + assert service is None + assert "404" in (reason or "") + + class TestResolveProviderService: _PAYLOAD = TestListModelProviderServices._PAYLOAD