sadpandajoe commented on code in PR #44406:
URL: https://github.com/apache/superset/pull/44406#discussion_r4217596322
##########
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:
`query_cache_key` is mocked to a constant here and the assertions only look
at the final query-object state, so moving `_apply_annotation_overrides` after
key derivation would still pass. That reintroduces the miss where a daily
override pulls in `current_user_id()` Jinja that the saved monthly query lacks,
and the second viewer gets the first viewer's rows. Could `query_cache_key` use
a `side_effect` that asserts the overridden grain, `time_range` and bounds are
already set when it is called?
##########
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:
With `CHART_DATA_INCLUDE_TIMING` on, annotation SQL now runs after
`data_acquisition_ns` is captured and before `payload_assembly_start_ns`
starts, so it lands in none of the reported phases. A cold chart with a 10 ms
dataframe and a 5 s annotation query reports roughly 10 ms of acquisition, and
a dataframe cache hit reports `data_acquisition_ms: null` even though
annotation SQL ran. The configuring-superset docs say `data_acquisition_ms`
"includes database work and annotation dependencies", and version 1 can't
change semantics. Should the annotation fetch be timed into
`data_acquisition_ns` (summed with the dataframe time when both run)?
##########
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:
This only asserts `cache_key` on the `_get_annotation_data_cached` call. The
helper-level `force_cached` test passes the flag directly and the payload test
replaces the helper, so dropping `force_cached=force_cached` at the call site
in `get_df_payload_result` would still pass. Async cache preflight with
`force_cached=True` and a missing annotation entry would then run the
annotation SQL instead of raising `CacheLoadError`. Should this also assert
`kwargs["force_cached"]` (and `force_query`), or add a `force_cached=True` case
with the real helper?
--
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]
