sadpandajoe commented on code in PR #42900:
URL: https://github.com/apache/superset/pull/42900#discussion_r4449276981
##########
superset/db_engine_specs/databricks.py:
##########
@@ -1013,6 +1032,261 @@ 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:
+ try:
+ cls._raise_statement_error(cursor)
+ except Exception as statement_error: # pylint:
disable=broad-except
+ raise cls.get_dbapi_mapped_exception(
+ statement_error
+ ) from statement_error
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+
+ @staticmethod
+ def _raise_statement_error(cursor: Any) -> None:
+ """
+ Raise the warehouse's error for a statement that didn't succeed.
+
+ ``get_async_execution_result`` reports any unsuccessful statement with
+ a generic "Operation status FAILED" error. The synchronous ``execute``
+ instead raises the server's error code and message, through the SEA
+ backend's state check. The driver exposes no public accessor for that
+ detail, so the same check is applied to the statement's latest status.
+ Returns without raising if the detail can't be retrieved, leaving the
+ caller to raise the driver's generic error.
+ """
+ backend = getattr(cursor, "backend", None)
+ poll_query = getattr(backend, "_poll_query", None)
+ check_state = getattr(
+ backend, "_check_command_not_in_failed_or_closed_state", None
+ )
+ command_id = getattr(cursor, "active_command_id", None)
+ if poll_query is None or check_state is None or command_id is None:
+ return
+ try:
+ status = poll_query(command_id).status
+ except SoftTimeLimitExceeded:
+ raise
+ except Exception: # pylint: disable=broad-except
+ return
+ check_state(status, command_id)
+
+ @classmethod
+ def _poll_statement(cls, cursor: Any, query: Query) -> bool:
+ """
+ Wait for the cursor's statement to finish, cancelling it on a stop.
+
+ A failed status request doesn't mean the statement finished, so it's
+ treated as still running and polled again; the statement stays
+ cancellable meanwhile. After ``sea_max_poll_failures`` consecutive
+ failures the statement is cancelled and the error is raised.
+
+ :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)
+ failures = 0
+ try:
+ while True:
+ try:
+ pending = cursor.is_query_pending()
+ failures = 0
+ except (SoftTimeLimitExceeded, SupersetTimeoutException):
+ raise
+ except Exception as ex: # pylint: disable=broad-except
+ failures += 1
+ if failures >= cls.sea_max_poll_failures:
+ cls._cancel_statement(cursor, query.id)
+ raise cls.get_dbapi_mapped_exception(ex) from ex
Review Comment:
If a user's OAuth token expires after `execute_async()` has been accepted,
the status poll starts failing with auth errors; the loop retries ten times and
then raises the mapped exception without ever calling `needs_oauth2()` /
`start_oauth2_dance()`. Previously the polling ran inside
`BaseEngineSpec.execute()`, whose handler starts the OAuth flow, so the user
was sent to re-authorize; now they get a generic failed query. The
`get_async_execution_result()` failure path at lines 1094-1102 has the same
gap. Could the OAuth check used for submission also cover the poll and
result-loading failures, with a test that drives an auth error through
`execute_with_cursor` for both?
--
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]
