diff --git a/superset/common/query_context_processor.py b/superset/common/query_context_processor.py index 5f98c82ad0b8..030088326b8a 100644 --- a/superset/common/query_context_processor.py +++ b/superset/common/query_context_processor.py @@ -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 @@ -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, ) @@ -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() @@ -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 @@ -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, @@ -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( + 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` @@ -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, @@ -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, @@ -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) + 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: """ @@ -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 @@ -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) diff --git a/tests/unit_tests/common/test_query_context_processor.py b/tests/unit_tests/common/test_query_context_processor.py index d6fea922b551..bd65272b86f1 100644 --- a/tests/unit_tests/common/test_query_context_processor.py +++ b/tests/unit_tests/common/test_query_context_processor.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +from collections.abc import Iterator from datetime import datetime, timedelta from typing import Any from unittest.mock import MagicMock, patch @@ -30,7 +31,7 @@ normalize_contribution_totals, QueryContextProcessor, ) -from superset.exceptions import QueryObjectValidationError +from superset.exceptions import CacheLoadError, QueryObjectValidationError from superset.utils import json as superset_json from superset.utils.core import GenericDataType from superset.utils.date_parser import get_past_or_future @@ -116,25 +117,259 @@ def processor(mock_query_context): return processor -def test_query_cache_key_binds_annotation_data_to_requesting_user(processor): - """The cache key for annotated queries must differ per requesting user.""" +def test_annotation_cache_key_binds_native_annotation_read_scope(processor) -> None: + """The annotation cache key for NATIVE layers must differ when the + requester's ``can_read`` (Annotation) access differs -- not who they are.""" query_obj = MagicMock() query_obj.annotation_layers = [{"sourceType": "NATIVE", "name": "a", "value": 1}] - with ( - patch( - "superset.common.query_context_processor.get_user_id", - side_effect=[1, 2], - ), - patch("superset.common.query_context_processor.security_manager"), - ): - processor.query_cache_key(query_obj) - processor.query_cache_key(query_obj) + # ``security_manager`` autodetects as an async spec under a bare + # ``patch()`` (its real object trips ``unittest.mock``'s coroutine + # inference), which would silently turn every attribute access into an + # ``AsyncMock`` returning a fresh unawaited coroutine per call -- always + # unequal to itself and never equal to a configured return value. Forcing + # ``new_callable=MagicMock`` keeps these synchronous, as the real object + # is. + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as security_manager: + security_manager.can_access.side_effect = [True, False] + processor.annotation_cache_key(query_obj) + processor.annotation_cache_key(query_obj) contexts = [ call.kwargs["annotation_context"] for call in query_obj.cache_key.call_args_list ] assert contexts[0] != contexts[1] +def test_annotation_cache_key_shares_across_same_access_scope() -> None: + """Two distinct requesters (separate processor/query-object instances, + standing in for two different requests) with identical access scope must + produce identical annotation-context material. Reusing a single + processor/query_obj across both calls (as this test previously did) + would pass trivially regardless of whether the key is scope-based or + identity-based, since nothing about "who's asking" would ever vary.""" + layer = {"sourceType": "NATIVE", "name": "a", "value": 1} + processor_a = QueryContextProcessor(MagicMock()) + processor_b = QueryContextProcessor(MagicMock()) + query_obj_a = MagicMock(annotation_layers=[layer]) + query_obj_b = MagicMock(annotation_layers=[layer]) + + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as security_manager: + security_manager.can_access.return_value = True + context_a = processor_a._annotation_cache_context(query_obj_a) + context_b = processor_b._annotation_cache_context(query_obj_b) + + assert context_a == context_b + + +def test_query_cache_key_does_not_bind_annotation_scope(processor) -> None: + """The dataframe cache key must stay shared across viewers of the same + chart, even when the query has annotation layers — only the separate + annotation cache key (see above) carries access-scope material.""" + query_obj = MagicMock() + query_obj.annotation_layers = [{"sourceType": "NATIVE", "name": "a", "value": 1}] + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ): + processor.query_cache_key(query_obj) + processor.query_cache_key(query_obj) + for call in query_obj.cache_key.call_args_list: + assert "annotation_context" not in call.kwargs + + +@pytest.fixture +def mock_annotation_chart() -> Iterator[MagicMock]: + """A found chart, wired as the referenced chart for + ``_annotation_source_scope`` tests -- factors out the repeated + ``ChartDAO.find_by_id`` patch those tests all need.""" + chart = MagicMock() + with patch( + "superset.common.query_context_processor.ChartDAO.find_by_id", + return_value=chart, + ): + yield chart + + +def test_annotation_source_scope_binds_datasource_access( + processor, mock_annotation_chart +) -> None: + """A chart-backed annotation layer's scope must differ when the + requester's access to the referenced datasource differs.""" + mock_annotation_chart.get_query_context.return_value = None + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as security_manager: + security_manager.can_access_datasource.side_effect = [True, False] + security_manager.get_rls_cache_key.return_value = [] + scope_a = processor._annotation_source_scope({"value": 1}) + scope_b = processor._annotation_source_scope({"value": 1}) + assert scope_a != scope_b + assert scope_a["access"] is True + assert scope_b["access"] is False + + +def test_annotation_source_scope_reuses_referenced_chart_cache_key( + processor, mock_annotation_chart +) -> None: + """When the referenced chart has a saved query context, its own cache + key(s) -- covering RLS and per-user Jinja/virtual-dataset material -- are + reused rather than re-derived.""" + mock_query_object = MagicMock() + mock_query_context = MagicMock() + mock_query_context.queries = [mock_query_object] + mock_query_context.query_cache_key.return_value = "referenced-chart-key" + mock_annotation_chart.get_query_context.return_value = mock_query_context + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as security_manager: + security_manager.can_access_datasource.return_value = True + scope = processor._annotation_source_scope({"value": 1}) + assert scope == {"access": True, "data_key": ["referenced-chart-key"]} + mock_query_context.query_cache_key.assert_called_once_with(mock_query_object) + + +def test_annotation_source_scope_uses_live_fetch_authorization( + processor, mock_annotation_chart +) -> None: + """When the referenced chart has a saved query context, ``access`` must + come from that context's own ``raise_for_access`` -- the same + authorization path the live fetch in ``get_viz_annotation_data`` uses -- + not the coarser, context-free ``can_access_datasource``. A requester + denied by ``can_access_datasource`` but granted via a bypass that depends + on the chart's own saved form_data (e.g. a dashboard/viewer-promiscuous + bypass) must get a scope distinct from one truly denied, or the latter + could read the former's cached payload.""" + from superset.exceptions import SupersetSecurityException + + mock_query_object = MagicMock() + mock_query_context = MagicMock() + mock_query_context.queries = [mock_query_object] + mock_query_context.query_cache_key.return_value = "referenced-chart-key" + mock_annotation_chart.get_query_context.return_value = mock_query_context + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as security_manager: + # Coarse proxy says "denied" for both -- the real authorization + # path must be what actually decides access. + security_manager.can_access_datasource.return_value = False + mock_query_context.raise_for_access.side_effect = [ + None, + SupersetSecurityException(MagicMock()), + ] + scope_a = processor._annotation_source_scope({"value": 1}) + scope_b = processor._annotation_source_scope({"value": 1}) + assert scope_a["access"] is True + assert scope_b["access"] is False + assert scope_a != scope_b + security_manager.can_access_datasource.assert_not_called() + + +def test_annotation_source_scope_applies_overrides_before_keying( + processor, mock_annotation_chart +) -> None: + """A time-grain/time-range override on the annotation layer must be + applied to the referenced chart's query objects *before* deriving the + cache key, mirroring ``get_viz_annotation_data`` exactly -- otherwise the + key can omit per-user Jinja/RLS material an override only introduces at a + finer grain.""" + mock_query_object = MagicMock() + mock_query_object.extras = {} + mock_query_context = MagicMock() + mock_query_context.queries = [mock_query_object] + mock_query_context.query_cache_key.return_value = "referenced-chart-key" + mock_annotation_chart.get_query_context.return_value = mock_query_context + layer = { + "value": 1, + "overrides": {"time_grain_sqla": "P1D", "time_range": "Last week"}, + } + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ): + processor._annotation_source_scope(layer) + assert mock_query_object.extras["time_grain_sqla"] == "P1D" + assert mock_query_object.from_dttm is not None + assert mock_query_object.to_dttm is not None + + +def test_annotation_source_scope_fails_closed_on_any_derivation_error( + processor, mock_annotation_chart +) -> None: + """A lookup failure must fail closed rather than silently deduping onto a + successfully-derived scope -- and not just for SupersetException: the RLS + lookup is a real DB query and get_extra_cache_keys() renders Jinja for + virtual datasets, so a driver or template error is just as likely as a + SupersetException here, and must not 500 the whole chart-data request.""" + mock_annotation_chart.get_query_context.side_effect = RuntimeError("db boom") + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as security_manager: + security_manager.can_access_datasource.return_value = True + security_manager.get_rls_cache_key.return_value = [] + scope = processor._annotation_source_scope({"value": 1}) + assert scope == {"access": False, "data_key": []} + + +def test_annotation_source_scope_fallback_lookup_also_fails_closed( + processor, mock_annotation_chart +) -> None: + """If the fallback's own RLS lookup fails too (e.g. the same DB outage + that failed the primary derivation), that must not escape either.""" + mock_annotation_chart.get_query_context.side_effect = RuntimeError("db boom") + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as security_manager: + security_manager.can_access_datasource.return_value = True + security_manager.get_rls_cache_key.side_effect = RuntimeError("still down") + scope = processor._annotation_source_scope({"value": 1}) + assert scope == {"access": False, "data_key": None} + + +def test_annotation_source_scope_none_when_chart_missing(processor) -> None: + with patch( + "superset.common.query_context_processor.ChartDAO.find_by_id", + return_value=None, + ): + scope = processor._annotation_source_scope({"value": 999}) + assert scope == {"access": None, "data_key": None} + + +def test_annotation_source_scope_uses_resolved_datasource_for_semantic_views( + processor, mock_annotation_chart +) -> None: + """``Slice.datasource`` 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. + ``resolved_datasource`` must be used instead so these charts still + participate in per-requester access scoping.""" + mock_annotation_chart.datasource = None + mock_annotation_chart.get_query_context.return_value = None + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as security_manager: + security_manager.can_access_datasource.side_effect = [True, False] + security_manager.get_rls_cache_key.return_value = [] + scope_a = processor._annotation_source_scope({"value": 1}) + scope_b = processor._annotation_source_scope({"value": 1}) + assert scope_a == {"access": True, "data_key": []} + assert scope_b == {"access": False, "data_key": []} + security_manager.can_access_datasource.assert_called_with( + mock_annotation_chart.resolved_datasource + ) + + def test_get_data_table_like(processor, mock_query_context): df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]}) coltypes = [GenericDataType.NUMERIC, GenericDataType.STRING] @@ -2226,6 +2461,182 @@ def test_get_df_payload_no_warning_when_not_memory_limited() -> None: assert result["warning"] is None +def test_get_df_payload_result_decouples_annotation_cache_from_dataframe_cache() -> ( + None +): + """ + The dataframe cache entry must stay shareable across viewers, and + annotation-layer data must be resolved through its own (per-user) cache + path -- not stored on the dataframe's cache entry -- so that two viewers + of the same annotated chart share one dataframe cache hit while each + still gets their own annotation-security-scoped payload. + """ + from superset.common.query_object import QueryObject + + mock_query_context = MagicMock() + mock_query_context.force = False + mock_datasource = MagicMock() + mock_datasource.column_names = ["col1"] + + processor = QueryContextProcessor(mock_query_context) + processor._qc_datasource = mock_datasource + + query_obj = QueryObject( + datasource=mock_datasource, + columns=["col1"], + annotation_layers=[ + { + "annotationType": "EVENT", + "sourceType": "NATIVE", + "name": "a", + "value": 1, + } + ], + ) + + class MockCache: + def __init__(self): + self.is_loaded = True + self.applied_filter_columns = ["col1"] + self.df = pd.DataFrame({"col1": [1, 2, 3]}) + self.query = "" + self.status = "success" + self.cache_dttm = "2024-01-01T00:00:00" + self.queried_dttm = "2024-01-01T00:00:00" + self.stacktrace = None + self.error_message = None + self.is_cached = True + self.sql_rowcount = 0 + self.cache_value = None + self.applied_template_filters = [] + self.rejected_filter_columns = [] + self.annotation_data = {"stale": "should not be used"} + self.bq_memory_limited = False + self.bq_memory_limited_row_count = 0 + self.result_persisted = False + self.set_query_result = MagicMock() + + mock_cache = MockCache() + + with ( + patch( + "superset.common.query_context_processor.QueryCacheManager" + ) as mock_cache_manager, + patch.object(query_obj, "validate", return_value=None), + patch.object(processor, "query_cache_key", return_value="df-key"), + patch.object(processor, "annotation_cache_key", return_value="ann-key"), + patch.object( + processor, "_get_annotation_data_cached", return_value={"a": [1, 2]} + ) as mock_get_annotation, + patch.object(processor, "get_cache_timeout", return_value=3600), + ): + mock_cache_manager.get.return_value = mock_cache + result = processor.get_df_payload(query_obj, force_cached=False) + + # The dataframe cache is a hit, so the (expensive) query is never re-run, + # and its cache entry is never rewritten. + mock_cache.set_query_result.assert_not_called() + + # Annotation data is resolved through the separate, per-user cache path, + # keyed by the dedicated annotation cache key -- not the dataframe's key. + mock_get_annotation.assert_called_once() + _, kwargs = mock_get_annotation.call_args + assert kwargs["cache_key"] == "ann-key" + + # The payload serves the freshly-resolved annotation data, not whatever + # (stale) value happened to sit on the dataframe's cache object. + assert result["annotation_data"] == {"a": [1, 2]} + + +def test_get_df_payload_result_annotation_refresh_independent_of_df_marker() -> None: + """ + A GTF forced refresh's idempotency marker is per (nonce, cache_key). Once + the dataframe's marker (keyed on the dataframe's own cache_key) is set, + force_query goes False for every later request carrying that nonce -- + but a *different* annotation access scope (its own annotation_key) has + never actually forced its own entry, and must still force its recompute + rather than silently reading whatever's already sitting under its key. + """ + from superset.common.query_object import QueryObject + + mock_query_context = MagicMock() + mock_query_context.force = True + mock_query_context.force_nonce = "nonce-1" + mock_datasource = MagicMock() + mock_datasource.column_names = ["col1"] + + processor = QueryContextProcessor(mock_query_context) + processor._qc_datasource = mock_datasource + + query_obj = QueryObject( + datasource=mock_datasource, + columns=["col1"], + annotation_layers=[ + { + "annotationType": "EVENT", + "sourceType": "NATIVE", + "name": "a", + "value": 1, + } + ], + ) + query_obj.force_nonce = None # falls back to the context-level nonce above + + class MockCache: + def __init__(self) -> None: + self.is_loaded: bool = True + self.applied_filter_columns: list[str] = ["col1"] + self.df: pd.DataFrame = pd.DataFrame({"col1": [1, 2, 3]}) + self.query: str = "" + self.status: str = "success" + self.cache_dttm: str = "2024-01-01T00:00:00" + self.queried_dttm: str = "2024-01-01T00:00:00" + self.stacktrace: str | None = None + self.error_message: str | None = None + self.is_cached: bool = True + self.sql_rowcount: int = 0 + self.cache_value: Any | None = None + self.applied_template_filters: list[Any] = [] + self.rejected_filter_columns: list[Any] = [] + self.annotation_data: dict[str, Any] = {} + self.bq_memory_limited: bool = False + self.bq_memory_limited_row_count: int = 0 + self.result_persisted: bool = False + self.set_query_result: MagicMock = MagicMock() + + mock_cache = MockCache() + + def marker_lookup(key: str) -> int | None: + """Return the force-nonce marker only for the dataframe's own cache_key.""" + return 1 if key == "gtf-force-nonce:nonce-1:df-key" else None + + with ( + patch( + "superset.common.query_context_processor.QueryCacheManager" + ) as mock_cache_manager, + patch("superset.common.query_context_processor.cache_manager") as cache_manager, + patch.object(query_obj, "validate", return_value=None), + patch.object(processor, "query_cache_key", return_value="df-key"), + patch.object(processor, "annotation_cache_key", return_value="ann-key"), + patch.object( + processor, "_get_annotation_data_cached", return_value={} + ) as mock_get_annotation, + patch.object(processor, "get_cache_timeout", return_value=3600), + ): + cache_manager.data_cache.get.side_effect = marker_lookup + mock_cache_manager.get.return_value = mock_cache + processor.get_df_payload(query_obj, force_cached=False) + + # The dataframe's own marker is set, so its force_query resolves False. + assert mock_cache_manager.get.call_args.kwargs["force_query"] is False + + # The annotation's marker was never set for this nonce, so its + # force_query must resolve True regardless of the dataframe's. + mock_get_annotation.assert_called_once() + _, kwargs = mock_get_annotation.call_args + assert kwargs["force_query"] is True + + def test_raise_for_access_evaluates_access_before_validate(): """ Access must be evaluated before the queries are validated, because query @@ -2618,6 +3029,94 @@ def test_mark_force_executed_noop_without_nonce(processor, mock_query_context): cache_manager.data_cache.set.assert_not_called() +# ============================================================================= +# Annotation-data cache decoupled from the dataframe cache +# ============================================================================= + + +def test_get_annotation_data_cached_reads_from_cache(processor): + """A hit on the annotation-specific key skips recomputation entirely.""" + with patch( + "superset.common.query_context_processor.cache_manager" + ) as cache_manager: + cache_manager.data_cache.get.return_value = {"annotation_data": {"a": 1}} + with patch.object(processor, "get_annotation_data") as mock_get: + result = processor._get_annotation_data_cached( + query_obj=MagicMock(), + cache_key="ak", + force_query=False, + force_cached=False, + timeout=60, + datasource_uid="ds", + ) + assert result == {"a": 1} + mock_get.assert_not_called() + + +def test_get_annotation_data_cached_computes_and_caches_on_miss(processor): + with ( + patch("superset.common.query_context_processor.cache_manager") as cache_manager, + patch("superset.common.query_context_processor.set_and_log_cache") as mock_set, + ): + cache_manager.data_cache.get.return_value = None + with patch.object( + processor, "get_annotation_data", return_value={"a": 1} + ) as mock_get: + result = processor._get_annotation_data_cached( + query_obj=MagicMock(), + cache_key="ak", + force_query=False, + force_cached=False, + timeout=60, + datasource_uid="ds", + ) + assert result == {"a": 1} + mock_get.assert_called_once() + mock_set.assert_called_once_with( + cache_manager.data_cache, "ak", {"annotation_data": {"a": 1}}, 60, "ds" + ) + + +def test_get_annotation_data_cached_force_cached_raises_on_miss(processor): + """``force_cached`` must never fall through to a live compute -- the same + contract ``QueryCacheManager.get`` enforces for the dataframe cache.""" + with patch( + "superset.common.query_context_processor.cache_manager" + ) as cache_manager: + cache_manager.data_cache.get.return_value = None + with pytest.raises(CacheLoadError): + processor._get_annotation_data_cached( + query_obj=MagicMock(), + cache_key="ak", + force_query=False, + force_cached=True, + timeout=60, + datasource_uid="ds", + ) + + +def test_get_annotation_data_cached_force_query_bypasses_cache_read(processor): + """A forced refresh recomputes rather than serving a stale cached value.""" + with ( + patch("superset.common.query_context_processor.cache_manager") as cache_manager, + patch("superset.common.query_context_processor.set_and_log_cache"), + ): + with patch.object( + processor, "get_annotation_data", return_value={"fresh": 1} + ) as mock_get: + result = processor._get_annotation_data_cached( + query_obj=MagicMock(), + cache_key="ak", + force_query=True, + force_cached=False, + timeout=60, + datasource_uid="ds", + ) + assert result == {"fresh": 1} + mock_get.assert_called_once() + cache_manager.data_cache.get.assert_not_called() + + def test_mark_force_executed_swallows_marker_write_error(processor, mock_query_context): """A marker write failure is best-effort — logged and swallowed.""" mock_query_context.force = True