rebenitez1802 commented on code in PR #45859:
URL: https://github.com/apache/superset/pull/45859#discussion_r4471743872
##########
superset/db_engine_specs/impala.py:
##########
@@ -197,20 +306,9 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None:
try:
status = cursor.status()
while status in unfinished_states:
- # End the read transaction first. Under REPEATABLE READ,
refreshing
- # inside one long transaction keeps returning the first
snapshot,
- # so a Stop committed by another request would never be seen.
- db.session.commit() # pylint:
disable=consider-using-transaction
- db.session.refresh(query)
- # Stop was requested: either before a cancel handle was
published
- # (early-cancel flag) or through SQL Lab's stop, which persists
- # STOPPED once cancel_query() succeeds.
- if (
- query.extra.get(QUERY_EARLY_CANCEL_KEY)
- or query.status == QueryStatus.STOPPED
- ):
+ if cls._stop_requested(query):
Review Comment:
Fixed by the master merge in 4c58f1c. #45845 changed `handle_cursor` to
catch `TIME_LIMIT_EXCEPTIONS`, which include `SupersetTimeoutException`, before
the broad handler. A timeout raised by the post-poll stop check now cancels the
operation and propagates, so the fetch never starts. 7c56930 adds
`test_timeout_during_the_post_poll_stop_check_cancels_the_operation`, which
covers both the Celery and the synchronous limit. Its synchronous case fails
against the pre-merge handler.
##########
superset/db_engine_specs/impala.py:
##########
@@ -152,7 +171,97 @@ def fetch_data(cls, cursor: Any, limit: int | None = None)
-> list[tuple[Any, ..
cursor.fetchall()
except Error as ex:
raise cls.get_dbapi_mapped_exception(ex) from ex
- return super().fetch_data(cursor, limit)
+ operation = getattr(cursor, "_last_operation", None)
+ try:
+ query = _STOP_WATCHED_OPERATIONS.get(operation)
+ except TypeError:
+ # The cursor exposes no driver operation that could be watched.
+ query = None
+ if query is None:
+ return super().fetch_data(cursor, limit)
+ with cls._stop_checked_fetch(cursor, operation, query):
+ return super().fetch_data(cursor, limit)
+
+ @classmethod
+ @contextmanager
+ def _stop_checked_fetch(
+ cls, cursor: Any, operation: Any, query: Query
+ ) -> Iterator[None]:
+ """
+ Check for a committed stop before each result RPC of a SQL Lab query.
+
+ Impala reports a streaming SELECT as finished once its rows can be
+ fetched, so most of such a query runs while the worker reads results,
+ after ``handle_cursor`` stopped polling. The driver issues one fetch
RPC
+ at a time, and while the coordinator's ``FETCH_ROWS_TIMEOUT_MS`` is
+ positive each returns within it even when no rows are ready, so a stop
+ goes unseen for at most about that plus the poll interval. A stop is
applied
+ on the worker's own cursor and thread, as in ``handle_cursor``, so no
+ second connection or HTTP call is involved.
+ """
+ fetch = operation.fetch
+ # Read once: after a failed check the instance can be expired in a
+ # session that cannot load it, so reading it then would raise again.
+ query_id = query.id
+ interval = app.config["DB_POLL_INTERVAL_SECONDS"].get(cls.engine, 5)
+ last_check: float | None = None
+ checking = True
+
+ def fetch_unless_stopped(*args: Any, **kwargs: Any) -> Any:
+ nonlocal last_check, checking
+ now = time.monotonic()
+ if checking and (last_check is None or now - last_check >=
interval):
+ last_check = now
+ try:
+ stopped = cls._stop_requested(query)
+ except (SoftTimeLimitExceeded, SupersetTimeoutException):
+ raise
+ except Exception: # pylint: disable=broad-except
+ # A metadata hiccup must not fail a query nobody stopped,
+ # so the failed transaction must not outlive the check.
+ try:
+ db.session.rollback() # pylint:
disable=consider-using-transaction
+ except Exception: # pylint: disable=broad-except
+ logger.debug("Query %s: rollback failed", query_id)
+ logger.warning(
+ "Query %s: stop check failed while fetching", query_id
+ )
+ checking = False
Review Comment:
Intentional, as Amin said, and pinned by
`test_failed_stop_check_does_not_fail_the_fetch`. A failed check falls back to
master's behaviour for the rest of that fetch, without failing a query nobody
stopped. A Stop committed later still ends the row STOPPED, and the results are
discarded once the fetch returns. What is lost is only the early cancel of the
Impala operation.
--
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]
