Skip to content
277 changes: 243 additions & 34 deletions superset/common/query_context_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,10 @@
from superset.daos.chart import ChartDAO
from superset.dataframe import df_to_records
from superset.exceptions import (
CacheLoadError,
QueryObjectValidationError,
SupersetException,
SupersetSecurityException,
SupersetTemplateException,
)
from superset.explorables.base import Explorable
Expand All @@ -63,7 +65,6 @@
get_column_name,
get_column_names_from_columns,
get_column_names_from_metrics,
get_user_id,
is_adhoc_column,
is_adhoc_metric,
)
Expand Down Expand Up @@ -257,11 +258,24 @@ def get_df_payload_result(
query_obj.validate()

cache_key = self.query_cache_key(query_obj)
annotation_key = self.annotation_cache_key(query_obj)
timeout = self.get_cache_timeout()
force_query = (
self._resolve_forced_query(query_obj, cache_key)
or timeout == CACHE_DISABLED_TIMEOUT
)
# Resolved separately from the dataframe's force_query, against its own
# annotation_key. A forced GTF refresh's nonce marker is per (nonce,
# cache_key): the dataframe's marker is keyed on cache_key, so once it's
# set, force_query goes False for every subsequent request carrying that
# nonce, including a request for a *different* annotation access scope
# (its own annotation_key) that has never actually forced-refreshed its
# own entry. Reusing the dataframe's force_query there would read that
# scope's existing annotation cache entry instead of refreshing it.
annotation_force_query = (
self._resolve_forced_query(query_obj, annotation_key)
or timeout == CACHE_DISABLED_TIMEOUT
)
query_planning_ns = max(0, time.perf_counter_ns() - query_planning_start_ns)

cache_resolution_start_ns = time.perf_counter_ns()
Expand Down Expand Up @@ -307,7 +321,6 @@ def get_df_payload_result(
)

query_result = self.get_query_result(query_obj)
annotation_data = self.get_annotation_data(query_obj)
except QueryObjectValidationError as ex:
cache.error_message = str(ex)
cache.status = QueryStatus.FAILED
Expand All @@ -320,7 +333,6 @@ def get_df_payload_result(
cache.set_query_result(
key=cache_key,
query_result=query_result,
annotation_data=annotation_data,
force_query=force_query,
timeout=self.get_cache_timeout(),
datasource_uid=self._qc_datasource.uid,
Expand All @@ -331,6 +343,26 @@ def get_df_payload_result(
# nonce reads the freshly-cached result instead of recomputing it.
self._mark_force_executed(query_obj, cache_key, cache.result_persisted)

# Annotation data is fetched per requesting user (and, for chart-backed
# layers, scoped by the referenced chart datasource's RLS), so it is
# resolved and cached under its own entry — independent of whether the
# (shareable) dataframe above was a hit or a miss — rather than forcing
# every viewer of the same chart onto their own full dataframe copy.
annotation_data: dict[str, Any] = {}
if query_obj and annotation_key and cache.status != QueryStatus.FAILED:
try:
annotation_data = self._get_annotation_data_cached(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two viewers with different annotation RLS scopes can now join the same dataframe-keyed async task, but it only warms the executing viewer's annotation entry; the other viewer's post-success synchronous read-back then runs the annotation query in the HTTP request and can time out even though the task succeeded. Should each subscriber's scoped annotation data be populated asynchronously before its task completion is reported?

query_obj=query_obj,
cache_key=annotation_key,
force_query=annotation_force_query,
force_cached=force_cached,
timeout=self.get_cache_timeout(),
datasource_uid=self._qc_datasource.uid,
)
except QueryObjectValidationError as ex:
cache.error_message = str(ex)
cache.status = QueryStatus.FAILED

payload_assembly_start_ns = time.perf_counter_ns()
# the N-dimensional DataFrame has converted into flat DataFrame
# by `flatten operator`, "comma" in the column is escaped by `escape_separator`
Expand Down Expand Up @@ -400,7 +432,7 @@ def get_df_payload_result(
"applied_template_filters": cache.applied_template_filters,
"applied_filter_columns": cache.applied_filter_columns,
"rejected_filter_columns": cache.rejected_filter_columns,
"annotation_data": cache.annotation_data,
"annotation_data": annotation_data,
"error": cache.error_message,
"is_cached": cache.is_cached,
"query": cache.query,
Expand All @@ -426,15 +458,15 @@ def get_df_payload_result(
def query_cache_key(self, query_obj: QueryObject, **kwargs: Any) -> str | None:
"""
Returns a QueryObject cache key for objects in self.queries

This key covers the dataframe alone. It intentionally does not bind
the requesting user's identity, so distinct viewers of the same chart
share one cache entry. See :meth:`annotation_cache_key` for the
separate, user-scoped key covering annotation-layer data.
"""
datasource = self._qc_datasource
extra_cache_keys = datasource.get_extra_cache_keys(query_obj.to_dict())

# Annotation data is cached on the same entry as the dataframe, so the
# key must also bind the annotation sources' security context.
if query_obj and query_obj.annotation_layers:
kwargs["annotation_context"] = self._annotation_cache_context(query_obj)

cache_key = (
query_obj.cache_key(
datasource=datasource.uid,
Expand All @@ -448,31 +480,186 @@ def query_cache_key(self, query_obj: QueryObject, **kwargs: Any) -> str | None:
)
return cache_key

def _annotation_cache_context(self, query_obj: QueryObject) -> dict[str, Any]:
def annotation_cache_key(self, query_obj: QueryObject) -> str | None:
"""
Cache key for this query's annotation-layer payload, or ``None`` when
the query has no annotation layers.

Annotation payloads are fetched under the requesting user's access
scope, which is a stricter security requirement than the dataframe
itself has. Keying them separately from :meth:`query_cache_key` keeps
that scoping from forcing every distinct viewer of an annotated chart
onto their own full copy of the (potentially much larger) shared
dataframe: users with the same access scope share this key too.
"""
Cache-key material binding cached annotation data to its security
context.
if not query_obj or not query_obj.annotation_layers:
return None
return self.query_cache_key(
query_obj, annotation_context=self._annotation_cache_context(query_obj)
)

Annotation payloads are fetched per requesting user and stored on the
same cache entry as the dataframe, so the key also binds the requesting
user and, for chart-backed layers, the RLS clauses of the referenced
chart's datasource.
def _annotation_cache_context(self, query_obj: QueryObject) -> dict[str, Any]:
"""
source_rls: dict[str, list[str] | None] = {}
Cache-key material binding annotation data to its security *scope* so
users with the same access share a cache entry and users with a
different scope — or no access — never read each other's data.

* NATIVE layers: the ``can_read`` permission on ``Annotation``, the
only user-dependent dimension of these global records.
* Chart-backed (``line``/``table``) layers: see
:meth:`_annotation_source_scope`.
"""
context: dict[str, Any] = {}

if any(
layer.get("sourceType") == "NATIVE" for layer in query_obj.annotation_layers
):
context["annotation_read"] = security_manager.can_access(
"can_read", "Annotation"
)

source_scope: dict[str, Any] = {}
for layer in query_obj.annotation_layers:
if layer.get("sourceType") not in ("line", "table"):
continue
layer_value = layer.get("value")
source_scope[str(layer_value)] = self._annotation_source_scope(layer)
if source_scope:
context["source_scope"] = source_scope

return context

def _annotation_source_scope(self, layer: dict[str, Any]) -> dict[str, Any]:
"""
Access and data-identity cache-key material for one chart-backed
annotation layer.

``access`` keeps a user denied the referenced chart's datasource from
reading an authorized user's cached payload. When the chart has a
saved query context, this runs the *same* authorization path
:meth:`get_viz_annotation_data` executes
(``QueryContext.raise_for_access``) rather than the coarser
:meth:`SecurityManager.can_access_datasource` — a requester whose
access comes from a dashboard/viewer-promiscuous-mode bypass (which
depends on the chart's own saved ``form_data``, e.g. its
``slice_id``/``dashboardId``) would otherwise still fail that coarser,
context-free check and collapse onto the same denied scope as a
genuinely unauthorized requester, letting the latter read the
former's cached payload. ``data_key`` is the annotation chart's own
query cache key(s) — derived from the same, override-applied query
objects actually executed (see :meth:`_apply_annotation_overrides`),
so it captures the datasource version, RLS clauses, and any per-user
Jinja/virtual-dataset RLS material exactly as the live fetch would,
including material an override only introduces at a finer grain.
Reusing this logic avoids re-deriving it and automatically inherits
any future correctness fixes made there.
"""
layer_value = layer.get("value")
datasource = None
try:
chart = (
ChartDAO.find_by_id(layer_value) if layer_value is not None else None
)
annotation_datasource = chart.datasource if chart else None
source_rls[str(layer.get("value"))] = (
security_manager.get_rls_cache_key(annotation_datasource)
if annotation_datasource
else None
# resolved_datasource, not datasource: the latter is pinned to
# table-backed datasources and resolves to None for a
# semantic-view-backed chart, which would otherwise collapse
# every requester onto the same {access: None, data_key: None}
# scope below regardless of their actual access.
datasource = chart.resolved_datasource if chart else None
if chart is None or datasource is None:
return {"access": None, "data_key": None}

annotation_query_context = chart.get_query_context()
if annotation_query_context is not None:
self._apply_annotation_overrides(annotation_query_context, layer)
try:
annotation_query_context.raise_for_access()
access: Any = True
except SupersetSecurityException:
access = False
data_key: Any = [
annotation_query_context.query_cache_key(query_object)
Comment thread
rusackas marked this conversation as resolved.
for query_object in annotation_query_context.queries
]
else:
# Fall back to the RLS-clause identity when the chart has no
# saved query context to key on.
access = security_manager.can_access_datasource(datasource)
data_key = security_manager.get_rls_cache_key(datasource)
except Exception: # noqa: BLE001 pylint: disable=broad-except
# Derivation can fail well beyond SupersetException: the DAO
# lookup and lazy ``datasource`` load are themselves live DB
# queries, and the RLS lookup / a virtual dataset's
# get_extra_cache_keys() renders Jinja, either of which can raise
# a driver/template error. None of that should ever 500 the whole
# chart-data request; fail closed instead so this scope can't
# silently dedupe onto a successfully-derived one.
logger.warning(
"Could not derive annotation cache key for chart %s; "
"falling back to a fail-closed scope",
layer_value,
exc_info=True,
)
try:
fallback_data_key = (
security_manager.get_rls_cache_key(datasource)
if datasource is not None
else None
)
except Exception: # noqa: BLE001 pylint: disable=broad-except
# The fallback's own lookup can fail the same way (e.g. the
# same DB outage that failed the primary derivation) -- don't
# let that escape either.
fallback_data_key = None
return {"access": False, "data_key": fallback_data_key}
return {"access": access, "data_key": data_key}

def _get_annotation_data_cached(
self,
query_obj: QueryObject,
cache_key: str,
force_query: bool,
force_cached: bool | None,
timeout: int | None,
datasource_uid: str | None,
) -> dict[str, Any]:
"""
Fetch this query's annotation-layer payload, cached under its own
(user/RLS-scoped) entry, separate from the shared dataframe cache.
"""
if not force_query:
try:
cached_value = cache_manager.data_cache.get(cache_key)
except Exception as ex: # noqa: BLE001 pylint: disable=broad-except
logger.warning(
"Error reading annotation cache: %s",
error_msg_from_exception(ex),
)
cached_value = None
if cached_value is not None:
current_app.config["STATS_LOGGER"].incr("loading_from_cache")
return cached_value.get("annotation_data", {})

if force_cached:
logger.warning(
"force_cached (annotation data): value not found for key %s",
cache_key,
)
return {"user_id": get_user_id(), "source_rls": source_rls}
raise CacheLoadError("Error loading annotation data from cache")

annotation_data = self.get_annotation_data(query_obj)
persisted = set_and_log_cache(
cache_manager.data_cache,
cache_key,
{"annotation_data": annotation_data},
timeout,
datasource_uid,
)
# Mirrors the dataframe path's own marker write: only if the fresh
# value actually persisted, so a follow-up request carrying the same
# nonce reads this entry instead of forcing another recompute.
self._mark_force_executed(query_obj, cache_key, persisted)
return annotation_data

def get_query_result(self, query_object: QueryObject) -> QueryResult:
"""
Expand Down Expand Up @@ -881,6 +1068,36 @@ def get_native_annotation_data(query_obj: QueryObject) -> dict[str, Any]:
annotation_data[layer_name] = result
return annotation_data

@staticmethod
def _apply_annotation_overrides(
query_context: QueryContext, annotation_layer: dict[str, Any]
) -> None:
"""
Apply a chart-backed annotation layer's per-request overrides (time
grain, time range) to its saved query context's queries, in place.

Shared by :meth:`get_viz_annotation_data` (which executes the
overridden query) and :meth:`_annotation_source_scope` (which derives
its cache key from it), so the key always reflects the query that
actually runs: an override can introduce per-user Jinja/RLS material
the saved, un-overridden query lacks (e.g. a template that only calls
``current_user_id()`` at a finer time grain), which deriving the key
from the un-overridden query would silently miss.
"""
if not (overrides := annotation_layer.get("overrides")):
return

if time_grain_sqla := overrides.get("time_grain_sqla"):
for query_object in query_context.queries:
query_object.extras["time_grain_sqla"] = time_grain_sqla

if time_range := overrides.get("time_range"):
from_dttm, to_dttm = get_since_until_from_time_range(time_range)

for query_object in query_context.queries:
query_object.from_dttm = from_dttm
query_object.to_dttm = to_dttm

@staticmethod
def get_viz_annotation_data( # noqa: C901
annotation_layer: dict[str, Any], force: bool
Expand Down Expand Up @@ -912,17 +1129,9 @@ def get_viz_annotation_data( # noqa: C901
)
)

if overrides := annotation_layer.get("overrides"):
if time_grain_sqla := overrides.get("time_grain_sqla"):
for query_object in query_context.queries:
query_object.extras["time_grain_sqla"] = time_grain_sqla

if time_range := overrides.get("time_range"):
from_dttm, to_dttm = get_since_until_from_time_range(time_range)

for query_object in query_context.queries:
query_object.from_dttm = from_dttm
query_object.to_dttm = to_dttm
QueryContextProcessor._apply_annotation_overrides(
query_context, annotation_layer
)

query_context.force = force
command = ChartDataCommand(query_context)
Expand Down
Loading
Loading