Let users choose a usage SQL warehouse

src/ucode/usage.py

@@ -35,6 +35,7 @@
     print_heading,
     print_note,
     print_warning,
+    prompt_for_selection,
     prompt_yes_no_default,
     render_box_table,
     spinner,
@@ -503,7 +504,13 @@ class ToolUsageTotals(NamedTuple):
     cost: Decimal | None
 
 
-TOOL_MODEL_TABLE_HEADERS = ["Model", "Requests", "Input (incl. cache)", "Output", "Cost (USD)"]
+TOOL_MODEL_TABLE_HEADERS = [
+    "Model",
+    "Requests",
+    "Input (incl. cache)",
+    "Output",
+    "Est. Cost (USD)",
+]
 
 
 def aggregate_tool_model_usage(records: list[dict[str, object]], tool: str) -> list[ModelUsage]:
@@ -692,6 +699,49 @@ def run_query_on_first_working_warehouse(
     raise last_error or RuntimeError("No SQL warehouse could run the usage query.")
 
 
+def select_sql_warehouse(candidates: list[SqlWarehouse]) -> SqlWarehouse | None:
+    """Ask which discovered warehouse should run the detailed usage query."""
+    selected_path = prompt_for_selection(
+        "Select a SQL warehouse for the usage query:",
+        [
+            (warehouse.http_path, f"{warehouse.label} ({warehouse.state})")
+            for warehouse in candidates
+        ],
+    )
+    return next(
+        (warehouse for warehouse in candidates if warehouse.http_path == selected_path),
+        None,
+    )
+
+
+def select_and_run_usage_query(
+    workspace: str,
+    token: str,
+    candidates: list[SqlWarehouse],
+    query: str,
+) -> tuple[str, list[str], list[tuple]] | None:
+    """Let the user choose warehouses until one runs the query or they cancel.
+
+    A failed warehouse is removed before the picker is shown again so the user cannot
+    accidentally retry the same unusable option forever. Returns ``None`` when the picker
+    is cancelled and raises the final warehouse error when no candidates remain.
+    """
+    remaining = list(candidates)
+    last_error: RuntimeError | None = None
+    while remaining:
+        selected = select_sql_warehouse(remaining)
+        if selected is None:
+            return None
+        try:
+            return run_query_on_first_working_warehouse(workspace, token, [selected], query)
+        except RuntimeError as exc:
+            last_error = exc
+            remaining = [
+                warehouse for warehouse in remaining if warehouse.http_path != selected.http_path
+            ]
+    raise last_error or RuntimeError("No SQL warehouse could run the usage query.")
+
+
 def _query_with_progress(
     workspace: str,
     token: str,
@@ -746,10 +796,18 @@ def usage(warehouse_id: str | None = None) -> int:
 
     with spinner("Discovering SQL warehouse..."):
         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()
-    )
+    if warehouse_id is None:
+        query_result = select_and_run_usage_query(
+            workspace, token, candidates, build_usage_report_query()
+        )
+        if query_result is None:
+            print_note("Usage details cancelled.")
+            return 0
+        resolved_http_path, columns, rows = query_result
+    else:
+        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)
 
@@ -793,6 +851,6 @@ def usage(warehouse_id: str | None = None) -> int:
         console.print(f"{label('Requests:')} {value(f'{totals.requests:,}')}")
         console.print(f"{label('Total tokens:')} {value(f'{totals.tokens:,}')}")
         if totals.cost is not None:
-            console.print(f"{label('Cost (USD):')} {value(format_cost_usd(totals.cost))}")
+            console.print(f"{label('Est. Cost (USD):')} {value(format_cost_usd(totals.cost))}")
         console.print(render_box_table(TOOL_MODEL_TABLE_HEADERS, rows, max_widths=table_widths))
     return 0

tests/test_usage.py

