diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 0d021e0..087ac89 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -221,7 +221,13 @@ def _http_get_json( _debug(f"GET {url}", f"HTTP {exc.code} {exc.reason}") if _debug_enabled() and body: _debug("body", body[:4000]) - return None, f"HTTP {exc.code} {exc.reason}" + reason = f"HTTP {exc.code} {exc.reason}" + # Surface the response body too — gateway auth failures return 400 + # with body `Invalid Token`, which is invisible without this. + body_excerpt = body.strip()[:200] + if body_excerpt: + reason = f"{reason}: {body_excerpt}" + return None, reason except urllib_error.URLError as exc: _debug(f"GET {url}", f"URLError: {exc.reason}") return None, f"network error: {exc.reason}" @@ -903,20 +909,50 @@ def ensure_ai_gateway_v2(workspace: str, token: str) -> None: Uses the dedicated v2 listing endpoint `GET /api/ai-gateway/v2/endpoints`: a 200 response (even with an empty list) means v2 is wired up on this workspace — a "no endpoints provisioned" case will surface naturally in - downstream discovery. 404 / 401 / 403 / network failures all raise a - clear error with the docs link instead of silently progressing. + downstream discovery. Failure branches: + + - 401 / 403 / 400 with `Invalid Token`: the token is bad for *this* + workspace. + - 404: AI Gateway V2 is not enabled on this workspace — point at the docs. + - other (5xx, network errors): surface the reason verbatim. """ hostname = workspace_hostname(workspace) url = f"https://{hostname}/api/ai-gateway/v2/endpoints?page_size=1" payload, reason = _http_get_json(url, token) if payload is not None: return + reason_str = reason or "unknown error" + if _looks_like_auth_failure(reason_str): + raise RuntimeError( + f"Databricks rejected the access token for {workspace} ({reason_str}). " + f"Try:\n" + f" databricks auth logout --host {workspace}\n" + f" databricks auth login --host {workspace}" + ) + if "HTTP 404" in reason_str: + raise RuntimeError( + "Databricks Unity AI Gateway is not enabled on this workspace " + f"({reason_str}). See {AI_GATEWAY_V2_DOCS_URL}" + ) raise RuntimeError( - "Databricks AI Gateway V2 is required but not available on this workspace " - f"({reason}). See {AI_GATEWAY_V2_DOCS_URL}" + "Databricks Unity AI Gateway probe failed on this workspace " + f"({reason_str}). See {AI_GATEWAY_V2_DOCS_URL}" ) +def _looks_like_auth_failure(reason: str) -> bool: + """True when the gateway response signals the token is not accepted. + + Covers 401/403 directly and the gateway's 400 + `Invalid Token` body + (which happens when the bearer is valid but issued for a different + workspace).""" + if "HTTP 401" in reason or "HTTP 403" in reason: + return True + if "HTTP 400" in reason and "invalid token" in reason.lower(): + return True + return False + + def discover_sql_warehouse_http_path( workspace: str, token: str, diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 16cd749..6857efc 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -501,27 +501,68 @@ def _mock_json_response(body: str): mock_resp.read.return_value = body.encode("utf-8") return mock_resp - def test_raises_on_404(self): - from unittest.mock import MagicMock, patch + @staticmethod + def _http_error(code: int, msg: str, body: str = ""): + import io + from unittest.mock import MagicMock from urllib.error import HTTPError - exc = HTTPError(url="", code=404, msg="Not Found", hdrs=MagicMock(), fp=None) + fp = io.BytesIO(body.encode("utf-8")) if body else None + return HTTPError(url="", code=code, msg=msg, hdrs=MagicMock(), fp=fp) + + def test_raises_on_404(self): + from unittest.mock import patch + + exc = self._http_error(404, "Not Found") with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): from ucode.databricks import ensure_ai_gateway_v2 - with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL): + with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL) as excinfo: ensure_ai_gateway_v2(WS, "fake-token") + assert "not enabled" in str(excinfo.value) - def test_raises_on_401(self): - from unittest.mock import MagicMock, patch - from urllib.error import HTTPError + def test_raises_on_401_with_auth_hint(self): + from unittest.mock import patch + + exc = self._http_error(401, "Unauthorized") + with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): + from ucode.databricks import ensure_ai_gateway_v2 + + with pytest.raises(RuntimeError, match="401") as excinfo: + ensure_ai_gateway_v2(WS, "fake-token") + message = str(excinfo.value) + assert "rejected" in message.lower() + assert "databricks auth login" in message + + def test_raises_on_400_invalid_token_with_auth_hint(self): + """400 + body `Invalid Token` is the misleading-error case from issue #84.""" + from unittest.mock import patch + + exc = self._http_error(400, "Bad Request", body="Invalid Token") + with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): + from ucode.databricks import ensure_ai_gateway_v2 - exc = HTTPError(url="", code=401, msg="Unauthorized", hdrs=MagicMock(), fp=None) + with pytest.raises(RuntimeError) as excinfo: + ensure_ai_gateway_v2(WS, "fake-token") + message = str(excinfo.value) + # The bug we are fixing: must NOT collapse to the generic + # "v2 not available" message — must call out the auth failure + # and point at re-login. + assert "Invalid Token" in message + assert "rejected" in message.lower() + assert "databricks auth login" in message + + def test_400_without_invalid_token_falls_through_to_generic(self): + """A 400 that is *not* an auth failure should still surface the body.""" + from unittest.mock import patch + + exc = self._http_error(400, "Bad Request", body="some other detail") with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): from ucode.databricks import ensure_ai_gateway_v2 - with pytest.raises(RuntimeError, match="401"): + with pytest.raises(RuntimeError, match=AI_GATEWAY_V2_DOCS_URL) as excinfo: ensure_ai_gateway_v2(WS, "fake-token") + assert "some other detail" in str(excinfo.value) def test_raises_on_url_error(self): from unittest.mock import patch @@ -561,6 +602,44 @@ def test_succeeds_with_empty_endpoints_list(self): ensure_ai_gateway_v2(WS, "fake-token") # should not raise +class TestHttpGetJsonReason: + """The `reason` string returned by `_http_get_json` must include the response body + so callers (e.g. ensure_ai_gateway_v2) can route on it. Before issue #84's fix + the body was logged only when UCODE_DEBUG=1 and dropped from the bubbled error.""" + + @staticmethod + def _http_error(code: int, msg: str, body: str = ""): + import io + from unittest.mock import MagicMock + from urllib.error import HTTPError + + fp = io.BytesIO(body.encode("utf-8")) if body else None + return HTTPError(url="", code=code, msg=msg, hdrs=MagicMock(), fp=fp) + + def test_reason_includes_body_on_http_error(self): + from unittest.mock import patch + + from ucode.databricks import _http_get_json + + exc = self._http_error(400, "Bad Request", body="Invalid Token") + with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): + payload, reason = _http_get_json("https://x/y", "tok") + assert payload is None + assert "HTTP 400" in reason + assert "Invalid Token" in reason + + def test_reason_without_body_is_status_only(self): + from unittest.mock import patch + + from ucode.databricks import _http_get_json + + exc = self._http_error(404, "Not Found") + with patch("ucode.databricks.urllib_request.urlopen", side_effect=exc): + payload, reason = _http_get_json("https://x/y", "tok") + assert payload is None + assert reason == "HTTP 404 Not Found" + + class TestParseDatabricksCliVersion: def test_parses_standard_format(self): assert _parse_databricks_cli_version("Databricks CLI v0.299.2") == (0, 299, 2)