Add --warehouse-id and fall through unusable SQL warehouses - #275
Add --warehouse-id and fall through unusable SQL warehouses#275andy-xu-db wants to merge 5 commits into
--warehouse-id and fall through unusable SQL warehouses#275Conversation
| with spinner("Discovering SQL warehouse..."): | ||
| resolved_http_path = discover_sql_warehouse_http_path(workspace, token, quiet=False) | ||
|
|
||
| with spinner("Querying system.ai_gateway.usage..."): |
There was a problem hiding this comment.
why was this removed? We should still display this message, otherwise the command will just appear as if its hanging while we query
| @app.command("usage") | ||
| def usage_cmd() -> None: | ||
| def usage_cmd( | ||
| warehouse_id: Annotated[ |
There was a problem hiding this comment.
Can you add this new field to the README.md
| ) -> str: | ||
| warehouse_id: str | None = None, | ||
| ) -> list[SqlWarehouse]: | ||
| """Candidate warehouses to run the usage query against, RUNNING ones first. |
There was a problem hiding this comment.
can we skip the non-running candidate warehouses?
There was a problem hiding this comment.
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.
| 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 |
There was a problem hiding this comment.
should we implement some sort of timeout here? how long does it usually take for it for return an error?
There was a problem hiding this comment.
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.
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
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.