rusackas commented on code in PR #44406:
URL: https://github.com/apache/superset/pull/44406#discussion_r4297984123
##########
superset/common/query_context_processor.py:
##########
@@ -449,34 +481,189 @@ 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]:
+ """
+ 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`.
"""
- source_rls: dict[str, list[str] | None] = {}
+ 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 ANNOTATION_SOURCE_TYPES_WITH_CHART_REFERENCE
):
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)
Review Comment:
That detector is gone now. The key is derived from the override-applied
query objects the live fetch actually runs, through `query_cache_key` /
`get_extra_cache_keys`, which renders the Jinja itself, so spacing like
`current_user_id ()` is handled by the template engine rather than a syntactic
check.
##########
superset/common/query_context_processor.py:
##########
@@ -332,6 +344,30 @@ 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 query_obj.annotation_layers
+ and cache.status != QueryStatus.FAILED
+ ):
+ try:
+ annotation_data = self._get_annotation_data_cached(
Review Comment:
Good catch. The annotation fetch is now timed and summed into
`data_acquisition_ns` (so a dataframe cache hit with annotation SQL no longer
reports null). Added a test for it.
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -117,25 +122,350 @@ 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
+
+
[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_uses_live_fetch_authorization(
+ processor: QueryContextProcessor, mock_annotation_chart: MagicMock
+) -> 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."""
+ 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:
+ # 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 is not None
+ assert scope_b is not None
+ 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: QueryContextProcessor, mock_annotation_chart: MagicMock
+) -> 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)
Review Comment:
Makes sense. `query_cache_key` now has a `side_effect` that asserts the
overridden grain, time range and bounds are already set when it's called, so
reordering would fail the test.
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -2238,6 +2568,252 @@ 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"
Review Comment:
Good point. The decouple test now asserts `force_cached` and `force_query`
are forwarded at the call site, and I added a case that runs the real helper
with `force_cached=True` and no annotation entry, expecting `CacheLoadError`
and no live fetch.
##########
superset/common/query_context_processor.py:
##########
@@ -72,7 +74,6 @@
get_column_name,
get_column_names_from_columns,
get_column_names_from_metrics,
- get_user_id,
is_adhoc_column,
Review Comment:
Those imports are already gone in the current head, and ruff (F401) passes
on the file.
--
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]
