fix(databricks): support cancelling SQL Lab queries on SEA connections - #42900
fix(databricks): support cancelling SQL Lab queries on SEA connections#42900eschutho wants to merge 1 commit into
Conversation
Databricks queries run through the Python Connector previously had no "stop query" support. The connector's only cancellation mechanism is Cursor.cancel(), which needs a full command identifier from the same backend session; the default Thrift backend's identifier includes a secret that's never exposed via any public/documented accessor, so it can't be reconstructed on the fresh cursor Superset uses to issue cancellation. The newer, opt-in SEA (Statement Execution API) backend uses a plain statement id instead, which can be captured and reused safely, so cancellation is implemented for that case only. Co-Authored-By: Claude <noreply@anthropic.com>
dosubot
Bot
added
data:connect:databricks
Code Review Agent Run #670807Actionable Suggestions - 0Filtered by Review RulesBito filtered these suggestions based on rules created automatically for your feedback. Manage rules.
Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
✅ Deploy Preview for superset-docs-preview ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #42900 +/- ##
==========================================
- Coverage 65.73% 65.23% -0.50%
==========================================
Files 2843 2842 -1
Lines 162672 161006 -1666
Branches 37255 36958 -297
==========================================
- Hits 106935 105036 -1899
- Misses 53645 53947 +302
+ Partials 2092 2023 -69
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
1 participant
SUMMARY
SQL Lab's "stop query" action has no effect on Databricks connections using the Python Connector driver (
databricks-sql-connector), the officially recommended driver for new deployments. The engine spec never implementedget_cancel_query_id/cancel_query, so a stopped query keeps running on the warehouse until it finishes or times out on its own.FIX
get_cancel_query_id/cancel_querytoDatabricksPythonConnectorEngineSpec, gated on the connection using Databricks' Statement Execution API (SEA) backend (use_sea=Truein connection params).Cursor.cancel(), which needs a full command identifier tied to the session that ran the query. Superset always issues cancellation from a brand-new connection (seesql_lab.cancel_query), so the identifier has to be serializable and reusable elsewhere.CommandId.from_sea_statement_id(...). That's what's captured and replayed here.CommandIdobject of the executing cursor — it's never exposed throughcursor.query_idor any other public/documented accessor. Rather than reconstructing that secret from private driver internals,get_cancel_query_idreturnsNonefor Thrift connections, so "stop query" fails explicitly instead of silently doing nothing.has_query_id_before_execute = Falseso the cancel id is captured after the statement executes — the driver only populatescursor.active_command_idonce a command has actually run.cancel_querylets any error from the fresh cancel cursor propagate instead of swallowing it into a genericFalse/"could not cancel" result, since silently reporting a failed cancel as if it were handled would be misleading.NOT CHANGED
has_implicit_cancel(the live cursor is closed to cancel server-side); out of scope here..cancel()(SQLCancel), but like the Thrift case above, Superset's cancel flow uses a fresh cursor/connection with no reference to the one running the query, so it's unclear this would actually cancel anything. Left unimplemented pending real-world verification.databricks-dbapiconnector — already flagged in this file as legacy ("Use Python Connector for new deployments"); not worth building cancel support for.TESTING INSTRUCTIONS
Added unit tests in
tests/unit_tests/db_engine_specs/test_databricks.pycovering:get_cancel_query_idreturns the SEA statement id whenuse_sea=Trueget_cancel_query_idreturnsNonewhen no command has executed yetget_cancel_query_idreturnsNoneon the default (non-SEA) backend — this is intended behavior, not a gapcancel_queryreconstructs the SEA command id and delegates to the driver'scursor.cancel()cancel_queryrejects a malformed cancel id without touching the cursorcancel_querypropagates errors from the cancel attempt instead of swallowing themRan
pytest tests/unit_tests/db_engine_specs/test_databricks.py tests/unit_tests/db_engine_specs/test_base.py(138 passed) and the fulltests/unit_tests/db_engine_specs/suite (1265 passed, 9 pre-existing/unrelated failures intest_bigquery.pyandtest_mysql.pyreproduced identically onmaster, caused by local environment dependency versions, not this change).Ran
pre-commit(ruff, ruff-format, mypy, pylint, db-engine-spec-metadata) on the changed files — all passed.ADDITIONAL INFORMATION
🤖 Generated with Claude Code