Skip to content

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

Open
andy-xu-db wants to merge 5 commits into
mainfrom
andy/usage-warehouse-override
Open

Add --warehouse-id and fall through unusable SQL warehouses#275
andy-xu-db wants to merge 5 commits into
mainfrom
andy/usage-warehouse-override

Conversation

@andy-xu-db

@andy-xu-db andy-xu-db commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Problem

ucode usage picks the first RUNNING SQL warehouse and aborts if it fails. On workspaces where a warehouse reports RUNNING but refuses connections, the command is unusable. For example:

• Using SQL warehouse ⚡ Reyden Warehouse (RUNNING).
ERROR Usage query failed: ENDPOINT_NOT_FOUND: SQL warehouse (…) does not exist at all in the database

There was no way to choose a different one. The error text told users to "pass --http-path", a flag that never existed.

Change

  • discover_sql_warehouses() returns all usable candidates, RUNNING first, instead of just one.
  • usage tries each in turn, warns on failure, and continues — so one bad warehouse no longer kills the report.
  • New --warehouse-id flag skips discovery entirely.

Testing

11 new unit tests; 1204 passed, ruff clean.

Verified on eng-ml-inference.staging, which has the broken warehouse above: auto-discovery now warns and falls through to a working one (exit 0, where main exits 1), and --warehouse-id pins directly.

@andy-xu-db
andy-xu-db marked this pull request as ready for review August 5, 2026 18:32
Comment thread src/ucode/usage.py
with spinner("Discovering SQL warehouse..."):
resolved_http_path = discover_sql_warehouse_http_path(workspace, token, quiet=False)

with spinner("Querying system.ai_gateway.usage..."):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why was this removed? We should still display this message, otherwise the command will just appear as if its hanging while we query

Comment thread src/ucode/cli.py
@app.command("usage")
def usage_cmd() -> None:
def usage_cmd(
warehouse_id: Annotated[

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add this new field to the README.md

Comment thread src/ucode/databricks.py
) -> str:
warehouse_id: str | None = None,
) -> list[SqlWarehouse]:
"""Candidate warehouses to run the usage query against, RUNNING ones first.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we skip the non-running candidate warehouses?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Currently if none of the warehouses are running, the code will start one up which works but can take minutes if not more. Should we allow that? If not then I'll remove the non-running candidate warehouses.

Comment thread src/ucode/usage.py
Comment on lines +456 to +462
try:
columns, rows = run_usage_query(workspace, warehouse.http_path, token, 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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we implement some sort of timeout here? how long does it usually take for it for return an error?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Errors are usually returned within a second. The main concern is if there aren't any running sql warehouses, which will cause the fallback to start up a new one. If we don't bother with that then I can add a timeout of a minute or two minutes.

AarushiShah-db
AarushiShah-db previously approved these changes Aug 6, 2026
@andy-xu-db
andy-xu-db enabled auto-merge (squash) August 6, 2026 17:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants