Add --warehouse-id and fall through unusable SQL warehouses

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 <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 |

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

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,41 +2755,45 @@ 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(
     workspace: str,
     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 [])]

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)
 

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")

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)
 
 
 # ---------------------------------------------------------------------------