feat(sql): bump execute_sql wait_timeout to 50s, expose get_statement_status + cancel_statement - #15
Open
limony-city wants to merge 1 commit into
Open
limony-city wants to merge 1 commit into
limony-city wants to merge 1 commit into
Conversation
…_status + cancel_statement
Before this change, `execute_sql` always submitted statements with `wait_timeout=10s`
(the API default) and there was no MCP tool to fetch the result for a statement
that the API returned in `PENDING` state. Long-running or cold-warehouse queries
became unreachable from the chat session — the only recourse was the Databricks
UI's query history or a raw curl against `/api/2.0/sql/statements/{id}`.
Changes:
- `sql.execute_statement`: accept a `wait_timeout` parameter, default `"50s"` —
the maximum synchronous wait supported by the Statement Execution API. This
catches the vast majority of queries inline.
- `execute_sql` MCP tool: forward the new `wait_timeout` arg so callers can
override (e.g. `"0s"` for fire-and-poll flows).
- New `get_statement_status` MCP tool: wraps the existing `sql.get_statement_status`
helper so agents can poll a PENDING statement to completion.
- New `cancel_statement` MCP tool: wraps `sql.cancel_statement` for symmetry —
letting agents free warehouse capacity when they abandon a query.
- Tests in `tests/test_sql.py` cover the new default, the override, the
endpoints the new tools call, and the tool-registration metadata.
limony-city
force-pushed
the
feat/sql-statement-polling
branch
from
June 7, 2026 11:01
6d4745f to
cea12b9
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The
execute_sqlMCP tool currently submits statements withwait_timeout=10s(the Databricks Statement Execution API default) and there is no companion MCP tool to fetch the result of a statement the API returned inPENDINGstate. Long-running or cold-warehouse queries become unreachable from the chat/agent session — recovery requires either the Databricks UI query history or a raw curl against/api/2.0/sql/statements/{id}.This PR fixes both problems with minimal surface change.
Changes
sql.execute_statement: newwait_timeoutparameter, default"50s"— the maximum synchronous wait supported by the Statement Execution API. Catches the vast majority of queries inline.execute_sqlMCP tool: forwards the newwait_timeoutarg so callers can override (e.g."0s"for fire-and-poll flows).get_statement_statusMCP tool: thin wrapper over the existingsql.get_statement_statushelper so agents can poll a PENDING statement to completion.cancel_statementMCP tool: wraps the existingsql.cancel_statementhelper for symmetry — lets agents free warehouse capacity when they abandon a query.Tests
New
tests/test_sql.pycovers:execute_statementdefaults towait_timeout="50s"and forwards explicit overrides.get_statement_status/cancel_statementcall the correct endpoints.execute_sqlwithwait_timeoutin its input schema,get_statement_status,cancel_statement) are registered in the server's tool list.```
tests/test_sql.py::test_execute_statement_defaults_to_50s_wait_timeout PASSED
tests/test_sql.py::test_execute_statement_honors_explicit_wait_timeout PASSED
tests/test_sql.py::test_get_statement_status_calls_correct_endpoint PASSED
tests/test_sql.py::test_cancel_statement_calls_correct_endpoint PASSED
tests/test_sql.py::test_new_sql_tools_are_registered PASSED
```
The five pre-existing failures on `master` (`test_server_structured`, `test_transcript`, `test_tool_metadata::test_list_tools_has_schemas`) are caused by API drift in the `mcp` SDK (`ToolManager.call_tool` no longer accepts `convert_result`; `Tool` no longer has `outputSchema`) and are unrelated to this PR.
Backwards compatibility
wait_timeoutparameter is optional with a default. Existing callers see no breakage in the API surface — only a longer (≤50s vs ≤10s) synchronous wait window, which is strictly better for usability and matches the Statement Execution API's documented maximum.