@@ -12,6 +12,7 @@
 from ucode.databricks import SqlWarehouse
 from ucode.ui import label, value
 from ucode.usage import (
+    TOOL_MODEL_TABLE_HEADERS,
     USAGE_BREAKDOWN_DAYS,
     USAGE_SUMMARY_DAYS,
     ModelPrice,
@@ -33,6 +34,8 @@
     render_budget_lines,
     render_usage_summary,
     run_query_on_first_working_warehouse,
+    select_and_run_usage_query,
+    select_sql_warehouse,
     simplify_model_name,
     summarize_models,
     usage,
@@ -464,6 +467,9 @@ def test_totals_cost_none_when_nothing_priced(self):
         _, totals = build_tool_model_rows(records, "claude", self._lookup())
         assert totals.cost is None
 
+    def test_cost_header_is_explicitly_estimated(self):
+        assert TOOL_MODEL_TABLE_HEADERS[-1] == "Est. Cost (USD)"
+
 
 class TestRenderBudgetLines:
     def test_no_lines_when_unavailable(self):
@@ -662,6 +668,9 @@ def fake_render_box_table(headers, table_rows, max_widths=None):
             usage_mod, "resolve_current_budget_spend", lambda *args, **kwargs: (None, "disabled")
         )
         monkeypatch.setattr(usage_mod, "prompt_yes_no_default", lambda *args, **kwargs: True)
+        monkeypatch.setattr(
+            usage_mod, "prompt_for_selection", lambda prompt, options: options[0][0]
+        )
         monkeypatch.setattr(
             usage_mod, "fetch_external_model_prices", lambda *args, **kwargs: ([], "disabled")
         )
@@ -688,6 +697,53 @@ def fake_render_box_table(headers, table_rows, max_widths=None):
         assert "gemini" not in "\n".join(printed).lower()
         assert "900" not in "\n".join(printed)
 
+    def test_queries_only_the_selected_warehouse(self, monkeypatch):
+        queried_paths: list[str] = []
+        candidates = [
+            SqlWarehouse("/sql/1.0/warehouses/first", "First", "RUNNING"),
+            SqlWarehouse("/sql/1.0/warehouses/second", "Second", "STOPPED"),
+        ]
+
+        monkeypatch.setattr(
+            usage_mod,
+            "load_state",
+            lambda: {"workspace": "https://workspace", "available_tools": []},
+        )
+        monkeypatch.setattr(usage_mod, "ensure_databricks_auth", lambda *args, **kwargs: None)
+        monkeypatch.setattr(usage_mod, "get_databricks_token", lambda *args, **kwargs: "token")
+        monkeypatch.setattr(
+            usage_mod,
+            "resolve_current_budget_spend",
+            lambda *args, **kwargs: ((Decimal("1"), Decimal("10")), None),
+        )
+        monkeypatch.setattr(usage_mod, "prompt_yes_no_default", lambda *args, **kwargs: True)
+        monkeypatch.setattr(
+            usage_mod, "discover_sql_warehouses", lambda *args, **kwargs: candidates
+        )
+
+        def choose_second(prompt, options):
+            assert options == [
+                (candidates[0].http_path, "First (RUNNING)"),
+                (candidates[1].http_path, "Second (STOPPED)"),
+            ]
+            return candidates[1].http_path
+
+        monkeypatch.setattr(usage_mod, "prompt_for_selection", choose_second)
+
+        def fake_query(workspace, http_path, token, query, on_connected=None):
+            queried_paths.append(http_path)
+            return ["requester_name"], [("user@example.com",)]
+
+        monkeypatch.setattr(usage_mod, "run_usage_query", fake_query)
+        monkeypatch.setattr(
+            usage_mod, "fetch_external_model_prices", lambda *args, **kwargs: ([], None)
+        )
+        monkeypatch.setattr(usage_mod, "print_note", lambda *args: None)
+        monkeypatch.setattr(usage_mod, "console", type("C", (), {"print": lambda *args: None})())
+
+        assert usage() == 0
+        assert queried_paths == [candidates[1].http_path]
+
     def test_shows_budget_before_prompt_and_skips_sql_when_declined(self, monkeypatch):
         events: list[str] = []
 
