eschutho commented on code in PR #42900:
URL: https://github.com/apache/superset/pull/42900#discussion_r4343680530
##########
superset/db_engine_specs/databricks.py:
##########
@@ -931,6 +946,206 @@ def adjust_engine_params(
return uri, connect_args
+ @staticmethod
+ def _uses_sea(cursor: Any) -> bool:
+ """Whether ``cursor`` runs over the Statement Execution API (SEA)."""
+ session = getattr(getattr(cursor, "connection", None), "session", None)
+ return bool(getattr(session, "use_sea", False))
+
+ @classmethod
+ def execute_with_cursor(cls, cursor: Any, sql: str, query: Query) -> None:
+ """
+ Run a statement so a stop can cancel it while it is still executing.
+
+ On SEA connections the statement is submitted with ``execute_async``,
+ which publishes its statement id before returning. The id is recorded
+ as the query's cancel id, and the statement is polled with
+ ``is_query_pending`` until it finishes or a stop is requested, in which
+ case the live cursor cancels it. Once the statement is no longer
+ running its id is retired, so a stop between statements of a
+ multi-statement script never targets an already finished statement.
+
+ Thrift connections keep the synchronous execution path, as their
+ statements can't be cancelled from the stop request's connection.
+ """
+ if not cls._uses_sea(cursor):
+ super().execute_with_cursor(cursor, sql, query)
+ return
+
+ logger.debug("Query %d: Running query: %s", query.id, sql)
+ if not query.extra.get(cls.sea_async_execution_key):
+ query.set_extra_json_key(cls.sea_async_execution_key, True)
+ db.session.commit() # pylint: disable=consider-using-transaction
+
+ if cls.arraysize:
+ cursor.arraysize = cls.arraysize
+ try:
+ cursor.execute_async(sql)
+ except Exception as ex:
+ database = query.database
+ if (
+ not is_oauth2_retry_active()
+ and database.is_oauth2_enabled()
+ and cls.needs_oauth2(ex)
+ ):
+ cls.start_oauth2_dance(database)
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+
+ try:
+ if (cancel_query_id := cls.get_cancel_query_id(cursor, query)) is
not None:
+ query.set_extra_json_key(QUERY_CANCEL_KEY, cancel_query_id)
+ db.session.commit() # pylint:
disable=consider-using-transaction
+ stopped = cls._poll_statement(cursor, query)
+ finally:
+ cls._retire_cancel_query_id(query)
+
+ if stopped:
+ return
+
+ try:
+ # Loads the finished statement's result set onto the cursor, so it
+ # can be fetched exactly as after a synchronous ``execute``.
+ cursor.get_async_execution_result()
+ except Exception as ex:
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+
+ @classmethod
+ def _poll_statement(cls, cursor: Any, query: Query) -> bool:
+ """
+ Wait for the cursor's statement to finish, cancelling it on a stop.
+
+ :return: True if the query was stopped, False if the statement finished
+ """
+ # Short statements finish within moments, so polling starts with a
+ # short interval that backs off to the configured one.
+ sleep_interval = 0.1
+ max_interval = app.config["DB_POLL_INTERVAL_SECONDS"].get(cls.engine,
1)
+ try:
+ while True:
+ try:
+ pending = cursor.is_query_pending()
+ except SoftTimeLimitExceeded:
+ raise
+ except Exception: # pylint: disable=broad-except
+ # Retrieving the results reports the underlying failure.
+ logger.warning("Query %d: polling statement state failed",
query.id)
+ pending = False
Review Comment:
Good catch @sadpandajoe, thanks. Fixed in d890c58fb6.
- A failed `is_query_pending()` request (5xx/429/connection reset, etc.) no
longer counts as the statement finishing. It's treated as still running, so the
cancel id stays published. The loop still checks for a stop on that iteration
and cancels the live statement if one was requested, then polls again with the
usual backoff.
- So polling can't spin forever on a warehouse that keeps failing, there's a
cap: after `sea_max_poll_failures` (10) consecutive failures the statement is
cancelled and the error is raised. A successful poll resets the counter.
- `SupersetTimeoutException` is no longer swallowed. Raised during a poll,
it cancels the statement and propagates, same as `SoftTimeLimitExceeded`.
New tests:
- `is_query_pending` raises once, the id is still published on the following
polls, and a later Stop cancels the live statement.
- A transient failure followed by normal completion still loads results.
- Persistent failures cancel the statement and raise.
- A timeout during a poll cancels the statement.
--
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]
