You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ru...@apache.org on 2026/09/17 22:46:23 UTC
(superset) branch fix/annotation-cache-scope created (now 19954fdb004)
This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a change to branch fix/annotation-cache-scope
in repository https://gitbox.apache.org/repos/asf/superset.git
at 19954fdb004 fix(charts): decouple annotation cache from the shared dataframe cache
This branch includes the following new commits:
new 19954fdb004 fix(charts): decouple annotation cache from the shared dataframe cache
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 01/01: fix(charts): decouple annotation cache from the shared dataframe cache
Posted by ru...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch fix/annotation-cache-scope
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 19954fdb004845b9cf25f7ffe0ccb97b8f7211df
Author: Superset Dev <de...@superset.apache.org>
AuthorDate: Thu Sep 17 15:45:53 2026 -0700
fix(charts): decouple annotation cache from the shared dataframe cache
Annotation-layer data is fetched per requesting user and, for chart-backed
layers, scoped by the referenced chart datasource's RLS -- a stricter
requirement than the dataframe itself has. Binding that context into the
dataframe's own cache key meant every distinct viewer of an annotated chart
got their own full copy of the (potentially much larger) dataframe, instead
of sharing one cache entry.
Split the two: the dataframe cache key goes back to depending only on its
own datasource/RLS/extra_cache_keys, while annotation data is resolved and
cached under a separate, still user/RLS-scoped key. This also restores
cross-user task dedup for annotated charts in the async (GTF) flow, which
had the same coupling problem.
Co-Authored-By: Evan Rusackas <ev...@preset.io>
Co-Authored-By: Claude Sonnet 5 <no...@anthropic.com>
---
superset/common/query_context_processor.py | 106 +++++++++--
.../common/test_query_context_processor.py | 202 ++++++++++++++++++++-
2 files changed, 289 insertions(+), 19 deletions(-)
diff --git a/superset/common/query_context_processor.py b/superset/common/query_context_processor.py
index 2c27fb9fb06..f296f8317e9 100644
--- a/superset/common/query_context_processor.py
+++ b/superset/common/query_context_processor.py
@@ -43,6 +43,7 @@ from superset.constants import CACHE_DISABLED_TIMEOUT, CacheRegion
from superset.daos.annotation_layer import AnnotationLayerDAO
from superset.daos.chart import ChartDAO
from superset.exceptions import (
+ CacheLoadError,
QueryObjectValidationError,
SupersetException,
)
@@ -254,6 +255,7 @@ class QueryContextProcessor:
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)
@@ -304,7 +306,6 @@ class QueryContextProcessor:
)
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
@@ -317,7 +318,6 @@ class QueryContextProcessor:
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,
@@ -328,6 +328,26 @@ class QueryContextProcessor:
# 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=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`
@@ -397,7 +417,7 @@ class QueryContextProcessor:
"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,
@@ -423,15 +443,15 @@ class QueryContextProcessor:
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,
@@ -445,15 +465,30 @@ class QueryContextProcessor:
)
return cache_key
+ 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 per requesting user and, for
+ chart-backed layers, scoped by the RLS clauses of the referenced
+ chart's datasource — a stricter security requirement than the
+ dataframe itself has. Keying them separately from
+ :meth:`query_cache_key` keeps that per-user scoping from forcing every
+ distinct viewer of an annotated chart onto their own full copy of the
+ (potentially much larger) shared dataframe.
+ """
+ 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)
+ )
+
def _annotation_cache_context(self, query_obj: QueryObject) -> dict[str, Any]:
"""
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.
+ context: the requesting user and, for chart-backed layers, the RLS
+ clauses of the referenced chart's datasource.
"""
source_rls: dict[str, list[str] | None] = {}
for layer in query_obj.annotation_layers:
@@ -471,6 +506,49 @@ class QueryContextProcessor:
)
return {"user_id": get_user_id(), "source_rls": source_rls}
+ 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,
+ )
+ raise CacheLoadError("Error loading annotation data from cache")
+
+ annotation_data = self.get_annotation_data(query_obj)
+ set_and_log_cache(
+ cache_manager.data_cache,
+ cache_key,
+ {"annotation_data": annotation_data},
+ timeout,
+ datasource_uid,
+ )
+ return annotation_data
+
def get_query_result(self, query_object: QueryObject) -> QueryResult:
"""
Returns a pandas dataframe based on the query object.
diff --git a/tests/unit_tests/common/test_query_context_processor.py b/tests/unit_tests/common/test_query_context_processor.py
index 5dc21d12b85..001e2021e7f 100644
--- a/tests/unit_tests/common/test_query_context_processor.py
+++ b/tests/unit_tests/common/test_query_context_processor.py
@@ -30,7 +30,7 @@ from superset.common.query_context_processor import (
normalize_contribution_totals,
QueryContextProcessor,
)
-from superset.exceptions import QueryObjectValidationError
+from superset.exceptions import CacheLoadError, QueryObjectValidationError
from superset.utils.core import GenericDataType
from superset.utils.date_parser import get_past_or_future
@@ -115,8 +115,8 @@ 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_annotation_data_to_requesting_user(processor):
+ """The annotation cache key must differ per requesting user."""
query_obj = MagicMock()
query_obj.annotation_layers = [{"sourceType": "NATIVE", "name": "a", "value": 1}]
with (
@@ -126,14 +126,33 @@ def test_query_cache_key_binds_annotation_data_to_requesting_user(processor):
),
patch("superset.common.query_context_processor.security_manager"),
):
- processor.query_cache_key(query_obj)
- processor.query_cache_key(query_obj)
+ 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_query_cache_key_does_not_bind_requesting_user(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) is user-scoped."""
+ 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)
+ for call in query_obj.cache_key.call_args_list:
+ assert "annotation_context" not in call.kwargs
+
+
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]
@@ -2186,6 +2205,91 @@ 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():
+ """
+ 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_raise_for_access_evaluates_access_before_validate():
"""
Access must be evaluated before the queries are validated, because query
@@ -2550,6 +2654,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