This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch ScreenshotCachePayload-serialization in repository https://gitbox.apache.org/repos/asf/superset.git
commit 417aca275974fc356de8a00d1b13385f5701df73 Author: Beto Dealmeida <[email protected]> AuthorDate: Wed Feb 5 17:45:36 2025 -0500 fix: ScreenshotCachePayload serialization --- superset/charts/api.py | 4 +-- superset/dashboards/api.py | 4 +-- superset/db_engine_specs/base.py | 7 ++-- superset/models/core.py | 22 +++++-------- superset/utils/screenshots.py | 55 ++++++++++++++++++++++++------- tests/unit_tests/utils/screenshot_test.py | 11 +++---- 6 files changed, 65 insertions(+), 38 deletions(-) diff --git a/superset/charts/api.py b/superset/charts/api.py index a600a3ca7f..292575feac 100644 --- a/superset/charts/api.py +++ b/superset/charts/api.py @@ -623,7 +623,7 @@ class ChartRestApi(BaseSupersetModelRestApi): if cache_payload.should_trigger_task(force): logger.info("Triggering screenshot ASYNC") - screenshot_obj.cache.set(cache_key, ScreenshotCachePayload()) + screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict()) cache_chart_thumbnail.delay( current_user=get_current_user(), chart_id=chart.id, @@ -755,7 +755,7 @@ class ChartRestApi(BaseSupersetModelRestApi): logger.info( "Triggering thumbnail compute (chart id: %s) ASYNC", str(chart.id) ) - screenshot_obj.cache.set(cache_key, ScreenshotCachePayload()) + screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict()) cache_chart_thumbnail.delay( current_user=current_user, chart_id=chart.id, diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py index c8c744ec63..c15610010b 100644 --- a/superset/dashboards/api.py +++ b/superset/dashboards/api.py @@ -1115,7 +1115,7 @@ class DashboardRestApi(BaseSupersetModelRestApi): if cache_payload.should_trigger_task(force): logger.info("Triggering screenshot ASYNC") - screenshot_obj.cache.set(cache_key, ScreenshotCachePayload()) + screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict()) cache_dashboard_screenshot.delay( username=get_current_user(), guest_token=( @@ -1296,7 +1296,7 @@ class DashboardRestApi(BaseSupersetModelRestApi): "Triggering thumbnail compute (dashboard id: %s) ASYNC", str(dashboard.id), ) - screenshot_obj.cache.set(cache_key, ScreenshotCachePayload()) + 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/db_engine_specs/base.py b/superset/db_engine_specs/base.py index f239ef2019..5d8682b7a2 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -77,7 +77,7 @@ from superset.utils.core import ColumnSpec, GenericDataType from superset.utils.hashing import md5_sha_from_str from superset.utils.json import redact_sensitive, reveal_sensitive from superset.utils.network import is_hostname_valid, is_port_open -from superset.utils.oauth2 import encode_oauth2_state +from superset.utils.oauth2 import check_for_oauth2, encode_oauth2_state if TYPE_CHECKING: from superset.connectors.sqla.models import TableColumn @@ -1861,10 +1861,9 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods if cls.arraysize: cursor.arraysize = cls.arraysize try: - cursor.execute(query) + with check_for_oauth2(database): + cursor.execute(query) except Exception as ex: - if database.is_oauth2_enabled() and cls.needs_oauth2(ex): - cls.start_oauth2_dance(database) raise cls.get_dbapi_mapped_exception(ex) from ex @classmethod diff --git a/superset/models/core.py b/superset/models/core.py index 96a1953fae..6ad87f61ef 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -889,15 +889,13 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable :return: schema list """ try: - with self.get_inspector( - catalog=catalog, - ssh_tunnel=ssh_tunnel, - ) as inspector: - return self.db_engine_spec.get_schema_names(inspector) + with check_for_oauth2(self): + with self.get_inspector( + catalog=catalog, + ssh_tunnel=ssh_tunnel, + ) as inspector: + return self.db_engine_spec.get_schema_names(inspector) except Exception as ex: - if self.is_oauth2_enabled() and self.db_engine_spec.needs_oauth2(ex): - self.start_oauth2_dance() - raise self.db_engine_spec.get_dbapi_mapped_exception(ex) from ex @cache_util.memoized_func( @@ -916,12 +914,10 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable :return: catalog list """ try: - with self.get_inspector(ssh_tunnel=ssh_tunnel) as inspector: - return self.db_engine_spec.get_catalog_names(self, inspector) + with check_for_oauth2(self): + with self.get_inspector(ssh_tunnel=ssh_tunnel) as inspector: + return self.db_engine_spec.get_catalog_names(self, inspector) except Exception as ex: - if self.is_oauth2_enabled() and self.db_engine_spec.needs_oauth2(ex): - self.start_oauth2_dance() - raise self.db_engine_spec.get_dbapi_mapped_exception(ex) from ex @property diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py index 86f5a94ce7..d998d928fe 100644 --- a/superset/utils/screenshots.py +++ b/superset/utils/screenshots.py @@ -20,7 +20,7 @@ import logging from datetime import datetime from enum import Enum from io import BytesIO -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypedDict from flask import current_app @@ -63,13 +63,37 @@ class StatusValues(Enum): ERROR = "Error" +class ScreenshotCachePayloadType(TypedDict): + image: bytes | None + timestamp: str + status: str + + class ScreenshotCachePayload: - def __init__(self, image: bytes | None = None): + def __init__( + self, + image: bytes | None = None, + status: str = StatusValues.PENDING, + timestamp: str = "", + ): self._image = image - self._timestamp = datetime.now().isoformat() - self.status = StatusValues.PENDING - if image: - self.status = StatusValues.UPDATED + self._timestamp = timestamp or datetime.now().isoformat() + self.status = StatusValues.UPDATED if image else status + + @classmethod + def from_dict(cls, payload: ScreenshotCachePayloadType) -> ScreenshotCachePayload: + return cls( + image=payload["image"], + status=StatusValues(payload["status"]), + timestamp=payload["timestamp"], + ) + + def to_dict(self) -> ScreenshotCachePayloadType: + return { + "image": self._image, + "timestamp": self._timestamp, + "status": self.status.value, + } def update_timestamp(self) -> None: self._timestamp = datetime.now().isoformat() @@ -171,14 +195,23 @@ class BaseScreenshot: thumb_size: WindowSize | None = None, ) -> ScreenshotCachePayload | None: cache_key = self.get_cache_key(window_size, thumb_size) - return self.get_from_cache_key(cache_key) + payload = self.get_from_cache_key(cache_key) + return ( + payload + if isinstance(payload, ScreenshotCachePayload) + else ScreenshotCachePayload.from_dict(payload) + ) @classmethod def get_from_cache_key(cls, cache_key: str) -> ScreenshotCachePayload | None: logger.info("Attempting to get from cache: %s", cache_key) if payload := cls.cache.get(cache_key): - # for backwards compatability, byte objects should be converted - if not isinstance(payload, ScreenshotCachePayload): + # Initially, only bytes were stored. This was changed to store an instance + # of ScreenshotCachePayload, but since it can't be serialized in all + # backends it was further changed to just a dict. + if isinstance(payload, dict): + payload = ScreenshotCachePayload.from_dict(payload) + elif isinstance(payload, bytes): payload = ScreenshotCachePayload(payload) return payload logger.info("Failed at getting from cache: %s", cache_key) @@ -217,7 +250,7 @@ class BaseScreenshot: 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) + self.cache.set(cache_key, cache_payload.to_dict()) image = None # Assuming all sorts of things can go wrong with Selenium try: @@ -239,7 +272,7 @@ class BaseScreenshot: logger.info("Caching thumbnail: %s", cache_key) with event_logger.log_context(f"screenshot.cache.{self.thumbnail_type}"): cache_payload.update(image) - self.cache.set(cache_key, cache_payload) + self.cache.set(cache_key, cache_payload.to_dict()) logger.info("Updated thumbnail cache; Status: %s", cache_payload.get_status()) return diff --git a/tests/unit_tests/utils/screenshot_test.py b/tests/unit_tests/utils/screenshot_test.py index 5d29d829a2..1af00a3636 100644 --- a/tests/unit_tests/utils/screenshot_test.py +++ b/tests/unit_tests/utils/screenshot_test.py @@ -26,7 +26,6 @@ from superset.utils.hashing import md5_sha_from_dict from superset.utils.screenshots import ( BaseScreenshot, ScreenshotCachePayload, - StatusValues, ) BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot" @@ -122,7 +121,7 @@ class TestComputeAndCache: self._setup_compute_and_cache(mocker, screenshot_obj) screenshot_obj.compute_and_cache(force=False) cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload.status == StatusValues.UPDATED + assert cache_payload["status"] == "Updated" def test_screenshot_error(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj) @@ -130,7 +129,7 @@ class TestComputeAndCache: get_screenshot.side_effect = Exception screenshot_obj.compute_and_cache(force=False) cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload.status == StatusValues.ERROR + assert cache_payload["status"] == "Error" def test_resize_error(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj) @@ -138,7 +137,7 @@ class TestComputeAndCache: resize_image.side_effect = Exception screenshot_obj.compute_and_cache(force=False) cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload.status == StatusValues.ERROR + assert cache_payload["status"] == "Error" def test_skips_if_computing(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj) @@ -156,7 +155,7 @@ class TestComputeAndCache: screenshot_obj.compute_and_cache(force=True) get_screenshot.assert_called_once() cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload.status == StatusValues.UPDATED + assert cache_payload["status"] == "Updated" def test_skips_if_updated(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj) @@ -178,7 +177,7 @@ class TestComputeAndCache: ) get_screenshot.assert_called_once() cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload._image != b"initial_value" + assert cache_payload["image"] != b"initial_value" def test_resize(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj)
