sadpandajoe commented on code in PR #44406:
URL: https://github.com/apache/superset/pull/44406#discussion_r4341302255
##########
superset/common/query_context_processor.py:
##########
@@ -478,73 +518,261 @@ 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 material binding cached annotation data to its security
- context.
-
- 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.
+ Cache key for this query's annotation-layer payload, or ``None`` when
+ the query has no annotation layers or when some layer's access scope
+ could not be resolved (the payload is then never cached, see
+ :meth:`_annotation_source_scope`).
+
+ 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.
+ """
+ if not query_obj or not query_obj.annotation_layers:
+ return None
+ if (annotation_context := self._annotation_cache_context(query_obj))
is None:
+ return None
+ return self.query_cache_key(query_obj,
annotation_context=annotation_context)
+
+ def _annotation_cache_context(
+ self, query_obj: QueryObject
+ ) -> dict[str, Any] | 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`, recorded per *layer* (not per
+ referenced chart) since two layers can reference the same chart
+ with different overrides, and so different scopes.
+
+ Returns ``None`` when any layer's scope cannot be resolved, meaning
+ the payload must not be cached at all.
"""
- source_rls: dict[str, list[str] | None] = {}
- source_versions: dict[str, tuple[str, str]] = {}
- metadata_versions: dict[str, str] = {}
+ 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: list[dict[str, Any]] = []
for layer in query_obj.annotation_layers:
if (
layer.get("sourceType")
not in ANNOTATION_SOURCE_TYPES_WITH_CHART_REFERENCE
):
continue
- layer_value = layer.get("value")
+ if (scope := self._annotation_source_scope(layer)) is None:
+ return None
+ source_scope.append({"value": str(layer.get("value")), **scope})
+ if source_scope:
+ context["source_scope"] = source_scope
+
+ return context
+
+ def _annotation_source_scope(self, layer: dict[str, Any]) -> dict[str,
Any] | None:
+ """
+ Access and data-identity cache-key material for one chart-backed
+ annotation layer, or ``None`` when it cannot be resolved.
+
+ ``None`` means "do not cache": a missing chart, a chart with neither
+ a saved query context nor a resolvable datasource, or any error while
+ deriving the scope would otherwise collapse an authorized and a
+ denied requester onto one shared identity, letting the latter read
+ the former's warmed payload. The caller fetches live instead. A
+ referenced chart whose own caching is disabled
+ (``CACHE_DISABLED_TIMEOUT``) is ``None`` too: its rows are never
+ shared through its own cache, so they must not be shared through the
+ host chart's annotation entry either.
+
+ ``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 and
+ result-type-prepared query objects actually executed (see
+ :meth:`_apply_annotation_overrides` and
+ :func:`~superset.common.query_actions.prepare_query_object`), 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 or
+ that a samples query's cleared metrics only expose at execution.
+ Reusing this logic avoids re-deriving it and automatically inherits
+ any future correctness fixes made there.
+ """
+ layer_value = layer.get("value")
+ try:
chart = (
ChartDAO.find_by_id(layer_value) if layer_value is not None
else None
)
- annotation_datasource: Datasource | None = (
- chart.resolved_datasource if chart else None
- )
- if chart and isinstance(chart.query_context, str):
+ if chart is None:
+ return None
+
+ # The saved query context is what the live fetch executes and
+ # authorizes against, so it is consulted first, independent of
+ # the chart's own datasource pointer: a chart can have its
+ # ``datasource_id`` cleared while its saved query context still
+ # targets a valid dataset, and the live fetch would still work.
+ annotation_query_context = chart.get_query_context()
+ if annotation_query_context is not None:
+ if (
+ annotation_query_context.get_cache_timeout()
+ == CACHE_DISABLED_TIMEOUT
+ ):
+ return None
+ self._apply_annotation_overrides(annotation_query_context,
layer)
try:
- source: dict[str, Any] = json.loads(chart.query_context)[
- "datasource"
+ annotation_query_context.raise_for_access()
+ access: Any = True
+ except SupersetSecurityException:
+ access = False
+ # A denied requester never reads or writes a cached payload,
+ # so skip deriving a data identity (which can render the
+ # dataset's Jinja) for a resource they cannot access.
+ # Keyed on the query object the fetch will execute: a samples
+ # or drill-detail result type rewrites the saved query (e.g.
+ # clears its metrics) before running it, and that can change
+ # the per-user Jinja material the key has to carry.
+ data_key: Any = (
+ [
+ annotation_query_context.query_cache_key(
+ prepare_query_object(
+ annotation_query_context,
+ query_object,
+ query_object.result_type
+ or annotation_query_context.result_type,
+ )
+ )
+ for query_object in annotation_query_context.queries
]
- annotation_datasource = DatasourceDAO.get_datasource(
- source["type"], source["id"]
- )
- except (
- ValueError,
- TypeError,
- KeyError,
- DatasourceNotFound,
- DatasourceTypeNotSupportedError,
- DatasourceValueIsIncorrect,
+ if access
+ else None
+ )
+ annotation_datasource: Any =
annotation_query_context.datasource
+ else:
+ # Fall back to the RLS-clause identity when the chart has no
+ # saved query context to key on. resolved_datasource, not
+ # datasource: the latter is pinned to table-backed
+ # datasources and resolves to None for a semantic-view-backed
+ # chart.
+ datasource = chart.resolved_datasource
+ if datasource is None:
+ return None
+ if CACHE_DISABLED_TIMEOUT in (
+ chart.cache_timeout,
+ datasource.cache_timeout,
):
- annotation_datasource = None
+ return None
+ access = security_manager.can_access_datasource(datasource)
+ data_key = security_manager.get_rls_cache_key(datasource)
+ annotation_datasource = datasource
+ scope: dict[str, Any] = {"access": access, "data_key": data_key}
+ # Semantic-view results are versioned by the layer's metadata
+ # generation and result discriminator, which the datasource-level
+ # key material above does not carry.
if isinstance(annotation_datasource, SemanticView):
- metadata_versions[str(layer_value)] = (
- annotation_datasource.metadata_generation
- )
- source_rls[str(layer.get("value"))] = (
- security_manager.get_rls_cache_key(annotation_datasource)
- if annotation_datasource
- else None
+ scope["metadata_version"] =
annotation_datasource.metadata_generation
+ # Separates entries written before NULL filter selections were
+ # preserved from those written after.
+ scope["semantic_filter_protocol"] = "semantic-null-filters-v1"
+ if (
+ discriminator :=
annotation_datasource.result_cache_discriminator
+ ) is not None:
+ scope["result_version"] = discriminator
+ 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. A fallback scope
+ # shared by every requester whose derivation failed would still
+ # be a shared identity, so the answer is "no caching", not a
+ # substitute key.
+ logger.warning(
+ "Could not derive annotation cache key for chart %s; "
+ "annotation data for this request will not be cached",
+ layer_value,
+ exc_info=True,
)
- if isinstance(annotation_datasource, SemanticView):
- discriminator: tuple[str, str] | None = (
- annotation_datasource.result_cache_discriminator
+ return None
+ return scope
+
+ def _get_annotation_data_cached(
+ self,
+ query_obj: QueryObject,
+ cache_key: str | None,
+ 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.
+
+ A ``None`` ``cache_key`` means some layer's access scope could not be
+ resolved (see :meth:`_annotation_source_scope`), so there is no key
+ under which a cached payload could be safely shared: an unresolved
+ scope looks the same for an authorized and a denied requester. The
+ payload is then fetched live and never cached. That holds under
+ ``force_cached`` too: computing here beats failing the read-back of
+ an otherwise successful async task.
+ """
+ if cache_key is None:
Review Comment:
When an annotation layer's scope is `None` (including the new cache-disabled
case where the referenced chart has `cache_timeout=-1`), this fetches the
annotation live even under `force_cached`, and never caches it. In async mode
that means a warm-dataframe load runs the source chart's query inside the HTTP
request in `_run_async`'s probe, and on a cold load the worker's annotation run
is discarded and the `enableAsyncMode: false` read-back runs it again in the
request, with no task timeout or cancellation. A slow source chart with caching
disabled can therefore hit the web request timeout even though the task
succeeded. Would it make sense to return a cache miss from the async probe for
uncacheable layers, or otherwise skip the in-request fetch when running async?
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -126,23 +135,471 @@ 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}]
+ # ``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])
+
+ # The requesting user differs between the two calls. Binding user identity
+ # into the context (by name, from either module) would make the contexts
+ # diverge, so equality below proves the key is scope-based.
with (
+ patch(
+ "superset.common.query_context_processor.security_manager",
+ new_callable=MagicMock,
+ ) as security_manager,
patch(
"superset.common.query_context_processor.get_user_id",
+ create=True,
side_effect=[1, 2],
),
- patch("superset.common.query_context_processor.security_manager"),
+ patch("superset.utils.core.get_user_id", side_effect=[1, 2]),
+ ):
+ 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)
- contexts = [
- call.kwargs["annotation_context"] for call in
query_obj.cache_key.call_args_list
+ for call in query_obj.cache_key.call_args_list:
+ assert "annotation_context" not in call.kwargs
+
+
[email protected]
+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: QueryContextProcessor, mock_annotation_chart: MagicMock
+) -> 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 is not None
+ assert scope_b is not None
+ 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: QueryContextProcessor, mock_annotation_chart: MagicMock
+) -> 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 = MagicMock()
+ mock_query_context: MagicMock = 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_keys_on_result_type_prepared_query(
+ processor: QueryContextProcessor, mock_annotation_chart: MagicMock
+) -> None:
+ """A referenced chart saved with a samples result type executes a
+ rewritten query (metrics cleared, every column selected), so its key
+ must be derived from that prepared query, not the saved one -- a Jinja
+ clause that only renders user-specific material without metrics would
+ otherwise give two viewers one key over different rows."""
+ from superset.common.query_object import QueryObject
+
+ mock_datasource = MagicMock()
+ mock_datasource.columns = []
+ saved_query = QueryObject(
+ datasource=mock_datasource, metrics=["count"], columns=["col1"]
+ )
+ mock_query_context = MagicMock()
+ mock_query_context.result_type = ChartDataResultType.SAMPLES
Review Comment:
This test sets `SAMPLES` on the context but leaves `saved_query.result_type`
as `None`, so it passes even if the `query_object.result_type or
annotation_query_context.result_type` argument to `prepare_query_object` is
reverted to the context-only value. That override is the latest fix, and a
saved `FULL` context with an individual `SAMPLES` query (metrics `["count"]`)
is not exercised. Could you add a case with a `FULL` context and a `SAMPLES`
query-level result type asserting `data_key` is derived from the
cleared-metrics query?
--
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]
