diff --git a/docs/admin_docs/configuration/cache.mdx b/docs/admin_docs/configuration/cache.mdx index 397ca4a408d6..318f7634989f 100644 --- a/docs/admin_docs/configuration/cache.mdx +++ b/docs/admin_docs/configuration/cache.mdx @@ -282,6 +282,48 @@ Requirements and limitations: - Cascading/dependent native filters and search-term variants of filter option queries are not warmed by this strategy. +### Warming Up Favorited Dashboards + +The `favorites` strategy warms the cache for each user's favorited dashboards. It is off by +default; enable it and schedule it with Celery beat: + +```python +CACHE_WARMUP_FAVORITES = True +CACHE_WARMUP_FAVORITES_MAX_DASHBOARDS = 5 # most recently favorited dashboards per user +CACHE_WARMUP_FAVORITES_QUEUE = "cache_warmup" # optional dedicated Celery queue + +class CeleryConfig(CeleryConfig): + beat_schedule = { + "cache-warmup-favorites": { + "task": "cache-warmup", + "schedule": crontab(minute=0, hour="*/6"), + "kwargs": { + "strategy_name": "favorites", + "since": "7 days ago", # only users who logged in since then + "max_users": 100, # users handled per run, most recent login first + }, + }, + } +``` + +Each run enqueues one task per selected user. The task runs as that user, so the cache entries +match what that user's dashboards request, and only charts whose query results are missing or +expired are computed. It does not use `SUPERSET_CACHE_WARMUP_USER`. + +Notes: + +- A Celery worker and a real `DATA_CACHE_CONFIG` backend are required (with `NullCache` every + run would recompute). +- A per-result lock, kept in `CACHE_CONFIG`, stops concurrent runs from computing the same + result twice; use a shared backend such as Redis for it to work across processes. +- Set `CACHE_WARMUP_FAVORITES_QUEUE` and run a worker for that queue to keep warm-up from + competing with reports, alerts and thumbnails. Each per-user task has a 15 minute time limit. +- Dashboard default column filters (native filter defaults and legacy defaults) are applied. + Other dashboard filter settings, such as time range, time column and grain, extras and ad hoc + filters, are not, so charts whose dashboard load uses them may still miss the cache. +- `CACHE_WARMUP_FAVORITES` is an ordinary configuration setting, not a feature flag. +- Queued per-user tasks expire after two hours if no worker picks them up. + ## Caching Thumbnails This is an optional feature that can be turned on by activating its [feature flag](/admin-docs/configuration/configuring-superset#feature-flags) on config: diff --git a/superset/commands/chart/warm_up_cache.py b/superset/commands/chart/warm_up_cache.py index 623869871551..7dde90fe2989 100644 --- a/superset/commands/chart/warm_up_cache.py +++ b/superset/commands/chart/warm_up_cache.py @@ -18,6 +18,8 @@ from typing import Any, cast, Optional, Union +from celery.exceptions import SoftTimeLimitExceeded + from superset.commands.base import BaseCommand from superset.commands.chart.data.get_data_command import ChartDataCommand from superset.commands.chart.exceptions import ( @@ -26,8 +28,9 @@ WarmUpCacheChartNotFoundError, ) from superset.common.db_query_status import QueryStatus +from superset.constants import CACHE_DISABLED_TIMEOUT from superset.exceptions import SupersetSecurityException -from superset.extensions import db, security_manager +from superset.extensions import cache_manager, db, security_manager from superset.models.slice import Slice from superset.semantic_layers.access import ( is_semantic_layers_enabled, @@ -37,6 +40,9 @@ from superset.utils.core import error_msg_from_exception, QueryObjectFilterClause from superset.views.utils import get_dashboard_extra_filters +# Seconds a per-result warm-up lock is held if the worker dies before releasing it. +WARMUP_LOCK_TTL = 900 + class ChartWarmUpCacheCommand(BaseCommand): def __init__( @@ -44,10 +50,14 @@ def __init__( chart_or_id: Union[int, Slice], dashboard_id: Optional[int], extra_filters: Optional[str], + only_if_stale: bool = False, ): self._chart_or_id = chart_or_id self._dashboard_id = dashboard_id self._extra_filters = extra_filters + # When set, charts whose query results are all cached are left untouched + # and missing results are computed without bypassing the cache. + self._only_if_stale = only_if_stale def _get_dashboard_filters(self, chart_id: int) -> list[dict[str, Any]]: """Retrieve dashboard filters from extra_filters or dashboard metadata.""" @@ -59,6 +69,42 @@ def _get_dashboard_filters(self, chart_id: int) -> list[dict[str, Any]]: return get_dashboard_extra_filters(chart_id, self._dashboard_id) + @staticmethod + def _missing_cache_keys(query_context: Any) -> list[str]: + """Return cache keys of the context's queries lacking an unexpired entry.""" + if query_context.get_cache_timeout() == CACHE_DISABLED_TIMEOUT: + # Results are never stored, so there is nothing to warm. + return [] + missing: list[str] = [] + for query in query_context.queries: + # Validation normalizes filters, matching the key used when querying. + query.validate() + cache_key = query_context.query_cache_key(query) + if not cache_key or cache_manager.data_cache.get(cache_key) is None: + missing.append(cache_key or "") + return missing + + def _acquire_warmup_lock(self, query_context: Any) -> tuple[bool, str | None]: + """ + Decide whether this worker should compute the chart's missing results. + + Returns ``(should_run, lock_key)``. ``should_run`` is False when everything is + cached or another worker holds the lock; otherwise the returned lock key must + be released after the computation. + """ + missing = self._missing_cache_keys(query_context) + if not missing: + return False, None + # Only one worker warms a given result at a time. + lock_key = f"cache-warmup-lock:{missing[0]}" + if not cache_manager.cache.add(lock_key, 1, timeout=WARMUP_LOCK_TTL): + return False, None + if not self._missing_cache_keys(query_context): + # Another worker stored the result while this one waited. + cache_manager.cache.delete(lock_key) + return False, None + return True, lock_key + def _warm_up_non_legacy_cache(self, chart: Slice) -> tuple[Any, Any]: """Warm up cache for non-legacy visualizations.""" query_context = chart.get_query_context() @@ -77,10 +123,20 @@ def _warm_up_non_legacy_cache(self, chart: Slice) -> tuple[Any, Any]: + query.filter ) - query_context.force = True - command = ChartDataCommand(query_context) - command.validate() - payload = command.run() + lock_key: str | None = None + if self._only_if_stale: + should_run, lock_key = self._acquire_warmup_lock(query_context) + if not should_run: + return None, QueryStatus.SUCCESS + else: + query_context.force = True + try: + command = ChartDataCommand(query_context) + command.validate() + payload = command.run() + finally: + if lock_key: + cache_manager.cache.delete(lock_key) # Report the first error. for query_result in cast(list[dict[str, Any]], payload["queries"]): @@ -108,6 +164,8 @@ def run(self, *, skip_disabled: bool = False) -> dict[str, Any]: try: error, status = self._warm_up_non_legacy_cache(chart) + except SoftTimeLimitExceeded: + raise except Exception as ex: # pylint: disable=broad-except error = error_msg_from_exception(ex) status = None diff --git a/superset/config.py b/superset/config.py index efba67f3f816..eff0ca576dc0 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1476,6 +1476,21 @@ def sync_theme_logo_href( # - Cascade/dependent filters and search-term variants are not warmed in this # version; only base option queries are supported. +# Gate for the `favorites` cache warm-up strategy, which warms charts on each +# user's favorited dashboards. Schedule it through the `cache-warmup` Celery beat +# task (see the `FavoritesStrategy` docstring in superset/tasks/cache.py). A +# separate task per user runs as that user and computes only query results that +# are missing or expired in DATA_CACHE_CONFIG. Requires a Celery worker and a real +# DATA_CACHE_CONFIG backend (with NullCache every run would recompute). +CACHE_WARMUP_FAVORITES = False +# Maximum number of favorited dashboards (most recently favorited first) warmed +# per user per run. +CACHE_WARMUP_FAVORITES_MAX_DASHBOARDS = 5 +# Celery queue for per-user favorites warm-up tasks. Point it at a dedicated, +# low-priority queue (served by its own worker) so a warm-up burst cannot delay +# other tasks. None uses the default queue. +CACHE_WARMUP_FAVORITES_QUEUE: str | None = None + # Time before selenium times out after trying to locate an element on the page and wait # for that element to load for a screenshot. SCREENSHOT_LOCATE_WAIT = int(timedelta(seconds=10).total_seconds()) diff --git a/superset/tasks/cache.py b/superset/tasks/cache.py index c2a6d4e6d3b9..e72b5c037d66 100644 --- a/superset/tasks/cache.py +++ b/superset/tasks/cache.py @@ -20,6 +20,7 @@ from dataclasses import dataclass from typing import Any, Optional, Union +from celery.exceptions import SoftTimeLimitExceeded from celery.utils.log import get_task_logger from flask import current_app from sqlalchemy import and_, func @@ -28,8 +29,9 @@ from superset import db, security_manager from superset.common.query_context import QueryContext from superset.daos.dashboard import DashboardDAO +from superset.exceptions import SupersetSecurityException from superset.extensions import celery_app -from superset.models.core import Log +from superset.models.core import FavStar, FavStarClassName, Log from superset.models.dashboard import Dashboard from superset.tags.models import Tag, TaggedObject from superset.tasks.native_filter_cache import ( @@ -92,6 +94,8 @@ class Strategy: # pylint: disable=too-few-public-methods name: str = "" uses_webdriver: bool = True + # Per-user strategies fan out one task per user from ``get_user_ids``. + per_user: bool = False def __init__(self) -> None: pass @@ -102,6 +106,9 @@ def get_urls(self) -> list[str]: def get_tasks(self) -> list[CacheWarmupTask]: return [] + def get_user_ids(self) -> list[int]: + return [] + class DummyStrategy(Strategy): # pylint: disable=too-few-public-methods """ @@ -316,11 +323,73 @@ def get_tasks(self) -> list[CacheWarmupTask]: return tasks +class FavoritesStrategy(Strategy): # pylint: disable=too-few-public-methods + """ + Warm up charts on each user's favorited dashboards. + + Cache keys depend on the requesting user (row level security, impersonation), + so each user's dashboards are warmed by a separate task that runs as that user. + Only results that are missing or expired are computed. Users are selected by + recent login, most recent first, up to ``max_users`` per run. Requires + ``CACHE_WARMUP_FAVORITES = True``. + + beat_schedule = { + 'cache-warmup-favorites': { + 'task': 'cache-warmup', + 'schedule': crontab(minute=0, hour='*/6'), + 'kwargs': { + 'strategy_name': 'favorites', + 'since': '7 days ago', + 'max_users': 100, + }, + }, + } + + """ + + name: str = "favorites" + uses_webdriver: bool = False + per_user: bool = True + + def __init__(self, since: str = "7 days ago", max_users: int = 100) -> None: + super().__init__() + self.since: Any = parse_human_datetime(since) if since else None + self.max_users: int = max(1, int(max_users)) + + def get_urls(self) -> list[str]: + return [] + + def get_user_ids(self) -> list[int]: + """Return ids of active users with favorited dashboards, newest login first.""" + query = ( + db.session.query(FavStar.user_id) + .join( + security_manager.user_model, + security_manager.user_model.id == FavStar.user_id, + ) + .filter( + FavStar.class_name == FavStarClassName.DASHBOARD, + security_manager.user_model.active.is_(True), + security_manager.user_model.last_login.isnot(None), + ) + ) + if self.since: + query = query.filter(security_manager.user_model.last_login >= self.since) + rows = ( + query.group_by(FavStar.user_id, security_manager.user_model.last_login) + .order_by(security_manager.user_model.last_login.desc()) + .limit(self.max_users) + .all() + ) + return [row.user_id for row in rows] + + strategies: list[type[Strategy]] = [ DummyStrategy, TopNDashboardsStrategy, DashboardTagsStrategy, NativeFilterOptionsStrategy, + FavoritesStrategy, ] strategy_registry: dict[str, type[Strategy]] = { strategy.name: strategy for strategy in strategies @@ -349,6 +418,27 @@ def _warmup_urls( _browser_manager._cleanup() # noqa: SLF001 +def _dispatch_per_user( + strategy: Strategy, results: dict[str, list[str]] +) -> Union[dict[str, list[str]], str]: + """Enqueue one warm-up task per user selected by a per-user strategy.""" + if not current_app.config.get("CACHE_WARMUP_FAVORITES"): + message = "CACHE_WARMUP_FAVORITES is disabled; nothing was warmed up." + logger.warning(message) + return message + queue: str | None = current_app.config.get("CACHE_WARMUP_FAVORITES_QUEUE") + for user_id in strategy.get_user_ids(): + try: + cache_warmup_user_favorites.apply_async( + args=(user_id,), queue=queue, expires=FAVORITES_TASK_EXPIRES + ) + results["success"].append(f"user:{user_id}") + except Exception: # noqa: BLE001 + logger.exception("Error enqueueing cache warm-up for user %s", user_id) + results["errors"].append(f"user:{user_id}") + return results + + @celery_app.task(name="cache-warmup") def cache_warmup( strategy_name: str, *args: Any, **kwargs: Any @@ -377,6 +467,9 @@ def cache_warmup( results: dict[str, list[str]] = {"success": [], "errors": []} + if strategy.per_user: + return _dispatch_per_user(strategy, results) + warmup_username: str | None = current_app.config.get("SUPERSET_CACHE_WARMUP_USER") if not warmup_username: message = ( @@ -426,3 +519,104 @@ def cache_warmup( results, ) return results + + +# Seconds after which an unstarted per-user task is discarded. +FAVORITES_TASK_EXPIRES = 7200 +# Favorites read per wanted dashboard, leaving room for inaccessible ones. +FAVORITES_FETCH_FACTOR = 4 + + +def _get_favorite_dashboards(user_id: int, limit: int) -> list[Dashboard]: + """Return the user's most recently favorited dashboards, newest first.""" + dashboard_ids: list[int] = [ + row.obj_id + for row in db.session.query(FavStar.obj_id) + .filter( + FavStar.user_id == user_id, + FavStar.class_name == FavStarClassName.DASHBOARD, + ) + .order_by(FavStar.dttm.desc()) + # Some favorites may be inaccessible or deleted, so read a bounded surplus. + .limit(limit * FAVORITES_FETCH_FACTOR) + .all() + ] + dashboards: list[Dashboard] = [] + if not dashboard_ids: + return dashboards + by_id: dict[int, Dashboard] = { + dashboard.id: dashboard + for dashboard in db.session.query(Dashboard) + .options(selectinload(Dashboard.slices)) + .filter(Dashboard.id.in_(dashboard_ids)) + .all() + } + for dashboard_id in dashboard_ids: + if (dashboard := by_id.get(dashboard_id)) is None: + continue + try: + security_manager.raise_for_access(dashboard=dashboard) + except SupersetSecurityException: + continue + dashboards.append(dashboard) + if len(dashboards) >= limit: + break + return dashboards + + +@celery_app.task( + name="cache-warmup-user-favorites", soft_time_limit=840, time_limit=900 +) +def cache_warmup_user_favorites(user_id: int) -> dict[str, list[int]]: + """ + Warm the cache for charts on a user's favorited dashboards. + + Runs as the given user, since cache keys depend on the requesting user (row + level security, impersonation). Charts whose cached results are all present + are skipped; only missing or expired results are computed. + """ + # pylint: disable=import-outside-toplevel + from superset.commands.chart.exceptions import ChartAccessDeniedError + from superset.commands.chart.warm_up_cache import ChartWarmUpCacheCommand + from superset.utils.core import override_user + + results: dict[str, list[int]] = {"success": [], "errors": []} + if not current_app.config.get("CACHE_WARMUP_FAVORITES"): + logger.warning("CACHE_WARMUP_FAVORITES is disabled; skipping user %s", user_id) + return results + user: Any = security_manager.get_user_by_id(user_id) + if not user: + logger.warning("Favorites cache warm-up: user %s not found", user_id) + return results + + limit: int = current_app.config["CACHE_WARMUP_FAVORITES_MAX_DASHBOARDS"] + with override_user(user, force=True): + for dashboard in _get_favorite_dashboards(user_id, limit): + for chart in dashboard.slices: + try: + outcome: dict[str, Any] = ChartWarmUpCacheCommand( + chart, dashboard.id, None, only_if_stale=True + ).run(skip_disabled=True) + except ChartAccessDeniedError: + logger.info( + "Skipping chart %s on dashboard %s: no access", + chart.id, + dashboard.id, + ) + continue + except SoftTimeLimitExceeded: + logger.warning( + "Favorites cache warm-up for user %s timed out", user_id + ) + return results + except Exception: # noqa: BLE001 + logger.exception( + "Error warming up chart %s on dashboard %s", + chart.id, + dashboard.id, + ) + results["errors"].append(chart.id) + continue + key: str = "errors" if outcome.get("viz_error") else "success" + results[key].append(chart.id) + return results diff --git a/tests/unit_tests/commands/chart/warm_up_cache_key_test.py b/tests/unit_tests/commands/chart/warm_up_cache_key_test.py new file mode 100644 index 000000000000..6d66ae25ead7 --- /dev/null +++ b/tests/unit_tests/commands/chart/warm_up_cache_key_test.py @@ -0,0 +1,118 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""The warm-up cache lookup must use the same key as request-time execution.""" + +from typing import Any +from unittest.mock import MagicMock, patch + +from superset.commands.chart.warm_up_cache import ChartWarmUpCacheCommand +from superset.common.query_context_processor import QueryContextProcessor +from superset.common.query_object import QueryObject + + +def _query_context( + time_range: str = "Last week", database_extra: str = "{}" +) -> MagicMock: + datasource = MagicMock() + datasource.database.extra = database_extra + datasource.database.impersonate_user = False + datasource.uid = "1__table" + datasource.changed_on = None + datasource.get_extra_cache_keys.return_value = [] + datasource.column_names = ["a"] + query = QueryObject( + datasource=datasource, + columns=["a"], + metrics=["count"], + time_range=time_range, + ) + query_context = MagicMock() + query_context.datasource = datasource + query_context.queries = [query] + query_context.get_cache_timeout.return_value = 300 + processor = QueryContextProcessor(query_context) + query_context.query_cache_key.side_effect = processor.query_cache_key + return query_context + + +def _lookup_keys(query_context: MagicMock, rls: list[str]) -> list[str]: + """Keys the warm-up path looks up, recorded from the cache read.""" + looked_up: list[str] = [] + with ( + patch("superset.commands.chart.warm_up_cache.cache_manager") as cm, + patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as sm, + ): + sm.get_rls_cache_key.return_value = rls + cm.data_cache.get.side_effect = lambda key: looked_up.append(key) + ChartWarmUpCacheCommand._missing_cache_keys(query_context) + return looked_up + + +def _request_key(query_context: MagicMock, rls: list[str]) -> Any: + """Key the request path derives: validate the query, then build the key.""" + query = query_context.queries[0] + query.validate() + with patch( + "superset.common.query_context_processor.security_manager", + new_callable=MagicMock, + ) as sm: + sm.get_rls_cache_key.return_value = rls + return query_context.query_cache_key(query) + + +def test_warmup_key_matches_request_key(app_context: None) -> None: + query_context = _query_context() + assert _lookup_keys(query_context, ["rls"]) == [ + _request_key(query_context, ["rls"]) + ] + + +def test_row_level_security_changes_the_key(app_context: None) -> None: + query_context = _query_context() + assert _lookup_keys(query_context, ["region = 'EU'"]) != _lookup_keys( + query_context, ["region = 'US'"] + ) + + +def test_relative_time_range_key_is_stable(app_context: None) -> None: + """Keys use the user-supplied relative range, not the resolved timestamps.""" + first = _lookup_keys(_query_context("Last week"), []) + second = _lookup_keys(_query_context("Last week"), []) + assert first == second + assert first != _lookup_keys(_query_context("Last month"), []) + + +def test_per_user_caching_key_follows_the_running_user(app_context: None) -> None: + """With per-user caching, the key depends on the user the task runs as.""" + from superset.utils.core import override_user + + query_context = _query_context(database_extra='{"per_user_caching": true}') + engine_spec = query_context.datasource.database.db_engine_spec + engine_spec.get_impersonation_key.side_effect = lambda user: f"user-{user.id}" + + with override_user(MagicMock(id=1), force=True): + first = _lookup_keys(query_context, []) + with override_user(MagicMock(id=2), force=True): + second = _lookup_keys(query_context, []) + with override_user(MagicMock(id=1), force=True): + again = _lookup_keys(query_context, []) + + assert first != second + assert first == again diff --git a/tests/unit_tests/commands/chart/warm_up_cache_test.py b/tests/unit_tests/commands/chart/warm_up_cache_test.py index 64cad9de7b2d..8eb8fe8a4c87 100644 --- a/tests/unit_tests/commands/chart/warm_up_cache_test.py +++ b/tests/unit_tests/commands/chart/warm_up_cache_test.py @@ -14,7 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from unittest.mock import Mock, patch +from typing import Any +from unittest.mock import MagicMock, Mock, patch import pytest @@ -344,3 +345,64 @@ def test_validate_with_nonexistent_chart_id(mock_db): with pytest.raises(WarmUpCacheChartNotFoundError): command.validate() + + +def _ctx(keys: list[str]) -> MagicMock: + query_context = MagicMock(queries=[MagicMock() for _ in keys]) + query_context.get_cache_timeout.return_value = 300 + query_context.query_cache_key.side_effect = keys + return query_context + + +def test_missing_cache_keys_none_when_cached() -> None: + with patch("superset.commands.chart.warm_up_cache.cache_manager") as cm: + cm.data_cache.get.return_value = {"df": 1} + assert ChartWarmUpCacheCommand._missing_cache_keys(_ctx(["k1", "k2"])) == [] + + +def test_missing_cache_keys_reports_expired_entry() -> None: + with patch("superset.commands.chart.warm_up_cache.cache_manager") as cm: + cm.data_cache.get.side_effect = [{"df": 1}, None] + assert ChartWarmUpCacheCommand._missing_cache_keys(_ctx(["k1", "k2"])) == ["k2"] + + +def test_missing_cache_keys_validates_before_keying() -> None: + """Keys are derived after query validation, as on the request path.""" + query_context = _ctx(["k1"]) + order: list[str] = [] + + def record_validate() -> None: + order.append("validate") + + def record_key(_query: Any) -> str: + order.append("key") + return "k1" + + query_context.queries[0].validate.side_effect = record_validate + query_context.query_cache_key.side_effect = record_key + with patch("superset.commands.chart.warm_up_cache.cache_manager") as cm: + cm.data_cache.get.return_value = {"df": 1} + ChartWarmUpCacheCommand._missing_cache_keys(query_context) + assert order == ["validate", "key"] + + +def test_missing_cache_keys_empty_when_caching_disabled() -> None: + query_context = _ctx(["k1"]) + query_context.get_cache_timeout.return_value = -1 + assert ChartWarmUpCacheCommand._missing_cache_keys(query_context) == [] + + +def test_run_propagates_soft_time_limit() -> None: + """A Celery soft time limit must reach the task instead of becoming a viz error.""" + from celery.exceptions import SoftTimeLimitExceeded + + chart = Mock(spec=Slice) + command = ChartWarmUpCacheCommand(chart, None, None) + with ( + patch.object(command, "validate"), + patch.object( + command, "_warm_up_non_legacy_cache", side_effect=SoftTimeLimitExceeded() + ), + pytest.raises(SoftTimeLimitExceeded), + ): + command.run() diff --git a/tests/unit_tests/tasks/test_cache.py b/tests/unit_tests/tasks/test_cache.py index 172b20009dca..94d1e211ab7f 100644 --- a/tests/unit_tests/tasks/test_cache.py +++ b/tests/unit_tests/tasks/test_cache.py @@ -68,7 +68,9 @@ def test_cache_warmup_user_not_found(app_context: None) -> None: "superset.tasks.cache.current_app", _fake_app({"SUPERSET_CACHE_WARMUP_USER": "bot"}), ), - mock.patch("superset.tasks.cache.security_manager") as mock_sm, + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ) as mock_sm, mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, ): mock_sm.find_user = mock.MagicMock(return_value=None) @@ -92,7 +94,9 @@ def test_cache_warmup_happy_path(app_context: None) -> None: "superset.tasks.cache.current_app", _fake_app({"SUPERSET_CACHE_WARMUP_USER": "bot"}), ), - mock.patch("superset.tasks.cache.security_manager") as mock_sm, + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ) as mock_sm, mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, mock.patch("superset.tasks.cache.DummyStrategy.get_urls", return_value=urls), ): @@ -120,7 +124,9 @@ def side_effect(url: str, _element: str, **kwargs: Any) -> Any: "superset.tasks.cache.current_app", _fake_app({"SUPERSET_CACHE_WARMUP_USER": "bot"}), ), - mock.patch("superset.tasks.cache.security_manager") as mock_sm, + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ) as mock_sm, mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, mock.patch("superset.tasks.cache.DummyStrategy.get_urls", return_value=urls), ): @@ -152,7 +158,9 @@ def side_effect(url: str, _element: str, **kwargs: Any) -> bytes | None: "superset.tasks.cache.current_app", _fake_app({"SUPERSET_CACHE_WARMUP_USER": "bot"}), ), - mock.patch("superset.tasks.cache.security_manager") as mock_sm, + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ) as mock_sm, mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, mock.patch("superset.tasks.cache.DummyStrategy.get_urls", return_value=urls), ): @@ -429,3 +437,197 @@ def test_native_filter_options_strategy_registered() -> None: from superset.tasks.cache import NativeFilterOptionsStrategy, strategy_registry assert strategy_registry["native_filter_options"] is NativeFilterOptionsStrategy + + +def test_favorites_strategy_registered() -> None: + from superset.tasks.cache import FavoritesStrategy, strategy_registry + + assert strategy_registry["favorites"] is FavoritesStrategy + assert FavoritesStrategy.per_user is True + + +def test_favorites_dispatch_disabled(app_context: None) -> None: + """With the gate off, the strategy enqueues nothing and needs no warmup user.""" + from superset.tasks.cache import cache_warmup + + with ( + mock.patch( + "superset.tasks.cache.current_app", + _fake_app({"CACHE_WARMUP_FAVORITES": False}), + ), + mock.patch("superset.tasks.cache.cache_warmup_user_favorites") as task, + mock.patch("superset.tasks.cache.FavoritesStrategy.get_user_ids") as users, + ): + result = cache_warmup("favorites") + + assert isinstance(result, str) + assert "CACHE_WARMUP_FAVORITES is disabled" in result + task.apply_async.assert_not_called() + users.assert_not_called() + + +def test_favorites_dispatch_enqueues_one_task_per_user(app_context: None) -> None: + from superset.tasks.cache import cache_warmup + + config = {"CACHE_WARMUP_FAVORITES": True, "CACHE_WARMUP_FAVORITES_QUEUE": "warm"} + with ( + mock.patch("superset.tasks.cache.current_app", _fake_app(config)), + mock.patch("superset.tasks.cache.cache_warmup_user_favorites") as task, + mock.patch( + "superset.tasks.cache.FavoritesStrategy.get_user_ids", return_value=[3, 4] + ), + ): + task.apply_async.side_effect = [None, RuntimeError("broker")] + result = cache_warmup("favorites", max_users=2) + + assert result == {"success": ["user:3"], "errors": ["user:4"]} + task.apply_async.assert_any_call(args=(3,), queue="warm", expires=7200) + + +def test_user_favorites_task_warms_stale_charts_only(app_context: None) -> None: + """Every favorited chart goes through the stale-only command, capped by config.""" + from superset.tasks.cache import cache_warmup_user_favorites + + chart_a, chart_b = mock.MagicMock(id=1), mock.MagicMock(id=2) + dashboard = mock.MagicMock(id=10, slices=[chart_a, chart_b]) + config = { + "CACHE_WARMUP_FAVORITES": True, + "CACHE_WARMUP_FAVORITES_MAX_DASHBOARDS": 3, + } + with ( + mock.patch("superset.tasks.cache.current_app", _fake_app(config)), + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ), + mock.patch( + "superset.tasks.cache._get_favorite_dashboards", return_value=[dashboard] + ) as favs, + mock.patch( + "superset.commands.chart.warm_up_cache.ChartWarmUpCacheCommand" + ) as command, + ): + command.return_value.run.side_effect = [ + {"chart_id": 1, "viz_error": None}, + {"chart_id": 2, "viz_error": "boom"}, + ] + result = cache_warmup_user_favorites(5) + + favs.assert_called_once_with(5, 3) + command.assert_any_call(chart_a, 10, None, only_if_stale=True) + assert result == {"success": [1], "errors": [2]} + + +def test_user_favorites_task_disabled(app_context: None) -> None: + from superset.tasks.cache import cache_warmup_user_favorites + + with ( + mock.patch( + "superset.tasks.cache.current_app", + _fake_app({"CACHE_WARMUP_FAVORITES": False}), + ), + mock.patch("superset.tasks.cache._get_favorite_dashboards") as favs, + ): + assert cache_warmup_user_favorites(5) == {"success": [], "errors": []} + favs.assert_not_called() + + +def test_user_favorites_task_unknown_user(app_context: None) -> None: + from superset.tasks.cache import cache_warmup_user_favorites + + with ( + mock.patch( + "superset.tasks.cache.current_app", + _fake_app({"CACHE_WARMUP_FAVORITES": True}), + ), + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ) as sm, + mock.patch("superset.tasks.cache._get_favorite_dashboards") as favs, + ): + sm.get_user_by_id.return_value = None + assert cache_warmup_user_favorites(99) == {"success": [], "errors": []} + favs.assert_not_called() + + +def test_get_favorite_dashboards_order_cap_and_access(app_context: None) -> None: + """Newest favorites first, inaccessible dashboards skipped, capped at limit.""" + from superset.exceptions import SupersetSecurityException + from superset.tasks.cache import _get_favorite_dashboards + + dashboards = {i: mock.MagicMock(id=i) for i in (1, 2, 3, 4)} + + def raise_for_access(dashboard: Any) -> None: + if dashboard.id == 2: + raise SupersetSecurityException(mock.MagicMock()) + + with ( + mock.patch("superset.tasks.cache.db") as db, + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ) as sm, + ): + sm.raise_for_access.side_effect = raise_for_access + query = db.session.query.return_value + # Favorites query (newest first), then dashboards lookup. + favs = query.filter.return_value.order_by.return_value.limit + favs.return_value.all.return_value = [ + mock.MagicMock(obj_id=i) for i in (3, 2, 1, 4) + ] + query.options.return_value.filter.return_value.all.return_value = list( + dashboards.values() + ) + result = _get_favorite_dashboards(5, 2) + + favs.assert_called_once_with(8) + + assert [d.id for d in result] == [3, 1] + + +def test_favorites_strategy_user_selection(app_context: None) -> None: + """Users come back newest login first, limited by max_users, from one query.""" + from superset.tasks.cache import FavoritesStrategy + + with ( + mock.patch("superset.tasks.cache.db") as db, + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ) as sm, + ): + sm.user_model.last_login.__ge__.return_value = True + query = db.session.query.return_value.join.return_value.filter.return_value + chain = query.filter.return_value.group_by.return_value.order_by.return_value + chain.limit.return_value.all.return_value = [ + mock.MagicMock(user_id=9), + mock.MagicMock(user_id=4), + ] + strategy = FavoritesStrategy(since="7 days ago", max_users=2) + assert strategy.get_user_ids() == [9, 4] + + chain.limit.assert_called_once_with(2) + + +def test_user_favorites_task_skips_inaccessible_chart(app_context: None) -> None: + """A chart the user cannot access is skipped without counting as an error.""" + from superset.commands.chart.exceptions import ChartAccessDeniedError + from superset.tasks.cache import cache_warmup_user_favorites + + chart = mock.MagicMock(id=1) + dashboard = mock.MagicMock(id=10, slices=[chart]) + config = { + "CACHE_WARMUP_FAVORITES": True, + "CACHE_WARMUP_FAVORITES_MAX_DASHBOARDS": 3, + } + with ( + mock.patch("superset.tasks.cache.current_app", _fake_app(config)), + mock.patch( + "superset.tasks.cache.security_manager", new_callable=mock.MagicMock + ), + mock.patch( + "superset.tasks.cache._get_favorite_dashboards", return_value=[dashboard] + ), + mock.patch( + "superset.commands.chart.warm_up_cache.ChartWarmUpCacheCommand" + ) as command, + ): + command.return_value.run.side_effect = ChartAccessDeniedError() + assert cache_warmup_user_favorites(5) == {"success": [], "errors": []}