sadpandajoe commented on code in PR #44406:
URL: https://github.com/apache/superset/pull/44406#discussion_r4345616115
##########
tests/unit_tests/tasks/test_query_cancel.py:
##########
@@ -159,6 +159,51 @@ def
test_capture_registers_abort_handler_once_when_id_available() -> None:
assert ordered.index("set_cancellation") < ordered.index("on_abort")
+def test_capture_ignores_cursor_from_a_different_database() -> None:
+ from superset.tasks.async_queries import _capture_query_cancellation
+
+ ctx = MagicMock()
+ qc = _query_context("42")
+ other_database = MagicMock()
+ other_database.id = qc.datasource.database.id + 1
+ with (
+ patch("superset.tasks.async_queries.get_context", return_value=ctx),
+ patch("superset.tasks.async_queries.current_app"),
+ ):
+ with _capture_query_cancellation(qc):
+ sink = _cancel_id_sink.get()
+ assert sink is not None
+ sink(MagicMock(), other_database)
+
+ ctx.on_abort.assert_not_called()
+ ctx.set_cancellation.assert_not_called()
+
+
+def test_notify_cursor_routes_by_database_through_the_real_seam() -> None:
+ """Driven through ``notify_cursor`` (the seam ``get_raw_connection`` calls)
+ rather than the sink directly: the task's own database registers the abort
+ handler exactly once, another database's cursor registers nothing."""
+ from superset.tasks.async_queries import _capture_query_cancellation
+
+ ctx = MagicMock()
+ qc = _query_context("42")
+ own_database = qc.datasource.database
+ other_database = MagicMock()
+ other_database.id = own_database.id + 1
+ with (
+ patch("superset.tasks.async_queries.get_context", return_value=ctx),
+ patch("superset.tasks.async_queries.current_app"),
+ ):
+ with _capture_query_cancellation(qc):
+ notify_cursor(MagicMock(), other_database)
Review Comment:
This still bypasses the production caller: `notify_cursor` is invoked
directly with the database, so reverting
`Database._execute_sql_with_mutation_and_logging` in `models/core.py` back to
`notify_cursor(cursor)` keeps every test green. In that case a warm host
dataframe with a cold annotation layer on another PostgreSQL server records the
annotation cursor against the host database, and abort or timeout can cancel
the wrong server's session. Could a test run the real execute path under
`capture_cancel_id` (or assert the `notify_cursor` call args from it) and check
the host database registers the handle while another database does not?
##########
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(
Review Comment:
Whitespace-separated identity macros still escape this key:
`QueryContext.query_cache_key` only renders the dataset's Jinja when
`SqlaTable.has_extra_cache_key_calls` matches `ExtraCache.regex`, and that
regex requires `current_user_id(` with no space. With an RLS clause like
`owner_id = {{ current_user_id () }}` on the annotation source's dataset, two
users with the same grants derive the same `data_key` (the raw RLS clause is
identical and no user material is added), so the second viewer is served the
first viewer's cached annotation rows. Before this PR the annotation context
bound the requester's user ID, which hid that gap. Should an annotation entry
only be shared when the key demonstrably carries the user-dependent material,
rather than relying on that syntactic gate?
--
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]
