codeant-ai-for-open-source[bot] commented on code in PR #45859:
URL: https://github.com/apache/superset/pull/45859#discussion_r4471748994
##########
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:
✅ **Customized review instruction saved!**
**Instruction:**
> In Impala stop-check logic, do not flag disabling further stop checks
after a metadata failure for the current fetch; a later Stop still marks the
query STOPPED and its results are discarded, with only early cancellation of
the Impala operation lost.
**Applied to:**
- `superset/db_engine_specs/impala.py`
---
💡 *To manage or update this instruction, visit: [CodeAnt AI
Settings](https://app.codeant.ai/org/settings/learnings)*
--
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]
