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,
@@ -692,6 +693,21 @@ 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 _query_with_progress(
     workspace: str,
     token: str,
@@ -746,6 +762,12 @@ def usage(warehouse_id: str | None = None) -> int:
 
     with spinner("Discovering SQL warehouse..."):
         candidates = discover_sql_warehouses(workspace, token, warehouse_id=warehouse_id)
+    if warehouse_id is None:
+        selected = select_sql_warehouse(candidates)
+        if selected is None:
+            print_note("Usage details cancelled.")
+            return 0
+        candidates = [selected]
 
     resolved_http_path, columns, rows = run_query_on_first_working_warehouse(
         workspace, token, candidates, build_usage_report_query()

tests/test_usage.py

@@ -33,6 +33,7 @@
     render_budget_lines,
     render_usage_summary,
     run_query_on_first_working_warehouse,
+    select_sql_warehouse,
     simplify_model_name,
     summarize_models,
     usage,
@@ -662,6 +663,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 +692,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 +781,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",)]