eschutho opened a new pull request, #42900:
URL: https://github.com/apache/superset/pull/42900
### 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 implemented
`get_cancel_query_id`/`cancel_query`, so a stopped query keeps running on the
warehouse until it finishes or times out on its own.
### FIX
- Added `get_cancel_query_id`/`cancel_query` to
`DatabricksPythonConnectorEngineSpec`, gated on the connection using
Databricks' Statement Execution API (SEA) backend (`use_sea=True` in connection
params).
- The driver's only cancellation primitive is `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 (see
`sql_lab.cancel_query`), so the identifier has to be serializable and reusable
elsewhere.
- On a SEA connection, that identifier is a plain statement id, which the
driver can turn back into a valid cancel target on any cursor via
`CommandId.from_sea_statement_id(...)`. That's what's captured and replayed
here.
- On the default Thrift backend, the equivalent identifier requires an
operation *secret* that the driver only holds on the live `CommandId` object of
the executing cursor — it's never exposed through `cursor.query_id` or any
other public/documented accessor. Rather than reconstructing that secret from
private driver internals, `get_cancel_query_id` returns `None` for Thrift
connections, so "stop query" fails explicitly instead of silently doing nothing.
- Added `has_query_id_before_execute = False` so the cancel id is captured
*after* the statement executes — the driver only populates
`cursor.active_command_id` once a command has actually run.
- `cancel_query` lets any error from the fresh cancel cursor propagate
instead of swallowing it into a generic `False`/"could not cancel" result,
since silently reporting a failed cancel as if it were handled would be
misleading.
### NOT CHANGED
- **Databricks Interactive Cluster (Hive-based) connections** — already work
via `has_implicit_cancel` (the live cursor is closed to cancel server-side);
out of scope here.
- **ODBC (SQL Endpoints)** — pyodbc cursors expose `.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.
- **Legacy `databricks-dbapi` connector** — already flagged in this file as
legacy ("Use Python Connector for new deployments"); not worth building cancel
support for.
- **Thrift backend (the default) on the Python Connector** — see above; this
PR only covers SEA-enabled connections. Connections using the default
configuration will see no behavior change.
### TESTING INSTRUCTIONS
Added unit tests in `tests/unit_tests/db_engine_specs/test_databricks.py`
covering:
- `get_cancel_query_id` returns the SEA statement id when `use_sea=True`
- `get_cancel_query_id` returns `None` when no command has executed yet
- `get_cancel_query_id` returns `None` on the default (non-SEA) backend —
this is intended behavior, not a gap
- `cancel_query` reconstructs the SEA command id and delegates to the
driver's `cursor.cancel()`
- `cancel_query` rejects a malformed cancel id without touching the cursor
- `cancel_query` propagates errors from the cancel attempt instead of
swallowing them
Ran `pytest tests/unit_tests/db_engine_specs/test_databricks.py
tests/unit_tests/db_engine_specs/test_base.py` (138 passed) and the full
`tests/unit_tests/db_engine_specs/` suite (1265 passed, 9
pre-existing/unrelated failures in `test_bigquery.py` and `test_mysql.py`
reproduced identically on `master`, 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
- [ ] Has associated issue:
- [ ] Required feature flags:
- [ ] Changes UI
- [ ] Includes DB Migration (follow approval process in
[SIP-59](https://github.com/apache/superset/issues/13351))
- [ ] Migration is atomic, supports rollback & is backwards-compatible
- [ ] Confirm DB migration upgrade and downgrade tested
- [ ] Runtime estimates and downtime expectations provided
- [x] Introduces new feature or API
- [ ] Removes existing feature or API
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
