diff --git a/README.md b/README.md index 5bb6dbf..64096c1 100644 --- a/README.md +++ b/README.md @@ -226,6 +226,7 @@ pick the new config up on their next ucode run. |---------|-------------| | `ucode status` | Show current workspace, base URLs, managed config files, and selected models | | `ucode usage` | Show AI Gateway usage summary, plus your budget spend against its alert threshold when the workspace reports one | +| `ucode usage --warehouse-id ` | Query a specific SQL warehouse instead of discovering one | | `ucode revert` | Clear saved state and restore backed-up config files | | `ucode configure --dry-run` | Preview config files without writing them | | `ucode configure --agents claude,codex` | Configure specific agents without the interactive picker | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 2a079ff..a68c90d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2388,11 +2388,16 @@ def revert_cmd() -> None: @app.command("usage") -def usage_cmd() -> None: +def usage_cmd( + warehouse_id: Annotated[ + str | None, + typer.Option("--warehouse-id", help="SQL warehouse to query, instead of discovering one."), + ] = None, +) -> None: """Show Databricks AI Gateway usage summary (last 7 days).""" try: install_databricks_cli() - usage_report() + usage_report(warehouse_id=warehouse_id) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 580a262..ec79b98 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -25,7 +25,7 @@ ) from decimal import Decimal, InvalidOperation from pathlib import Path -from typing import Literal, cast, overload +from typing import Literal, NamedTuple, cast, overload from urllib import error as urllib_error from urllib import request as urllib_request from urllib.parse import urlencode, urlparse @@ -2710,12 +2710,27 @@ def _parse_decimal(value: object) -> Decimal | None: return None -def discover_sql_warehouse_http_path( +class SqlWarehouse(NamedTuple): + http_path: str + label: str + state: str + + +def discover_sql_warehouses( workspace: str, token: str, *, - quiet: bool = False, -) -> str: + warehouse_id: str | None = None, +) -> list[SqlWarehouse]: + """Candidate warehouses to run the usage query against, RUNNING ones first. + + Several are returned because a warehouse can report RUNNING and still refuse + connections, so callers fall through to the next one. An explicit + `warehouse_id` skips discovery entirely. + """ + if warehouse_id: + return [SqlWarehouse(_warehouse_http_path(warehouse_id), warehouse_id, "REQUESTED")] + hostname = workspace_hostname(workspace) request = urllib_request.Request( f"https://{hostname}/api/2.0/sql/warehouses", @@ -2740,33 +2755,30 @@ def discover_sql_warehouse_http_path( warehouses = payload.get("warehouses") if not isinstance(warehouses, list) or not warehouses: raise RuntimeError( - "No SQL warehouses found in this workspace. Create one or pass `--http-path`." + "No SQL warehouses found in this workspace. Create one or pass `--warehouse-id`." ) - running = [w for w in warehouses if isinstance(w, dict) and w.get("state") == "RUNNING"] - chosen = ( - running[0] - if running - else next( - (w for w in warehouses if isinstance(w, dict) and w.get("id")), - None, - ) - ) - if not chosen: + candidates: list[SqlWarehouse] = [] + for entry in warehouses: + if not isinstance(entry, dict): + continue + entry_id = entry.get("id") + if not isinstance(entry_id, str) or not entry_id.strip(): + continue + name = entry.get("name") + state = entry.get("state", "UNKNOWN") + label = name if isinstance(name, str) and name else entry_id + candidates.append(SqlWarehouse(_warehouse_http_path(entry_id), label, str(state))) + + if not candidates: raise RuntimeError("No usable SQL warehouse was returned by Databricks.") + # Stopped warehouses work too, but cold-starting one costs minutes. + candidates.sort(key=lambda w: w.state != "RUNNING") + return candidates - warehouse_id = chosen.get("id") - if not isinstance(warehouse_id, str) or not warehouse_id.strip(): - raise RuntimeError("Databricks returned a warehouse without an ID.") - warehouse_name = chosen.get("name") - warehouse_state = chosen.get("state", "UNKNOWN") - label_value = ( - warehouse_name if isinstance(warehouse_name, str) and warehouse_name else warehouse_id - ) - if not quiet: - print_note(f"Using SQL warehouse `{label_value}` ({warehouse_state}).") - return f"/sql/1.0/warehouses/{warehouse_id}" +def _warehouse_http_path(warehouse_id: str) -> str: + return f"/sql/1.0/warehouses/{warehouse_id.strip()}" def run_usage_query( @@ -2774,7 +2786,14 @@ def run_usage_query( http_path: str, token: str, query: str, + on_connected: Callable[[], None] | None = None, ) -> tuple[list[str], list[tuple]]: + """Run `query` on one warehouse. + + `on_connected` fires once the connection opens — the point a stopped + warehouse has finished starting — so callers can update their progress + message. + """ try: logging.getLogger("databricks.sql").setLevel(logging.ERROR) from databricks import sql @@ -2790,6 +2809,8 @@ def run_usage_query( http_path=http_path, access_token=token, ) as connection: + if on_connected is not None: + on_connected() with connection.cursor() as cursor: cursor.execute(query) columns = [desc[0] for desc in (cursor.description or [])] diff --git a/src/ucode/usage.py b/src/ucode/usage.py index b2d90f0..8d690a0 100644 --- a/src/ucode/usage.py +++ b/src/ucode/usage.py @@ -12,8 +12,9 @@ from typing import cast from ucode.databricks import ( + SqlWarehouse, apply_pat_environment, - discover_sql_warehouse_http_path, + discover_sql_warehouses, ensure_databricks_auth, get_databricks_token, resolve_current_budget_spend, @@ -31,6 +32,7 @@ muted, print_heading, print_note, + print_warning, render_box_table, spinner, value, @@ -39,6 +41,11 @@ USAGE_BREAKDOWN_DAYS = 7 USAGE_SUMMARY_DAYS = 30 +QUERY_MESSAGE = "Querying system.ai_gateway.usage..." +STARTUP_MESSAGE = "Starting up warehouse..." +# `REQUESTED` is an explicit --warehouse-id, whose state we never looked up. +WARM_WAREHOUSE_STATES = ("RUNNING", "REQUESTED") + def build_usage_report_query() -> str: return f""" @@ -460,7 +467,53 @@ def render_usage_summary( return "\n".join(lines) -def usage() -> int: +def run_query_on_first_working_warehouse( + workspace: str, + token: str, + candidates: list[SqlWarehouse], + query: str, +) -> tuple[str, list[str], list[tuple]]: + """Run `query` on the first candidate that accepts the connection. + + Returns the warehouse's http path alongside the result so later queries + reuse it. Raises the last error when every candidate fails. + """ + last_error: RuntimeError | None = None + for warehouse in candidates: + print_note(f"Using SQL warehouse `{warehouse.label}` ({warehouse.state}).") + try: + # Inside the loop so the spinner stops before any warning prints. + columns, rows = _query_with_progress(workspace, token, warehouse, query) + except RuntimeError as exc: + last_error = exc + print_warning(f"SQL warehouse `{warehouse.label}` is unusable: {exc}") + continue + return warehouse.http_path, columns, rows + raise last_error or RuntimeError("No SQL warehouse could run the usage query.") + + +def _query_with_progress( + workspace: str, + token: str, + warehouse: SqlWarehouse, + query: str, +) -> tuple[list[str], list[tuple]]: + """Run the query, reporting a cold start until the connection opens. + + A warehouse that isn't already up costs minutes to start, so the spinner + says that until `run_usage_query` reports it connected. + """ + connected = warehouse.state in WARM_WAREHOUSE_STATES + + def mark_connected() -> None: + nonlocal connected + connected = True + + with spinner(lambda: QUERY_MESSAGE if connected else STARTUP_MESSAGE): + return run_usage_query(workspace, warehouse.http_path, token, query, mark_connected) + + +def usage(warehouse_id: str | None = None) -> int: # Late import to avoid circular import (agents → state, but usage uses TOOL_SPECS for displays). from ucode.agents import TOOL_SPECS @@ -476,15 +529,11 @@ def usage() -> int: token = get_databricks_token(workspace, profile) with spinner("Discovering SQL warehouse..."): - resolved_http_path = discover_sql_warehouse_http_path(workspace, token, quiet=False) - - with spinner("Querying system.ai_gateway.usage..."): - columns, rows = run_usage_query( - workspace, - resolved_http_path, - token, - build_usage_report_query(), - ) + candidates = discover_sql_warehouses(workspace, token, warehouse_id=warehouse_id) + + resolved_http_path, columns, rows = run_query_on_first_working_warehouse( + workspace, token, candidates, build_usage_report_query() + ) records = parse_usage_rows(columns, rows) requester_name = find_requester_name(workspace, resolved_http_path, token, records) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 192e871..155dcae 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -27,6 +27,7 @@ build_skills_mcp_url, build_tool_base_url, classify_model_family, + discover_sql_warehouses, ensure_databricks_cli_version, ensure_pat_bearer, get_databricks_token, @@ -41,6 +42,22 @@ WS = "https://example.databricks.com" +class _FakeResponse: + """Minimal urlopen context manager returning a JSON body.""" + + def __init__(self, payload: dict): + self._body = json.dumps(payload).encode("utf-8") + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def read(self): + return self._body + + class TestWorkspaceHostname: def test_extracts_hostname(self): assert workspace_hostname(WS) == "example.databricks.com" @@ -2611,3 +2628,70 @@ def test_non_object_payload_is_no_spend(self, monkeypatch): spend, reason = resolve_current_budget_spend("https://ws", "token") assert spend is None assert "not a JSON object" in reason + + +class TestDiscoverSqlWarehouses: + def _payload(self, *entries: dict) -> dict: + return {"warehouses": list(entries)} + + def test_explicit_id_skips_discovery(self, monkeypatch): + def fail(*a, **k): + raise AssertionError("discovery should not be called") + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", fail) + assert discover_sql_warehouses(WS, "token", warehouse_id="abc") == [ + db_mod.SqlWarehouse("/sql/1.0/warehouses/abc", "abc", "REQUESTED") + ] + + def test_running_sorted_before_stopped(self, monkeypatch): + payload = self._payload( + {"id": "s1", "name": "stopped", "state": "STOPPED"}, + {"id": "r1", "name": "running", "state": "RUNNING"}, + ) + monkeypatch.setattr( + db_mod.urllib_request, "urlopen", lambda *a, **k: _FakeResponse(payload) + ) + result = discover_sql_warehouses(WS, "token") + assert [w.label for w in result] == ["running", "stopped"] + + def test_returns_all_candidates(self, monkeypatch): + payload = self._payload( + {"id": "a", "name": "A", "state": "RUNNING"}, + {"id": "b", "name": "B", "state": "RUNNING"}, + ) + monkeypatch.setattr( + db_mod.urllib_request, "urlopen", lambda *a, **k: _FakeResponse(payload) + ) + assert len(discover_sql_warehouses(WS, "token")) == 2 + + def test_skips_entries_without_id(self, monkeypatch): + payload = self._payload( + {"name": "no id", "state": "RUNNING"}, + {"id": "b", "name": "B", "state": "RUNNING"}, + ) + monkeypatch.setattr( + db_mod.urllib_request, "urlopen", lambda *a, **k: _FakeResponse(payload) + ) + assert [w.label for w in discover_sql_warehouses(WS, "token")] == ["B"] + + def test_falls_back_to_id_as_label(self, monkeypatch): + payload = self._payload({"id": "abc", "state": "RUNNING"}) + monkeypatch.setattr( + db_mod.urllib_request, "urlopen", lambda *a, **k: _FakeResponse(payload) + ) + assert discover_sql_warehouses(WS, "token")[0].label == "abc" + + def test_empty_list_raises_with_flag_hint(self, monkeypatch): + monkeypatch.setattr( + db_mod.urllib_request, "urlopen", lambda *a, **k: _FakeResponse({"warehouses": []}) + ) + with pytest.raises(RuntimeError, match="--warehouse-id"): + discover_sql_warehouses(WS, "token") + + def test_only_unusable_entries_raises(self, monkeypatch): + payload = self._payload({"name": "no id", "state": "RUNNING"}) + monkeypatch.setattr( + db_mod.urllib_request, "urlopen", lambda *a, **k: _FakeResponse(payload) + ) + with pytest.raises(RuntimeError, match="No usable SQL warehouse"): + discover_sql_warehouses(WS, "token") diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 3f47d95..ebd2ab3 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -25,7 +25,7 @@ from ucode.databricks import ( build_shared_base_urls, build_tool_base_url, - discover_sql_warehouse_http_path, + discover_sql_warehouses, ensure_ai_gateway_v2, fetch_ai_gateway_claude_models, fetch_codex_models, @@ -247,10 +247,11 @@ def test_configure_shared_state_and_reload( class TestSqlWarehouseDiscovery: def test_discovers_http_path(self, e2e_workspace, e2e_token): try: - http_path = discover_sql_warehouse_http_path(e2e_workspace, e2e_token, quiet=True) + candidates = discover_sql_warehouses(e2e_workspace, e2e_token) except RuntimeError as exc: pytest.skip(f"No SQL warehouse available: {exc}") - assert http_path.startswith("/sql/1.0/warehouses/") + assert candidates + assert all(w.http_path.startswith("/sql/1.0/warehouses/") for w in candidates) # --------------------------------------------------------------------------- diff --git a/tests/test_usage.py b/tests/test_usage.py index e3f2037..d6ff2a3 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -2,10 +2,14 @@ from __future__ import annotations +import contextlib from datetime import date, datetime, timedelta from decimal import Decimal +import pytest + import ucode.usage as usage_mod +from ucode.databricks import SqlWarehouse from ucode.ui import label, value from ucode.usage import ( USAGE_BREAKDOWN_DAYS, @@ -24,6 +28,7 @@ parse_usage_rows, render_budget_lines, render_usage_summary, + run_query_on_first_working_warehouse, simplify_model_name, summarize_model_tokens, summarize_models, @@ -515,8 +520,8 @@ def fake_render_box_table(headers, table_rows, max_widths=None): monkeypatch.setattr(usage_mod, "get_databricks_token", lambda *args, **kwargs: "token") monkeypatch.setattr( usage_mod, - "discover_sql_warehouse_http_path", - lambda *args, **kwargs: "/sql/1.0/warehouses/abc", + "discover_sql_warehouses", + lambda *args, **kwargs: [SqlWarehouse("/sql/1.0/warehouses/abc", "wh", "RUNNING")], ) monkeypatch.setattr(usage_mod, "run_usage_query", lambda *args, **kwargs: (columns, rows)) monkeypatch.setattr( @@ -532,8 +537,132 @@ def fake_render_box_table(headers, table_rows, max_widths=None): assert "Codex · Last 7 Days" in headings assert "Claude Code · Last 7 Days" in headings assert all("Gemini" not in heading for heading in headings) - assert notes == [f"No usage for Claude Code in the last {USAGE_BREAKDOWN_DAYS} days."] + assert notes == [ + "Using SQL warehouse `wh` (RUNNING).", + f"No usage for Claude Code in the last {USAGE_BREAKDOWN_DAYS} days.", + ] assert len(rendered_tables) == 1 assert rendered_tables[0][0][2] == "100" assert "gemini" not in "\n".join(printed).lower() assert "900" not in "\n".join(printed) + + +class TestRunQueryOnFirstWorkingWarehouse: + _COLUMNS = ["requester_name"] + _ROWS = [("user@example.com",)] + + def _warehouses(self, *labels: str) -> list[SqlWarehouse]: + return [SqlWarehouse(f"/sql/1.0/warehouses/{label}", label, "RUNNING") for label in labels] + + def test_returns_first_working_warehouse(self, monkeypatch): + monkeypatch.setattr(usage_mod, "print_note", lambda *a: None) + monkeypatch.setattr( + usage_mod, "run_usage_query", lambda *a, **k: (self._COLUMNS, self._ROWS) + ) + http_path, columns, rows = run_query_on_first_working_warehouse( + "https://ws", "token", self._warehouses("a", "b"), "SELECT 1" + ) + assert http_path == "/sql/1.0/warehouses/a" + assert (columns, rows) == (self._COLUMNS, self._ROWS) + + def test_falls_through_to_next_warehouse(self, monkeypatch): + warnings: list[str] = [] + monkeypatch.setattr(usage_mod, "print_note", lambda *a: None) + monkeypatch.setattr(usage_mod, "print_warning", warnings.append) + attempted: list[str] = [] + + def flaky(workspace, http_path, token, query, on_connected=None): + attempted.append(http_path) + if http_path.endswith("dead"): + raise RuntimeError("ENDPOINT_NOT_FOUND") + return self._COLUMNS, self._ROWS + + monkeypatch.setattr(usage_mod, "run_usage_query", flaky) + http_path, _, _ = run_query_on_first_working_warehouse( + "https://ws", "token", self._warehouses("dead", "alive"), "SELECT 1" + ) + assert http_path == "/sql/1.0/warehouses/alive" + assert attempted == ["/sql/1.0/warehouses/dead", "/sql/1.0/warehouses/alive"] + assert len(warnings) == 1 + assert "dead" in warnings[0] + + def test_raises_last_error_when_all_fail(self, monkeypatch): + monkeypatch.setattr(usage_mod, "print_note", lambda *a: None) + monkeypatch.setattr(usage_mod, "print_warning", lambda *a: None) + + def always_fail(workspace, http_path, token, query, on_connected=None): + raise RuntimeError(f"boom {http_path[-1]}") + + monkeypatch.setattr(usage_mod, "run_usage_query", always_fail) + with pytest.raises(RuntimeError, match="boom b"): + run_query_on_first_working_warehouse( + "https://ws", "token", self._warehouses("a", "b"), "SELECT 1" + ) + + def test_raises_when_no_candidates(self, monkeypatch): + monkeypatch.setattr(usage_mod, "print_note", lambda *a: None) + with pytest.raises(RuntimeError, match="No SQL warehouse could run"): + run_query_on_first_working_warehouse("https://ws", "token", [], "SELECT 1") + + +class TestUsageWarehouseIdPassthrough: + def test_forwards_warehouse_id_to_discovery(self, monkeypatch): + captured = {} + + def fake_discover(workspace, token, *, warehouse_id=None): + captured["warehouse_id"] = warehouse_id + return [SqlWarehouse("/sql/1.0/warehouses/xyz", "xyz", "REQUESTED")] + + monkeypatch.setattr( + usage_mod, "load_state", lambda: {"workspace": "https://ws", "available_tools": []} + ) + monkeypatch.setattr(usage_mod, "ensure_databricks_auth", lambda *a, **k: None) + monkeypatch.setattr(usage_mod, "get_databricks_token", lambda *a, **k: "token") + monkeypatch.setattr(usage_mod, "discover_sql_warehouses", fake_discover) + monkeypatch.setattr(usage_mod, "run_usage_query", lambda *a, **k: (["c"], [])) + monkeypatch.setattr(usage_mod, "print_note", lambda *a: None) + monkeypatch.setattr(usage_mod, "console", type("C", (), {"print": lambda *a: None})()) + + assert usage(warehouse_id="xyz") == 0 + assert captured["warehouse_id"] == "xyz" + + +class TestQueryProgressMessage: + def _messages(self, monkeypatch, state: str, connect: bool) -> list[str]: + """Spinner messages rendered for a warehouse in `state`.""" + seen: list[str] = [] + + @contextlib.contextmanager + def fake_spinner(message): + seen.append(message() if callable(message) else message) + yield + seen.append(message() if callable(message) else message) + + def fake_query(workspace, http_path, token, query, on_connected=None): + if connect and on_connected is not None: + on_connected() + return ["c"], [] + + monkeypatch.setattr(usage_mod, "spinner", fake_spinner) + monkeypatch.setattr(usage_mod, "run_usage_query", fake_query) + usage_mod._query_with_progress( + "https://ws", "token", SqlWarehouse("/p", "wh", state), "SELECT 1" + ) + return seen + + def test_running_shows_query_message(self, monkeypatch): + assert self._messages(monkeypatch, "RUNNING", connect=True) == [ + usage_mod.QUERY_MESSAGE, + usage_mod.QUERY_MESSAGE, + ] + + def test_requested_shows_query_message(self, monkeypatch): + # An explicit --warehouse-id; its real state was never looked up. + assert self._messages(monkeypatch, "REQUESTED", connect=True)[0] == usage_mod.QUERY_MESSAGE + + def test_stopped_starts_with_startup_message(self, monkeypatch): + assert self._messages(monkeypatch, "STOPPED", connect=False)[0] == usage_mod.STARTUP_MESSAGE + + def test_stopped_switches_to_query_once_connected(self, monkeypatch): + seen = self._messages(monkeypatch, "STOPPED", connect=True) + assert seen == [usage_mod.STARTUP_MESSAGE, usage_mod.QUERY_MESSAGE]