eschutho commented on code in PR #42900:
URL: https://github.com/apache/superset/pull/42900#discussion_r4343292725
##########
superset/db_engine_specs/databricks.py:
##########
@@ -775,6 +776,11 @@ class
DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec):
parameters_schema = DatabricksPythonConnectorSchema()
+ # The driver only populates `cursor.active_command_id` once a statement has
+ # actually been executed, so the cancel id can't be captured up front like
+ # it can for engines where it's tied to the session rather than the query.
+ has_query_id_before_execute = False
Review Comment:
Thanks @rusackas, @sadpandajoe. Went with the full poll-and-cancel loop here
rather than scoping down (e1c44a49aa).
How it works on SEA connections
(`DatabricksPythonConnectorEngineSpec.execute_with_cursor`):
- Each statement is submitted with `cursor.execute_async()`, which publishes
the SEA statement id before returning. That id is recorded as the query's
cancel id while the statement is still running, so Stop's `cancel_query` (fresh
connection, by statement id) targets the in-flight statement.
- The worker then polls `cursor.is_query_pending()`, backing off from 0.1s
to `DB_POLL_INTERVAL_SECONDS`. On each poll it refreshes the query. If it sees
`STOPPED` or the early-cancel flag, it cancels through the live cursor
(`cursor.cancel()`). The same happens on `SoftTimeLimitExceeded`. This mirrors
Impala's `handle_cursor`.
- Once the statement is no longer running, its id is retired from `extra`.
On a multi-statement script, a Stop between blocks therefore never cancels the
finished block's id and reports success. In that gap `prepare_cancel_query`
stages `QUERY_EARLY_CANCEL_KEY`, the worker skips the remaining blocks, and a
statement being submitted at that moment is cancelled on its first poll.
- If the statement finishes normally, `get_async_execution_result()` loads
its result set, so `description`/`fetch*` behave as after a synchronous
`execute`.
- Thrift connections keep the synchronous path unchanged. The early-cancel
flag is only staged for queries that ran through the SEA loop, so a Thrift Stop
still reports that cancellation isn't available instead of a false success.
Tests cover a finished statement, cancel mid-statement, a multi-statement
script where only the current statement's id is cancelled, a Stop between
statements, the soft time limit, and Thrift fallback.
One trade-off to flag: when a statement *fails* on the warehouse, the
driver's public `get_async_execution_result()` raises a generic "Operation
status FAILED" error rather than the server's error message that the
synchronous path surfaced. I haven't found a public driver accessor for that
detail. Happy to take suggestions there.
--
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]
