rusackas commented on code in PR #44546:
URL: https://github.com/apache/superset/pull/44546#discussion_r4080866524
##########
superset/common/query_context_processor.py:
##########
@@ -445,31 +464,150 @@ 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 ("line", "table"):
continue
layer_value = layer.get("value")
- chart = (
- ChartDAO.find_by_id(layer_value) if layer_value is not None
else None
+ source_scope[str(layer_value)] =
self._annotation_source_scope(layer_value)
+ if source_scope:
+ context["source_scope"] = source_scope
+
+ return context
+
+ def _annotation_source_scope(self, layer_value: 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. ``data_key`` is the
+ annotation chart's own query cache key(s), which already capture the
+ datasource version, RLS clauses, and any per-user Jinja/virtual-dataset
+ RLS material — reusing it here avoids re-deriving that logic and
+ automatically inherits any future correctness fixes made there.
+ """
+ chart = ChartDAO.find_by_id(layer_value) if layer_value is not None
else None
+ datasource = chart.datasource if chart else None
Review Comment:
Good catch, moved the DAO lookup and datasource load inside the try so a
failure there fails closed instead of aborting the whole request.
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -115,25 +115,168 @@ 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):
Review Comment:
Added return types to those.
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -115,25 +115,168 @@ 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):
+ """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():
+ """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):
+ """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():
+ """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
Review Comment:
Added, thanks.
##########
tests/unit_tests/common/test_query_context_processor.py:
##########
@@ -115,25 +115,168 @@ 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):
+ """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():
+ """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):
+ """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():
+ """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
+):
+ """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(1)
+ scope_b = processor._annotation_source_scope(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
+):
+ """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(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_fails_closed_on_any_derivation_error(
+ processor, mock_annotation_chart
+):
+ """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(1)
+ assert scope == {"access": False, "data_key": []}
Review Comment:
Added an assertion on the warning via caplog.
--
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]
