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
The following commit(s) were added to
refs/heads/ScreenshotCachePayload-serialization by this push:
new 993b519e3e fix: ScreenshotCachePayload serialization
993b519e3e is described below
commit 993b519e3ed1bd494328765b6f9a97539cfd8a2b
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)
