sadpandajoe commented on code in PR #43407:
URL: https://github.com/apache/superset/pull/43407#discussion_r3886730564
##########
superset/tasks/async_queries.py:
##########
@@ -16,126 +16,312 @@
# under the License.
from __future__ import annotations
-import dataclasses
import logging
-from typing import Any, TYPE_CHECKING
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Any, Iterator, TYPE_CHECKING
+from uuid import UUID
-from celery.exceptions import SoftTimeLimitExceeded
from flask import current_app
from flask_appbuilder.security.sqla.models import User
-from marshmallow import ValidationError
+from superset_core.tasks.types import TaskOptions, TaskScope
-from superset.charts.data.form_data import set_form_data
-from superset.charts.schemas import ChartDataQueryContextSchema
-from superset.exceptions import (
- SupersetErrorException,
- SupersetErrorsException,
+from superset.common.query_serialization import (
+ load_serialized_query,
+ serialize_query,
+ SerializedQuery,
)
+from superset.constants import CacheRegion
+from superset.exceptions import SupersetException
from superset.extensions import (
- async_query_manager,
- celery_app,
security_manager,
)
+from superset.tasks.ambient_context import get_context
+from superset.tasks.decorators import task
+from superset.tasks.query_cancel import (
+ cancel_chart_query,
+ capture_cancel_id,
+ capture_cancel_query_id,
+)
from superset.utils.core import override_user
-from superset.utils.error_sanitization import sanitize_error_dicts
if TYPE_CHECKING:
+ from superset_core.tasks.models import Task as CoreTask
+
from superset.common.query_context import QueryContext
+ from superset.common.query_object import QueryObject
+ from superset.models.tasks import Task
+ from superset.security.guest_token import GuestToken
logger = logging.getLogger(__name__)
-query_timeout = current_app.config[
- "SQLLAB_ASYNC_TIME_LIMIT_SEC"
-] # TODO: new config key
+# GTF task type for the chart-data fan-out. Each QueryObject runs as its own
SHARED
+# task keyed by its query_cache_key (safe cross-user dedup — the key encodes
+# RLS/impersonation). The client polls /api/v1/task/status_changes (filtered
to this
+# type) and aggregates the tasks' statuses itself; GTF owns completion
emission (there
+# is no coordinator task). The atomic unit is a QueryObject, not
chart-specific, so the
+# type is versioned to allow the serialization/execution contract to evolve.
+CHART_QUERY_TASK = "superset.query_object_v1"
+CACHE_KEY_PAYLOAD_KEY = "cache_key"
-def _create_query_context_from_form(form_data: dict[str, Any]) -> QueryContext:
- """
- Create the query context from the form data.
- :param form_data: The task form data
- :returns: The query context
- :raises ValidationError: If the request is incorrect
+def _resolve_user(user_id: int | None, guest_token: "GuestToken | None") ->
User:
+ """Resolve the acting user for an async chart-data task.
+
+ The GTF executor does not impersonate on its own, so each task establishes
the
+ request user itself, which is what RLS and database impersonation key off.
"""
+ if user_id:
+ return security_manager.get_user_by_id(user_id)
+ if guest_token:
+ return security_manager.get_guest_user_from_token(guest_token)
+ return security_manager.get_anonymous_user()
+
- try:
- return ChartDataQueryContextSchema().load(form_data)
- except KeyError as ex:
- raise ValidationError("Request is incorrect") from ex
+def _inject_contribution_totals(
+ query_obj: "QueryObject", totals_cache_key: str
+) -> None:
+ """Inject ``contribution_totals`` from the cached totals query into
``query_obj``.
+ A contribution query normalizes its metrics against column sums from a
separate
+ "totals" query. In the per-query task model the totals query runs as its
own
+ task (a ``depends_on`` prerequisite) and caches its dataframe; here we
read that
+ cached dataframe and inject the sums into this query's contribution
+ post-processing before it runs — the same result the synchronous
+ ``ensure_totals_available`` produces, but reading the cache the
prerequisite
+ populated instead of re-running the totals query. ``contribution_totals``
is
+ stripped from the cache key, so this affects only the result, not the key.
+ """
+ from superset.common.utils.query_cache_manager import QueryCacheManager
-def _load_user_from_job_metadata(job_metadata: dict[str, Any]) -> User:
- if user_id := job_metadata.get("user_id"):
- # logged in user
- user = security_manager.get_user_by_id(user_id)
- elif guest_token := job_metadata.get("guest_token"):
- # embedded guest user
- user = security_manager.get_guest_user_from_token(guest_token)
- del job_metadata["guest_token"]
- else:
- # default to anonymous user if no user is found
- user = security_manager.get_anonymous_user()
- return user
+ cache = QueryCacheManager.get(key=totals_cache_key,
region=CacheRegion.DATA)
+ if not cache.is_loaded or cache.df is None:
+ # The depends_on prerequisite guarantees the totals task succeeded and
wrote
+ # this cache entry, so a miss is unexpected (e.g. it was evicted
between the
+ # totals task finishing and this task reading). Fail loudly rather than
+ # caching a silently un-normalized result the client would then
re-request:
+ # this task's single query cannot reproduce the synchronous path's
+ # ensure_totals_available (it has no totals query to run).
+ raise SupersetException(
+ f"Contribution totals not found in cache under {totals_cache_key}"
+ )
+ df = cache.df
+ totals = {col: df[col].sum() for col in df.columns if df[col].dtype.kind
in "biufc"}
+ for post_processing in query_obj.post_processing or []:
+ if post_processing.get("operation") == "contribution":
+ post_processing.setdefault("options", {})["contribution_totals"] =
totals
-def _handle_soft_time_limit(
- job_metadata: dict[str, Any], ex: Exception, activity: str
-) -> None:
+def _get_dependency_cache_key() -> str:
+ """
+ Return the cache key published by a prerequisite chart-data task.
+
+ A dependent task reaches this point only after the scheduler's all-success
+ dependency gate has passed, so prerequisite payloads are expected to
contain
+ any output metadata their task body committed.
"""
- SoftTimeLimitExceeded is raised both by a genuine timeout and by a
- user-initiated cancel (revoke sends SIGUSR1). The cancel endpoint has
- already emitted the terminal event for the latter - it has to, since a task
- revoked while still queued never reaches this handler - so only a timeout
- is reported here, and without one the client would wait forever.
+ for payload in get_context().get_dependency_payloads():
+ cache_key = payload.get(CACHE_KEY_PAYLOAD_KEY)
+ if isinstance(cache_key, str):
+ return cache_key
+ raise SupersetException("Prerequisite task did not publish a cache key")
+
+
+@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.
"""
- if async_query_manager.is_job_cancelled(job_metadata["job_id"]):
- logger.info("Cancelled by the user while %s", activity)
+ database = getattr(query_context.datasource, "database", None)
+ if database is None:
+ yield
return
- logger.warning("A timeout occurred while %s, error: %s", activity, ex)
- async_query_manager.update_job(
- job_metadata,
- async_query_manager.STATUS_ERROR,
- errors=[{"message": f"A timeout occurred while {activity}"}],
- )
+ ctx = get_context()
+ app = current_app._get_current_object() # noqa: SLF001
+ captured = False
+
+ def _sink(cursor: Any) -> None:
+ nonlocal captured
+ if captured:
+ return
+ cancel_id = capture_cancel_query_id(database, cursor)
+ if cancel_id is None:
+ return
+ captured = True
+
+ # Persist the handle so the orphan reaper can cancel the query if this
+ # worker dies. Set it before on_abort: registering the first handler
+ # flushes the whole property cache, persisting the handle in that
write.
+ ctx.set_cancellation(database.id, cancel_id)
+ # Registering the first abort handler marks the task abortable and
starts
+ # the abort listener; on abort it cancels the query on a fresh
connection.
+ def _cancel() -> None:
+ cancel_chart_query(database, cancel_id, app)
-@celery_app.task(name="load_chart_data_into_cache",
soft_time_limit=query_timeout)
-def load_chart_data_into_cache(
- job_metadata: dict[str, Any],
- form_data: dict[str, Any],
+ ctx.on_abort(_cancel)
+
+ with capture_cancel_id(_sink):
+ yield
+
+
+@task(name=CHART_QUERY_TASK, scope=TaskScope.SHARED)
+def execute_chart_query(
+ serialized_query: SerializedQuery,
+ user_id: int | None = None,
+ guest_token: "GuestToken | None" = None,
+ requires_totals: bool = False,
) -> None:
- # pylint: disable=import-outside-toplevel
- from superset.commands.chart.data.get_data_command import ChartDataCommand
-
- with override_user(_load_user_from_job_metadata(job_metadata),
force=False):
- try:
- set_form_data(form_data)
- query_context = _create_query_context_from_form(form_data)
- command = ChartDataCommand(query_context)
- result = command.run(cache=True)
- cache_key = result["cache_key"]
- result_url = f"/api/v1/chart/data/{cache_key}"
- async_query_manager.update_job(
- job_metadata,
- async_query_manager.STATUS_DONE,
- result_url=result_url,
- )
- except SoftTimeLimitExceeded as ex:
- _handle_soft_time_limit(job_metadata, ex, "loading chart data")
- raise
- except Exception as ex:
- # Extract SIP-40 style errors when available
- if isinstance(ex, SupersetErrorException):
- errors = [dataclasses.asdict(ex.error)]
- elif isinstance(ex, SupersetErrorsException):
- errors = [dataclasses.asdict(error) for error in ex.errors]
- else:
- # Fallback for non-Superset exceptions
- error = str(ex.message if hasattr(ex, "message") else ex)
- errors = [{"message": error}]
- async_query_manager.update_job(
- job_metadata,
- async_query_manager.STATUS_ERROR,
- errors=sanitize_error_dicts(errors),
+ """Execute a single chart-data query and cache it under its
query_cache_key.
+
+ The atomic async unit: reconstruct the one query (canonical serialization),
+ optionally inject contribution totals from a prerequisite totals task,
then run
+ the existing per-query execution/caching path so a re-request reads the
same
+ DATA-cache entry. The query runs under ``_capture_query_cancellation`` so
an
+ abort/timeout can cancel it on engines that support query cancellation.
+ """
+ with override_user(_resolve_user(user_id, guest_token), force=False):
+ query_context = load_serialized_query(serialized_query)
+ # Floor the result-cache TTL: async caches the result for a follow-up
+ # request to read back (see get_cache_timeout).
+ query_context.is_async_execution = True
+ query_obj = query_context.queries[0]
+ if requires_totals:
+ _inject_contribution_totals(query_obj, _get_dependency_cache_key())
+ # Executes on cache miss and writes CacheRegion.DATA under
query_cache_key.
+ with _capture_query_cancellation(query_context):
+ result = query_context.get_df_payload_result(query_obj)
+ if cache_key := result.payload.get(CACHE_KEY_PAYLOAD_KEY):
+ # Write synchronously: a dependent contribution query reads this
+ # cache key via get_dependency_payloads once the DAG gate releases,
+ # so it must not sit in the throttle buffer.
+ get_context().update_task(
+ payload={CACHE_KEY_PAYLOAD_KEY: cache_key}, immediate=True
)
- raise
+
+
+def _query_task_cache_key(query_context: "QueryContext", index: int) -> str |
None:
+ """Compute a query's cache key exactly as its task will.
+
+ ``execute_chart_query`` validates each query before keying (see
+ ``get_df_payload_result``), so validate here too — otherwise the SHARED
task's
+ ``task_key`` (used for cross-user dedup) could diverge from the key the
task
+ actually caches under.
+ """
+ query_obj = query_context.queries[index]
+ query_obj.validate()
+ return query_context.query_cache_key(query_obj)
+
+
+def submit_chart_data_query_tasks(
+ query_context: "QueryContext",
+ user_id: int | None,
+) -> dict[str, Any]:
+ """Fan a chart-data request out into one GTF task per ``QueryObject``.
+
+ Each ``QueryObject`` runs as its own SHARED task keyed by its
``query_cache_key``
+ (safe cross-user dedup — the key encodes RLS/impersonation), writing the
per-query
+ DATA cache a later re-request reads back. A contribution query
``depends_on`` the
+ totals query's task and reads its cached result to normalize.
+
+ There is no coordinator task: the client polls
``/api/v1/task/status_changes`` and
+ aggregates the query tasks' own honest statuses itself (all ``SUCCESS`` →
re-issue
+ the request, now served entirely from the per-query cache; any terminal
non-success
+ → error). Completion is emitted per task by GTF (via the coordination
service), so
+ that is also what the websocket transport subscribes to.
+
+ Returns the HTTP 202 body ``{"task_ids": [...], "cursor": "..."}`` — the
query
+ tasks' UUIDs, in query order, plus the server-issued polling cursor. The
client
+ uses those values to poll through the GTF task API. Client aborts may
unsubscribe
+ from shared work or abort pending work; engine-level query cancellation is
outside
+ this chart async path.
+ """
+ guest_user = security_manager.get_current_guest_user_if_guest()
+ guest_token = guest_user.guest_token if guest_user else None
+
+ # Capture a status-poll cursor BEFORE any task is created. Because it
+ # predates every task's creation (and therefore every terminal transition),
+ # the client can poll `status_changes` from it and is guaranteed to observe
+ # each task's completion — closing the race where a task finishes before
the
+ # client's waiter/poll is established. Returned in the 202 (one small
value,
+ # unlike echoing every task id back).
+ poll_cursor = datetime.now()
Review Comment:
This cursor keeps microseconds even though `tasks.changed_on` is truncated
to seconds on MySQL, so a task that completes in the cursor's second can be
missed by the first poll and then permanently skipped by the advanced
watermark. Could this be floored to the storage precision, as
`get_statuses_changed_since` does for its generated cursors?
--
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]
