sadpandajoe commented on code in PR #42900:
URL: https://github.com/apache/superset/pull/42900#discussion_r4343560508
##########
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:
If a single status request fails (a 5xx, 429, or connection reset on the
poll), this treats the statement as finished: the loop returns,
`_retire_cancel_query_id` drops the cancel id, and
`get_async_execution_result()` then runs the driver's own blocking wait, which
never looks at the stop flag. A Stop clicked after that only stages the
early-cancel flag, `cancel_query` reports success, and the warehouse statement
keeps running (a soft time limit in that window skips the cancel too). The same
swallow also catches `SupersetTimeoutException` raised while the poll request
is in flight. Would it make sense to keep the statement marked pending and keep
polling on a transient error, and add a test where `is_query_pending` raises
once and the statement is still cancellable afterwards?
--
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]
