codeant-ai-for-open-source[bot] commented on code in PR #43627: URL: https://github.com/apache/superset/pull/43627#discussion_r3876989389
########## superset/tasks/query_cancel.py: ########## @@ -0,0 +1,131 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Engine-level cancellation for GTF chart-data query tasks. + +The chart-data execution path (``Database.get_df``) has none of SQL Lab's +cancel plumbing, so a GTF abort/timeout could only mark the task terminal while +the warehouse query kept running. This module adds the missing seam: + +- ``capture_cancel_id`` registers a cursor sink for the duration of a task's + query; ``notify_cursor`` (called from ``get_df`` before the blocking execute) + hands the live cursor to that sink so the task can capture an engine cancel id + via ``db_engine_spec.get_cancel_query_id`` — the same contract SQL Lab uses. +- ``cancel_chart_query`` kills the backend session over a *fresh* connection + (``db_engine_spec.cancel_query``), which unblocks the task's blocked ``get_df``. + +Only engines that return a cancel id before execution participate; others +capture nothing and are simply not cancellable (the abort still frees the task). +""" + +from __future__ import annotations + +import logging +from contextlib import closing, contextmanager +from contextvars import ContextVar +from typing import Any, Callable, Iterator, TYPE_CHECKING + +from superset.stats_logger import BaseStatsLogger + +if TYPE_CHECKING: + from flask import Flask + + from superset.models.core import Database + +logger = logging.getLogger(__name__) + +# Set by a chart-data task for the span of its query. When present, +# Database._execute_sql_with_mutation_and_logging hands the sink the live cursor +# before the blocking execute so the task can capture an engine cancel id. +# Absent (None) for every other get_df caller, so this is a no-op elsewhere. +_cancel_id_sink: ContextVar["Callable[[Any], None] | None"] = ContextVar( + "gtf_cancel_id_sink", default=None +) + + +@contextmanager +def capture_cancel_id(sink: "Callable[[Any], None]") -> Iterator[None]: + """Register a cursor sink for the duration of a chart-data query execution.""" + token = _cancel_id_sink.set(sink) + try: + yield + finally: + _cancel_id_sink.reset(token) + + +def notify_cursor(cursor: Any) -> None: + """Hand the live cursor to the active sink, if any. + + Called from ``get_df`` before the query executes. Best-effort: a capture + failure must never break query execution — it only forfeits cancellability. + """ + sink = _cancel_id_sink.get() + if sink is None: + return + try: + sink(cursor) + except Exception: # noqa: BLE001 pylint: disable=broad-except + logger.warning("Cancel-id capture failed", exc_info=True) + + +def cancel_chart_query( + database: "Database", cancel_query_id: str, app: "Flask | None" = None +) -> bool: + """Cancel a running chart-data warehouse query over a fresh connection. + + Runs ``db_engine_spec.cancel_query`` against a new connection to the same + database, terminating the backend session that the task's blocked ``get_df`` + is waiting on. Invoked from the task's abort handler (on the abort-listener + thread), so it opens its own connection rather than touching the busy one. + Best-effort and fully logged; the task's terminal transition is authoritative. + + :param database: the database the query is running against + :param cancel_query_id: engine cancel handle captured at query start + :param app: Flask app for config/DB access from the background thread + :returns: True if the engine reported the query cancelled + """ + from flask import current_app + + stats_logger: BaseStatsLogger = (app or current_app).config.get( + "STATS_LOGGER", BaseStatsLogger() + ) + spec = database.db_engine_spec + try: + with database.get_sqla_engine() as engine: + with closing(engine.raw_connection()) as conn: + with closing(conn.cursor()) as cursor: + # query is unused by the explicit-id specs (they cancel by + # the captured id alone); the chart path has no Query model. + cancelled = spec.cancel_query(cursor, None, cancel_query_id) # type: ignore[arg-type] Review Comment: **Suggestion:** The abort callback runs on a background thread with only the Flask application context, not the `override_user` context used by the executing task. Calling `database.get_sqla_engine()` therefore constructs a fresh OAuth2 connection without the task user's token, causing cancellation to fail for OAuth2-enabled databases and leaving the original warehouse query running. Preserve the task user context or pass the required OAuth2 credentials into the cancellation operation. [security] <details> <summary><b>Severity Level:</b> Major ⚠️</summary> ```mdx - ❌ OAuth2 chart-query cancellation can fail. - ⚠️ Original warehouse queries continue after task aborts. - ⚠️ Per-user database authentication is absent on cancellation. ``` </details> [](https://docs.codeant.ai/cli/resolve-pr-comments-skill) [](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=52118ca2453f49fb85959c1f310cd50f&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) [](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=52118ca2453f49fb85959c1f310cd50f&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) <details> <summary><b>Prompt for AI Agent 🤖 </b></summary> ```mdx This is a comment left during a code review. **Path:** superset/tasks/query_cancel.py **Line:** 107:112 **Comment:** *Security: The abort callback runs on a background thread with only the Flask application context, not the `override_user` context used by the executing task. Calling `database.get_sqla_engine()` therefore constructs a fresh OAuth2 connection without the task user's token, causing cancellation to fail for OAuth2-enabled databases and leaving the original warehouse query running. Preserve the task user context or pass the required OAuth2 credentials into the cancellation operation. Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise. Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix ``` </details> <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F43627&comment_hash=89164201e8c2c050d9edd128ce0cd73a8d0929e91124440b4756501e60fcccf9&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F43627&comment_hash=89164201e8c2c050d9edd128ce0cd73a8d0929e91124440b4756501e60fcccf9&reaction=dislike'>👎</a> ########## superset/tasks/async_queries.py: ########## @@ -120,10 +123,47 @@ def _get_dependency_cache_key() -> str: raise SupersetException("Prerequisite task did not publish a cache key") -# No timeout is set on these tasks yet: GTF enforces a timeout by aborting the -# task, but chart-data queries have no abort handler to cancel the underlying -# warehouse query, so a timeout would only mark the task failed while the query -# kept running. Restore a timeout once query cancellation is implemented. +@contextmanager +def _capture_query_cancellation(query_context: "QueryContext") -> Iterator[None]: + """Enable engine-level cancellation of this task's warehouse query. + + Captures the engine cancel id off the live cursor (for engines that expose + one before execution) and registers an abort handler that kills the backend + session over a fresh connection, unblocking the task's ``get_df``. Engines + without cancel support capture nothing and the task stays non-abortable, so + an abort/timeout simply frees the task without killing the (uncancellable) + query — matching the pre-cancellation behavior for those engines. + """ + database = getattr(query_context.datasource, "database", None) + if database is None: + yield + return + + ctx = get_context() + app = current_app._get_current_object() # noqa: SLF001 + captured = False + + def _sink(cursor: Any) -> None: + nonlocal captured + if captured: + return + # query is unused by the explicit-id specs; the chart path has no Query. + cancel_id = database.db_engine_spec.get_cancel_query_id(cursor, None) Review Comment: **Suggestion:** The engine-spec cancellation contract requires a real SQL Lab `Query`, but this passes `None`. Engines such as Ocient dereference `query.id` during `cancel_query`, and engines such as Impala access `query.database`; the resulting exception is swallowed by the cancellation path, so the task is marked aborted while the warehouse query continues running. Provide the actual query object or restrict this path to engines whose cancellation methods do not require it. [api mismatch] <details> <summary><b>Severity Level:</b> Major ⚠️</summary> ```mdx - ❌ Ocient query cancellation fails on every registered cancellable query. - ❌ Impala cancellation cannot access required database metadata. - ⚠️ Aborted tasks can leave warehouse queries consuming resources. ``` </details> [](https://docs.codeant.ai/cli/resolve-pr-comments-skill) [](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=4a03511e37aa42c3b63f1a9b739f93ef&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) [](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=4a03511e37aa42c3b63f1a9b739f93ef&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset) <details> <summary><b>Prompt for AI Agent 🤖 </b></summary> ```mdx This is a comment left during a code review. **Path:** superset/tasks/async_queries.py **Line:** 151:151 **Comment:** *Api Mismatch: The engine-spec cancellation contract requires a real SQL Lab `Query`, but this passes `None`. Engines such as Ocient dereference `query.id` during `cancel_query`, and engines such as Impala access `query.database`; the resulting exception is swallowed by the cancellation path, so the task is marked aborted while the warehouse query continues running. Provide the actual query object or restrict this path to engines whose cancellation methods do not require it. Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise. Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix ``` </details> <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F43627&comment_hash=0551a5d09dd1b79a50eb782dc2c99eda82f8b1d020da675d39a5712abfc0c8c2&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F43627&comment_hash=0551a5d09dd1b79a50eb782dc2c99eda82f8b1d020da675d39a5712abfc0c8c2&reaction=dislike'>👎</a> -- 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]
