codeant-ai-for-open-source[bot] commented on code in PR #45859:
URL: https://github.com/apache/superset/pull/45859#discussion_r4444055672
##########
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:**
> Do not flag disabling subsequent stop checks after a failed stop check in
the Impala fetch flow; the fetch intentionally completes and discards results
if the query is stopped.
**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]