@@ -730,6 +786,22 @@ def decline(prompt, *, default):
         assert usage() == 0
 
 
+class TestSelectSqlWarehouse:
+    def test_returns_selected_candidate(self, monkeypatch):
+        candidates = [
+            SqlWarehouse("/sql/1.0/warehouses/a", "Alpha", "RUNNING"),
+            SqlWarehouse("/sql/1.0/warehouses/b", "Beta", "STOPPED"),
+        ]
+        monkeypatch.setattr(
+            usage_mod, "prompt_for_selection", lambda prompt, options: candidates[1].http_path
+        )
+        assert select_sql_warehouse(candidates) == candidates[1]
+
+    def test_returns_none_when_cancelled(self, monkeypatch):
+        monkeypatch.setattr(usage_mod, "prompt_for_selection", lambda prompt, options: None)
+        assert select_sql_warehouse([SqlWarehouse("/path", "Warehouse", "RUNNING")]) is None
+
+
 class TestRunQueryOnFirstWorkingWarehouse:
     _COLUMNS = ["requester_name"]
     _ROWS = [("user@example.com",)]
@@ -788,6 +860,54 @@ def test_raises_when_no_candidates(self, monkeypatch):
             run_query_on_first_working_warehouse("https://ws", "token", [], "SELECT 1")
 
 
+class TestSelectAndRunUsageQuery:
+    _COLUMNS = ["requester_name"]
+    _ROWS = [("user@example.com",)]
+
+    def _warehouses(self) -> list[SqlWarehouse]:
+        return [
+            SqlWarehouse("/sql/1.0/warehouses/dead", "Dead", "RUNNING"),
+            SqlWarehouse("/sql/1.0/warehouses/alive", "Alive", "RUNNING"),
+        ]
+
+    def test_prompts_again_without_failed_warehouse(self, monkeypatch):
+        candidates = self._warehouses()
+        picker_options: list[list[str]] = []
+
+        def choose(remaining):
+            picker_options.append([warehouse.label for warehouse in remaining])
+            return remaining[0]
+
+        def query(workspace, http_path, token, query, on_connected=None):
+            if http_path.endswith("dead"):
+                raise RuntimeError("warehouse unavailable")
+            return self._COLUMNS, self._ROWS
+
+        monkeypatch.setattr(usage_mod, "select_sql_warehouse", choose)
+        monkeypatch.setattr(usage_mod, "run_usage_query", query)
+        monkeypatch.setattr(usage_mod, "print_note", lambda *a: None)
+        monkeypatch.setattr(usage_mod, "print_warning", lambda *a: None)
+
+        result = select_and_run_usage_query("https://ws", "token", candidates, "SELECT 1")
+
+        assert result == (candidates[1].http_path, self._COLUMNS, self._ROWS)
+        assert picker_options == [["Dead", "Alive"], ["Alive"]]
+
+    def test_can_cancel_after_a_warehouse_fails(self, monkeypatch):
+        candidates = self._warehouses()
+        choices = iter([candidates[0], None])
+        monkeypatch.setattr(usage_mod, "select_sql_warehouse", lambda remaining: next(choices))
+        monkeypatch.setattr(
+            usage_mod,
+            "run_usage_query",
+            lambda *a, **k: (_ for _ in ()).throw(RuntimeError("warehouse unavailable")),
+        )
+        monkeypatch.setattr(usage_mod, "print_note", lambda *a: None)
+        monkeypatch.setattr(usage_mod, "print_warning", lambda *a: None)
+
+        assert select_and_run_usage_query("https://ws", "token", candidates, "SELECT 1") is None
+
+
 class TestUsageWarehouseIdPassthrough:
     def test_forwards_warehouse_id_to_discovery(self, monkeypatch):
         captured = {}