You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by mi...@apache.org on 2026/06/26 16:41:10 UTC
(superset) branch master updated: fix(thumbnails): add deduplication to dashboard thumbnail Celery tasks (#38576)
This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new a5dbb394e5f fix(thumbnails): add deduplication to dashboard thumbnail Celery tasks (#38576)
a5dbb394e5f is described below
commit a5dbb394e5f80a30e36c626178ad90bfe595d378
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Fri Jun 26 10:41:02 2026 -0600
fix(thumbnails): add deduplication to dashboard thumbnail Celery tasks (#38576)
Co-authored-by: Claude Sonnet 4.6 <no...@anthropic.com>
---
superset/commands/distributed_lock/acquire.py | 36 ++---
superset/config.py | 3 +
superset/dashboards/api.py | 2 -
superset/distributed_lock/__init__.py | 6 +-
superset/exceptions.py | 8 ++
superset/tasks/thumbnails.py | 10 +-
superset/utils/screenshots.py | 104 +++++++++------
tests/unit_tests/tasks/test_thumbnails.py | 145 +++++++++++++++++++++
.../unit_tests/utils/test_screenshot_cache_fix.py | 115 +++++++++++-----
9 files changed, 335 insertions(+), 94 deletions(-)
diff --git a/superset/commands/distributed_lock/acquire.py b/superset/commands/distributed_lock/acquire.py
index c25501f7927..6cd9a878099 100644
--- a/superset/commands/distributed_lock/acquire.py
+++ b/superset/commands/distributed_lock/acquire.py
@@ -19,11 +19,10 @@ from __future__ import annotations
import logging
from datetime import datetime, timedelta, timezone
-from functools import partial
from typing import Any
import redis
-from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.exc import IntegrityError, SQLAlchemyError
from superset.commands.distributed_lock.base import (
BaseDistributedLockCommand,
@@ -31,17 +30,32 @@ from superset.commands.distributed_lock.base import (
get_redis_client,
)
from superset.daos.key_value import KeyValueDAO
-from superset.exceptions import AcquireDistributedLockFailedException
+from superset.exceptions import (
+ AcquireDistributedLockFailedException,
+ LockAlreadyHeldException,
+)
from superset.key_value.exceptions import (
KeyValueCodecEncodeException,
KeyValueUpsertFailedError,
)
from superset.key_value.types import KeyValueResource
-from superset.utils.decorators import on_error, transaction
+from superset.utils.decorators import transaction
logger = logging.getLogger(__name__)
+def _kv_lock_on_error(ex: Exception) -> None:
+ """Translate KV-backend exceptions into the appropriate lock exception type."""
+ if isinstance(ex, IntegrityError):
+ # Unique-constraint violation: lock is already held by another process.
+ raise LockAlreadyHeldException("KV lock already held") from ex
+ if isinstance(
+ ex, (KeyValueCodecEncodeException, KeyValueUpsertFailedError, SQLAlchemyError)
+ ):
+ raise AcquireDistributedLockFailedException(f"KV lock failed: {ex}") from ex
+ raise ex
+
+
class AcquireDistributedLock(BaseDistributedLockCommand):
"""
Acquire a distributed lock with automatic backend selection.
@@ -85,7 +99,7 @@ class AcquireDistributedLock(BaseDistributedLockCommand):
if not acquired:
logger.debug("Redis lock on %s already taken", self.redis_lock_key)
- raise AcquireDistributedLockFailedException("Lock already taken")
+ raise LockAlreadyHeldException("Lock already taken")
logger.debug(
"Acquired Redis lock: %s (TTL=%ds)",
@@ -99,17 +113,7 @@ class AcquireDistributedLock(BaseDistributedLockCommand):
f"Redis lock failed: {ex}"
) from ex
- @transaction(
- on_error=partial(
- on_error,
- catches=(
- KeyValueCodecEncodeException,
- KeyValueUpsertFailedError,
- SQLAlchemyError,
- ),
- reraise=AcquireDistributedLockFailedException,
- ),
- )
+ @transaction(on_error=_kv_lock_on_error)
def _acquire_kv(self) -> None:
"""Acquire lock using KeyValue table (database)."""
# Delete expired entries first to prevent stale locks from blocking
diff --git a/superset/config.py b/superset/config.py
index b5e9618a195..474b5250bf8 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -1183,6 +1183,9 @@ THUMBNAIL_CACHE_CONFIG: CacheConfig = {
"CACHE_NO_NULL_WARNING": True,
}
THUMBNAIL_ERROR_CACHE_TTL = int(timedelta(days=1).total_seconds())
+# How long to treat a COMPUTING cache entry as an active lease before considering
+# the worker stuck. Should exceed the task soft_time_limit (300 s) by a margin.
+THUMBNAIL_COMPUTING_CACHE_TTL = int(timedelta(seconds=360).total_seconds())
# Cache warmup user — must be set explicitly before enabling the cache-warmup
# Celery task. Intentionally defaults to None so operators pick a dedicated
diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index b8c437b7416..e90a3316700 100644
--- a/superset/dashboards/api.py
+++ b/superset/dashboards/api.py
@@ -1475,7 +1475,6 @@ class DashboardRestApi(CustomTagsOptimizationMixin, BaseSupersetModelRestApi):
if cache_payload.should_trigger_task(force):
logger.info("Triggering screenshot ASYNC")
- screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict())
cache_dashboard_screenshot.delay(
username=get_current_user(),
guest_token=(
@@ -1671,7 +1670,6 @@ class DashboardRestApi(CustomTagsOptimizationMixin, BaseSupersetModelRestApi):
"Triggering thumbnail compute (dashboard id: %s) ASYNC",
str(dashboard.id),
)
- screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict())
cache_dashboard_thumbnail.delay(
current_user=current_user,
dashboard_id=dashboard.id,
diff --git a/superset/distributed_lock/__init__.py b/superset/distributed_lock/__init__.py
index ab00143ba37..df7240d8dfd 100644
--- a/superset/distributed_lock/__init__.py
+++ b/superset/distributed_lock/__init__.py
@@ -46,8 +46,10 @@ def DistributedLock( # noqa: N802
to prevent deadlocks from crashed processes.
:param kwargs: Additional key parameters to differentiate locks
:yields: UUID identifying this lock acquisition
- :raises AcquireDistributedLockFailedException: If lock is already held
- or Redis connection fails
+ :raises LockAlreadyHeldException: If the lock is already held by another process
+ (subclass of AcquireDistributedLockFailedException)
+ :raises AcquireDistributedLockFailedException: If the lock backend fails
+ (Redis/DB connection error)
"""
# pylint: disable=import-outside-toplevel
from superset.commands.distributed_lock.acquire import AcquireDistributedLock
diff --git a/superset/exceptions.py b/superset/exceptions.py
index 15fe372abe6..be21c4a8aac 100644
--- a/superset/exceptions.py
+++ b/superset/exceptions.py
@@ -441,6 +441,14 @@ class AcquireDistributedLockFailedException(Exception): # noqa: N818
"""
+class LockAlreadyHeldException(AcquireDistributedLockFailedException): # noqa: N818
+ """
+ Raised when a distributed lock is already held by another process (lock contention).
+ Subclass of AcquireDistributedLockFailedException so existing callers that catch
+ the base exception continue to work unchanged.
+ """
+
+
class ReleaseDistributedLockFailedException(Exception): # noqa: N818
"""
Exception to signalize failure to release lock.
diff --git a/superset/tasks/thumbnails.py b/superset/tasks/thumbnails.py
index 9ff5ccfedf8..c80630e3114 100644
--- a/superset/tasks/thumbnails.py
+++ b/superset/tasks/thumbnails.py
@@ -27,7 +27,10 @@ from superset.extensions import celery_app
from superset.security.guest_token import GuestToken
from superset.tasks.utils import get_executor
from superset.utils.core import override_user
-from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
+from superset.utils.screenshots import (
+ ChartScreenshot,
+ DashboardScreenshot,
+)
from superset.utils.urls import get_url_path
from superset.utils.webdriver import WindowSize
@@ -99,12 +102,15 @@ def cache_dashboard_thumbnail(
user = security_manager.find_user(username)
with override_user(user):
screenshot = DashboardScreenshot(url, dashboard.digest)
+ resolved_cache_key = cache_key or screenshot.get_cache_key(
+ window_size, thumb_size
+ )
screenshot.compute_and_cache(
user=user,
window_size=window_size,
thumb_size=thumb_size,
force=force,
- cache_key=cache_key,
+ cache_key=resolved_cache_key,
)
diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py
index 3d680faf171..be965a5aee3 100644
--- a/superset/utils/screenshots.py
+++ b/superset/utils/screenshots.py
@@ -26,7 +26,11 @@ from typing import cast, TYPE_CHECKING, TypedDict
from flask import current_app as app
from superset import feature_flag_manager, thumbnail_cache
-from superset.exceptions import ScreenshotImageNotAvailableException
+from superset.distributed_lock import DistributedLock
+from superset.exceptions import (
+ LockAlreadyHeldException,
+ ScreenshotImageNotAvailableException,
+)
from superset.extensions import event_logger
from superset.utils.hashing import hash_from_dict
from superset.utils.urls import modify_url_query
@@ -119,7 +123,6 @@ class ScreenshotCachePayload:
def computing(self) -> None:
self.update_timestamp()
- self._image = None
self.status = StatusValues.COMPUTING
def update(self, image: bytes) -> None:
@@ -152,9 +155,7 @@ class ScreenshotCachePayload:
def is_computing_stale(self) -> bool:
"""Check if a COMPUTING status is stale (task likely failed or stuck)."""
- # Use the same TTL as error cache - if computing takes longer than this,
- # it's likely stuck and should be retried
- computing_ttl = app.config["THUMBNAIL_ERROR_CACHE_TTL"]
+ computing_ttl = app.config["THUMBNAIL_COMPUTING_CACHE_TTL"]
return (
datetime.now() - datetime.fromisoformat(self.get_timestamp())
).total_seconds() >= computing_ttl
@@ -280,43 +281,66 @@ class BaseScreenshot:
:return: Image payload
"""
cache_key = cache_key or self.get_cache_key(window_size, thumb_size)
- cache_payload = self.get_from_cache_key(cache_key) or ScreenshotCachePayload()
- if not cache_payload.should_trigger_task(force=force):
- logger.info(
- "Skipping compute - already processed for thumbnail: %s", cache_key
- )
- return
-
- window_size = window_size or self.window_size
- thumb_size = thumb_size or self.thumb_size
- logger.info("Processing url for thumbnail: %s", cache_key)
- cache_payload.computing()
- image = None
- # Assuming all sorts of things can go wrong with Selenium
try:
- logger.info("trying to generate screenshot")
- with event_logger.log_context(f"screenshot.compute.{self.thumbnail_type}"):
- image = self.get_screenshot(user=user, window_size=window_size)
- except Exception as ex: # pylint: disable=broad-except
- logger.warning("Failed at generating thumbnail %s", ex, exc_info=True)
- cache_payload.error()
- if image and window_size != thumb_size:
- try:
- image = self.resize_image(image, thumb_size=thumb_size)
- except Exception as ex: # pylint: disable=broad-except
- logger.warning("Failed at resizing thumbnail %s", ex, exc_info=True)
- cache_payload.error()
+ with DistributedLock(
+ namespace="thumbnail",
+ key=cache_key,
+ ttl_seconds=app.config["THUMBNAIL_COMPUTING_CACHE_TTL"],
+ ):
+ cache_payload = (
+ self.get_from_cache_key(cache_key) or ScreenshotCachePayload()
+ )
+ if not cache_payload.should_trigger_task(force=force):
+ logger.info(
+ "Skipping compute - already processed for thumbnail: %s",
+ cache_key,
+ )
+ return
+
+ window_size = window_size or self.window_size
+ thumb_size = thumb_size or self.thumb_size
+ logger.info("Processing url for thumbnail: %s", cache_key)
+ cache_payload.computing()
+ self.cache.set(cache_key, cache_payload.to_dict())
image = None
-
- # Cache the result (success or error) to avoid immediate retries
- if image:
- with event_logger.log_context(f"screenshot.cache.{self.thumbnail_type}"):
- cache_payload.update(image)
-
- logger.info("Caching thumbnail: %s", cache_key)
- self.cache.set(cache_key, cache_payload.to_dict())
- logger.info("Updated thumbnail cache; Status: %s", cache_payload.get_status())
- return
+ # Assuming all sorts of things can go wrong with Selenium
+ try:
+ logger.info("trying to generate screenshot")
+ with event_logger.log_context(
+ f"screenshot.compute.{self.thumbnail_type}"
+ ):
+ image = self.get_screenshot(user=user, window_size=window_size)
+ except Exception as ex: # pylint: disable=broad-except
+ logger.warning(
+ "Failed at generating thumbnail %s", ex, exc_info=True
+ )
+ cache_payload.error()
+ if image and window_size != thumb_size:
+ try:
+ image = self.resize_image(image, thumb_size=thumb_size)
+ except Exception as ex: # pylint: disable=broad-except
+ logger.warning(
+ "Failed at resizing thumbnail %s", ex, exc_info=True
+ )
+ cache_payload.error()
+ image = None
+
+ if image:
+ with event_logger.log_context(
+ f"screenshot.cache.{self.thumbnail_type}"
+ ):
+ cache_payload.update(image)
+
+ logger.info("Caching thumbnail: %s", cache_key)
+ self.cache.set(cache_key, cache_payload.to_dict())
+ logger.info(
+ "Updated thumbnail cache; Status: %s", cache_payload.get_status()
+ )
+ except LockAlreadyHeldException:
+ logger.info(
+ "Skipping duplicate thumbnail task for %s - lock already held",
+ cache_key,
+ )
@classmethod
def resize_image(
diff --git a/tests/unit_tests/tasks/test_thumbnails.py b/tests/unit_tests/tasks/test_thumbnails.py
new file mode 100644
index 00000000000..8dee18149f6
--- /dev/null
+++ b/tests/unit_tests/tasks/test_thumbnails.py
@@ -0,0 +1,145 @@
+# 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.
+
+from collections.abc import Callable, Iterator
+from typing import Any
+from unittest.mock import ANY, MagicMock, patch
+
+import pytest
+
+
+@pytest.fixture
+def mock_thumbnail_cache() -> Iterator[None]:
+ """Enable thumbnail cache mocking so tasks don't exit early."""
+ with patch("superset.tasks.thumbnails.thumbnail_cache", True):
+ yield
+
+
+@pytest.fixture
+def mock_dashboard() -> MagicMock:
+ """Return a mocked Dashboard with id=1 and a fixed digest."""
+ dashboard = MagicMock()
+ dashboard.id = 1
+ dashboard.digest = "test_digest"
+ return dashboard
+
+
+def _make_screenshot_mock(cache_key: str = "test_cache_key") -> MagicMock:
+ """Return a MagicMock screenshot whose get_cache_key returns cache_key."""
+ screenshot = MagicMock()
+ screenshot.get_cache_key.return_value = cache_key
+ return screenshot
+
+
+# ---------------------------------------------------------------------------
+# Helpers shared across tests
+# ---------------------------------------------------------------------------
+
+_COMMON_PATCHES = [
+ "superset.tasks.thumbnails.get_executor",
+ "superset.tasks.thumbnails.security_manager",
+ "superset.tasks.thumbnails.override_user",
+ "superset.tasks.thumbnails.get_url_path",
+ "superset.tasks.thumbnails.DashboardScreenshot",
+]
+
+
+def _apply_patches(fn: Callable[..., Any]) -> Callable[..., Any]:
+ """Stack the five common patches onto a test function."""
+ for p in reversed(_COMMON_PATCHES):
+ fn = patch(p)(fn)
+ return fn
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@patch("superset.tasks.thumbnails.thumbnail_cache", None)
+def test_cache_dashboard_thumbnail_skips_when_no_cache() -> None:
+ """Task exits early when no thumbnail cache is configured."""
+ from superset.tasks.thumbnails import cache_dashboard_thumbnail
+
+ result = cache_dashboard_thumbnail(current_user="admin", dashboard_id=1, force=True)
+
+ assert result is None
+
+
+@_apply_patches
+def test_cache_dashboard_thumbnail_always_delegates_to_compute_and_cache(
+ mock_screenshot_cls: MagicMock,
+ mock_get_url_path: MagicMock,
+ mock_override_user: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_executor: MagicMock,
+ mock_dashboard: MagicMock,
+ mock_thumbnail_cache: None,
+) -> None:
+ """Task always calls compute_and_cache; dedup lives inside compute_and_cache."""
+ from superset.tasks.thumbnails import cache_dashboard_thumbnail
+
+ mock_screenshot = _make_screenshot_mock()
+ mock_screenshot_cls.return_value = mock_screenshot
+ mock_get_executor.return_value = (None, "admin")
+ mock_security_manager.find_user.return_value = MagicMock()
+
+ with patch("superset.models.dashboard.Dashboard.get", return_value=mock_dashboard):
+ cache_dashboard_thumbnail(
+ current_user="admin",
+ dashboard_id=1,
+ force=False,
+ cache_key="test_cache_key",
+ )
+
+ mock_screenshot.compute_and_cache.assert_called_once()
+
+
+@_apply_patches
+def test_cache_dashboard_thumbnail_resolves_cache_key_when_not_provided(
+ mock_screenshot_cls: MagicMock,
+ mock_get_url_path: MagicMock,
+ mock_override_user: MagicMock,
+ mock_security_manager: MagicMock,
+ mock_get_executor: MagicMock,
+ mock_dashboard: MagicMock,
+ mock_thumbnail_cache: None,
+) -> None:
+ """Task resolves cache_key from the screenshot object when none is given."""
+ from superset.tasks.thumbnails import cache_dashboard_thumbnail
+
+ mock_screenshot = _make_screenshot_mock("resolved_cache_key")
+ mock_screenshot_cls.return_value = mock_screenshot
+ mock_get_executor.return_value = (None, "admin")
+ mock_security_manager.find_user.return_value = MagicMock()
+
+ with patch("superset.models.dashboard.Dashboard.get", return_value=mock_dashboard):
+ cache_dashboard_thumbnail(
+ current_user="admin",
+ dashboard_id=1,
+ force=False,
+ cache_key=None,
+ )
+
+ mock_screenshot.get_cache_key.assert_called_once()
+ mock_screenshot.compute_and_cache.assert_called_once_with(
+ user=ANY,
+ window_size=None,
+ thumb_size=None,
+ force=False,
+ cache_key="resolved_cache_key",
+ )
diff --git a/tests/unit_tests/utils/test_screenshot_cache_fix.py b/tests/unit_tests/utils/test_screenshot_cache_fix.py
index 8d0946abdf9..cf62a01d348 100644
--- a/tests/unit_tests/utils/test_screenshot_cache_fix.py
+++ b/tests/unit_tests/utils/test_screenshot_cache_fix.py
@@ -27,6 +27,7 @@ from unittest.mock import MagicMock, patch
import pytest
from pytest_mock import MockerFixture
+from superset.exceptions import LockAlreadyHeldException
from superset.utils.screenshots import (
BaseScreenshot,
ScreenshotCachePayload,
@@ -34,6 +35,7 @@ from superset.utils.screenshots import (
)
BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot"
+DISTRIBUTED_LOCK_PATH = "superset.utils.screenshots.DistributedLock"
class MockCache:
@@ -56,7 +58,7 @@ class MockCache:
@pytest.fixture
-def mock_user():
+def mock_user() -> MagicMock:
"""Fixture to create a mock user."""
user = MagicMock()
user.id = 1
@@ -64,7 +66,7 @@ def mock_user():
@pytest.fixture
-def screenshot_obj():
+def screenshot_obj() -> BaseScreenshot:
"""Fixture to create a BaseScreenshot object."""
url = "http://example.com"
digest = "sample_digest"
@@ -74,8 +76,11 @@ def screenshot_obj():
class TestCacheOnlyOnSuccess:
"""Test that cache is only saved when image generation succeeds."""
- def _setup_mocks(self, mocker: MockerFixture, screenshot_obj):
+ def _setup_mocks(
+ self, mocker: MockerFixture, screenshot_obj: BaseScreenshot
+ ) -> MagicMock:
"""Helper method to set up common mocks."""
+ mocker.patch(DISTRIBUTED_LOCK_PATH)
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
get_screenshot = mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"image_data"
@@ -91,6 +96,7 @@ class TestCacheOnlyOnSuccess:
self, mocker: MockerFixture, screenshot_obj, mock_user
):
"""Test that error status is cached when screenshot generation fails."""
+ mocker.patch(DISTRIBUTED_LOCK_PATH)
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
get_screenshot = mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot",
@@ -149,38 +155,37 @@ class TestCacheOnlyOnSuccess:
assert cached_value["status"] == "Updated"
assert cached_value["image"] is not None
- def test_no_intermediate_cache_during_computing(
- self, mocker: MockerFixture, screenshot_obj, mock_user
- ):
- """Test that cache is not saved during COMPUTING state."""
+ def test_computing_status_written_to_cache_early(
+ self,
+ mocker: MockerFixture,
+ screenshot_obj: BaseScreenshot,
+ mock_user: MagicMock,
+ ) -> None:
+ """compute_and_cache writes COMPUTING to cache before taking the screenshot
+ so concurrent tasks can detect it and avoid duplicate work."""
+ mocker.patch(DISTRIBUTED_LOCK_PATH)
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
BaseScreenshot.cache = MockCache()
- # Mock get_screenshot to check cache state during execution
- def check_cache_during_screenshot(*args, **kwargs):
- # At this point, we're in COMPUTING state
- # Cache should NOT be set yet
+ def check_cache_during_screenshot(*args: object, **kwargs: object) -> bytes:
cache_key = screenshot_obj.get_cache_key()
cached_value = BaseScreenshot.cache.get(cache_key)
- # Cache should be empty during screenshot generation
- assert cached_value is None, (
- "Cache should not be saved during COMPUTING state"
+ assert cached_value is not None, (
+ "Cache should be set to COMPUTING before screenshot starts"
)
+ assert cached_value["status"] == "Computing"
return b"image_data"
mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot",
side_effect=check_cache_during_screenshot,
)
- # Mock resize to avoid PIL errors with fake image data
mocker.patch(
BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
)
- # Execute compute_and_cache
screenshot_obj.compute_and_cache(user=mock_user, force=True)
- # After completion, cache should be set with UPDATED status
cache_key = screenshot_obj.get_cache_key()
cached_value = BaseScreenshot.cache.get(cache_key)
assert cached_value is not None
@@ -191,10 +196,10 @@ class TestShouldTriggerTask:
"""Test the should_trigger_task method improvements."""
@patch("superset.utils.screenshots.app")
- def test_trigger_on_stale_computing_status(self, mock_app):
+ def test_trigger_on_stale_computing_status(self, mock_app: MagicMock) -> None:
"""Test that stale COMPUTING status triggers recomputation."""
# Set TTL to 300 seconds
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {"THUMBNAIL_COMPUTING_CACHE_TTL": 300}
# Create payload with COMPUTING status from 400 seconds ago (stale)
old_timestamp = (datetime.now() - timedelta(seconds=400)).isoformat()
@@ -209,7 +214,7 @@ class TestShouldTriggerTask:
def test_no_trigger_on_fresh_computing_status(self, mock_app):
"""Test that fresh COMPUTING status does not trigger recomputation."""
# Set TTL to 300 seconds
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {"THUMBNAIL_COMPUTING_CACHE_TTL": 300}
# Create payload with COMPUTING status from 100 seconds ago (fresh)
fresh_timestamp = (datetime.now() - timedelta(seconds=100)).isoformat()
@@ -246,8 +251,10 @@ class TestShouldTriggerTask:
@patch("superset.utils.screenshots.app")
def test_trigger_on_expired_error(self, mock_app):
"""Test that expired ERROR status triggers task."""
- # Set TTL to 300 seconds
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {
+ "THUMBNAIL_COMPUTING_CACHE_TTL": 300,
+ "THUMBNAIL_ERROR_CACHE_TTL": 300,
+ }
# Create payload with ERROR status from 400 seconds ago (expired)
old_timestamp = (datetime.now() - timedelta(seconds=400)).isoformat()
@@ -258,10 +265,12 @@ class TestShouldTriggerTask:
assert payload.should_trigger_task(force=False) is True
@patch("superset.utils.screenshots.app")
- def test_no_trigger_on_fresh_error(self, mock_app):
+ def test_no_trigger_on_fresh_error(self, mock_app: MagicMock) -> None:
"""Test that fresh ERROR status does not trigger task."""
- # Set TTL to 300 seconds
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {
+ "THUMBNAIL_COMPUTING_CACHE_TTL": 300,
+ "THUMBNAIL_ERROR_CACHE_TTL": 300,
+ }
# Create payload with ERROR status from 100 seconds ago (fresh)
fresh_timestamp = (datetime.now() - timedelta(seconds=100)).isoformat()
@@ -288,7 +297,7 @@ class TestIsComputingStale:
@patch("superset.utils.screenshots.app")
def test_computing_is_stale(self, mock_app):
"""Test that old COMPUTING status is detected as stale."""
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {"THUMBNAIL_COMPUTING_CACHE_TTL": 300}
# Timestamp from 400 seconds ago
old_timestamp = (datetime.now() - timedelta(seconds=400)).isoformat()
@@ -301,7 +310,7 @@ class TestIsComputingStale:
@patch("superset.utils.screenshots.app")
def test_computing_is_not_stale(self, mock_app):
"""Test that fresh COMPUTING status is not stale."""
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {"THUMBNAIL_COMPUTING_CACHE_TTL": 300}
# Timestamp from 100 seconds ago
fresh_timestamp = (datetime.now() - timedelta(seconds=100)).isoformat()
@@ -314,7 +323,7 @@ class TestIsComputingStale:
@patch("superset.utils.screenshots.app")
def test_computing_exactly_at_ttl(self, mock_app):
"""Test boundary condition at exactly TTL."""
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {"THUMBNAIL_COMPUTING_CACHE_TTL": 300}
# Timestamp from exactly 300 seconds ago
exact_timestamp = (datetime.now() - timedelta(seconds=300)).isoformat()
@@ -328,7 +337,7 @@ class TestIsComputingStale:
@patch("superset.utils.screenshots.app")
def test_computing_just_past_ttl(self, mock_app):
"""Test boundary condition just past TTL."""
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {"THUMBNAIL_COMPUTING_CACHE_TTL": 300}
# Timestamp from 301 seconds ago (just past TTL)
past_ttl_timestamp = (datetime.now() - timedelta(seconds=301)).isoformat()
@@ -350,6 +359,7 @@ class TestIntegrationCacheBugFix:
Integration test: Failed screenshot should cache error status
to prevent immediate retries, not leave corrupted cache with image=None.
"""
+ mocker.patch(DISTRIBUTED_LOCK_PATH)
mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot",
side_effect=Exception("Network error"),
@@ -373,13 +383,18 @@ class TestIntegrationCacheBugFix:
@patch("superset.utils.screenshots.app")
def test_stale_computing_triggers_retry(
- self, mock_app, mocker: MockerFixture, screenshot_obj, mock_user
- ):
+ self,
+ mock_app: MagicMock,
+ mocker: MockerFixture,
+ screenshot_obj: BaseScreenshot,
+ mock_user: MagicMock,
+ ) -> None:
"""
Integration test: Stale COMPUTING status should trigger retry
to recover from stuck tasks.
"""
- mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
+ mock_app.config = {"THUMBNAIL_COMPUTING_CACHE_TTL": 300}
+ mocker.patch(DISTRIBUTED_LOCK_PATH)
BaseScreenshot.cache = MockCache()
# Create stale COMPUTING entry and seed it in the cache
@@ -408,3 +423,39 @@ class TestIntegrationCacheBugFix:
assert cached_value is not None
assert cached_value["status"] == "Updated"
assert cached_value["image"] is not None
+
+ def test_concurrent_task_skips_when_lock_already_held(
+ self,
+ mocker: MockerFixture,
+ screenshot_obj: BaseScreenshot,
+ mock_user: MagicMock,
+ ) -> None:
+ """compute_and_cache exits without rendering when the distributed lock
+ is already held by another worker — atomically preventing duplicate Selenium."""
+ mock_lock = mocker.patch(DISTRIBUTED_LOCK_PATH)
+ mock_lock.return_value.__enter__.side_effect = LockAlreadyHeldException(
+ "lock held"
+ )
+ get_screenshot = mocker.patch(BASE_SCREENSHOT_PATH + ".get_screenshot")
+ BaseScreenshot.cache = MockCache()
+
+ screenshot_obj.compute_and_cache(user=mock_user, force=False)
+
+ get_screenshot.assert_not_called()
+
+ def test_computing_preserves_previous_image(
+ self,
+ mocker: MockerFixture,
+ screenshot_obj: BaseScreenshot,
+ mock_user: MagicMock,
+ ) -> None:
+ """computing() must not wipe the cached image so a stale thumbnail remains
+ visible while a refresh is in progress."""
+ old_image = b"old_thumbnail_bytes"
+ payload = ScreenshotCachePayload(image=old_image)
+ assert payload._image == old_image
+
+ payload.computing()
+
+ assert payload._image == old_image
+ assert payload.status == StatusValues.COMPUTING