diff --git a/superset/charts/api.py b/superset/charts/api.py index a33c0bc943dc..66c0781ffb31 100644 --- a/superset/charts/api.py +++ b/superset/charts/api.py @@ -80,6 +80,7 @@ from superset.commands.importers.v1.utils import get_contents_from_bundle from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod from superset.daos.chart import ChartDAO +from superset.exceptions import ScreenshotImageNotAvailableException from superset.extensions import event_logger from superset.models.slice import Slice from superset.tasks.thumbnails import cache_chart_thumbnail @@ -704,8 +705,12 @@ def screenshot(self, pk: int, digest: str) -> WerkzeugResponse: if cache_payload := ChartScreenshot.get_from_cache_key(digest): if cache_payload.status == StatusValues.UPDATED: + try: + image = cache_payload.get_image() + except ScreenshotImageNotAvailableException: + return self.response_404() return Response( - FileWrapper(cache_payload.get_image()), + FileWrapper(image), mimetype="image/png", direct_passthrough=True, ) @@ -791,8 +796,12 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse: task_status=cache_payload.get_status(), ) self.incr_stats("from_cache", self.thumbnail.__name__) + try: + image = cache_payload.get_image() + except ScreenshotImageNotAvailableException: + return self.response_404() return Response( - FileWrapper(cache_payload.get_image()), + FileWrapper(image), mimetype="image/png", direct_passthrough=True, ) diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py index 54e05df8b755..a739e44c2715 100644 --- a/superset/dashboards/api.py +++ b/superset/dashboards/api.py @@ -101,6 +101,7 @@ TabsPayloadSchema, thumbnail_query_schema, ) +from superset.exceptions import ScreenshotImageNotAvailableException from superset.extensions import event_logger from superset.models.dashboard import Dashboard from superset.models.embedded_dashboard import EmbeddedDashboard @@ -1199,8 +1200,9 @@ def screenshot(self, pk: int, digest: str) -> WerkzeugResponse: # fetch the dashboard screenshot using the current user and cache if set if cache_payload := DashboardScreenshot.get_from_cache_key(digest): - image = cache_payload.get_image() - if not image: + try: + image = cache_payload.get_image() + except ScreenshotImageNotAvailableException: return self.response_404() if download_format == "pdf": pdf_img = image.getvalue() @@ -1324,8 +1326,12 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse: ) self.incr_stats("from_cache", self.thumbnail.__name__) + try: + image = cache_payload.get_image() + except ScreenshotImageNotAvailableException: + return self.response_404() return Response( - FileWrapper(cache_payload.get_image()), + FileWrapper(image), mimetype="image/png", direct_passthrough=True, ) diff --git a/superset/exceptions.py b/superset/exceptions.py index 6dddfb8302cf..1b3216da0c72 100644 --- a/superset/exceptions.py +++ b/superset/exceptions.py @@ -486,3 +486,7 @@ def __init__(self) -> None: level=ErrorLevel.ERROR, ) super().__init__(error) + + +class ScreenshotImageNotAvailableException(SupersetException): + status = 404 diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py index 23493471b847..e11253014ba7 100644 --- a/superset/utils/screenshots.py +++ b/superset/utils/screenshots.py @@ -26,6 +26,7 @@ from flask import current_app as app from superset import feature_flag_manager, thumbnail_cache +from superset.exceptions import ScreenshotImageNotAvailableException from superset.extensions import event_logger from superset.utils.hashing import md5_sha_from_dict from superset.utils.urls import modify_url_query @@ -121,10 +122,10 @@ def error( self.update_timestamp() self.status = StatusValues.ERROR - def get_image(self) -> BytesIO | None: - if not self._image: - return None - return BytesIO(self._image) + def get_image(self) -> BytesIO: + if self._image is None: + raise ScreenshotImageNotAvailableException() + return BytesIO(cast(bytes, self._image)) def get_timestamp(self) -> str: return self._timestamp diff --git a/tests/unit_tests/utils/screenshot_test.py b/tests/unit_tests/utils/screenshot_test.py index fa928fc73fb6..d747a9005531 100644 --- a/tests/unit_tests/utils/screenshot_test.py +++ b/tests/unit_tests/utils/screenshot_test.py @@ -192,3 +192,50 @@ def test_resize(self, mocker: MockerFixture, screenshot_obj): force=False, window_size=(1, 1), thumb_size=thumb_size ) resize_image.assert_called_once() + + +class TestScreenshotCachePayloadGetImage: + """Test the get_image method behavior including exception handling""" + + def test_get_image_returns_bytesio_when_image_exists(self): + """Test that get_image returns BytesIO object when image data exists""" + image_data = b"test image data" + payload = ScreenshotCachePayload(image=image_data) + + result = payload.get_image() + + assert result is not None + assert result.read() == image_data + + def test_get_image_raises_exception_when_no_image(self): + """Test get_image raises ScreenshotImageNotAvailableException when no image""" + from superset.exceptions import ScreenshotImageNotAvailableException + + payload = ScreenshotCachePayload() # No image data + + with pytest.raises(ScreenshotImageNotAvailableException): + payload.get_image() + + def test_get_image_raises_exception_when_image_is_none(self): + """Test that get_image raises exception when image is explicitly set to None""" + from superset.exceptions import ScreenshotImageNotAvailableException + + payload = ScreenshotCachePayload(image=None) + + with pytest.raises(ScreenshotImageNotAvailableException): + payload.get_image() + + def test_get_image_multiple_reads(self): + """Test that get_image returns fresh BytesIO each time""" + image_data = b"test image data" + payload = ScreenshotCachePayload(image=image_data) + + result1 = payload.get_image() + result2 = payload.get_image() + + # Both should be valid BytesIO objects + assert result1.read() == image_data + assert result2.read() == image_data + + # Should be different BytesIO instances + assert result1 is not result2 diff --git a/tests/unit_tests/utils/test_screenshot_exception_handling.py b/tests/unit_tests/utils/test_screenshot_exception_handling.py new file mode 100644 index 000000000000..4d1f6bee7be7 --- /dev/null +++ b/tests/unit_tests/utils/test_screenshot_exception_handling.py @@ -0,0 +1,126 @@ +# 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. + +""" +Tests for screenshot exception handling in API endpoints. +""" + +from unittest.mock import MagicMock + +import pytest + +from superset.exceptions import ScreenshotImageNotAvailableException +from superset.utils.screenshots import ScreenshotCachePayload + + +class TestScreenshotAPIExceptionHandling: + """Test that API endpoints properly handle ScreenshotImageNotAvailableException""" + + def test_dashboard_screenshot_api_handles_exception(self): + """Test dashboard screenshot API returns 404 when get_image raises exception""" + from superset.dashboards.api import DashboardRestApi + + # Create mock API instance + api = DashboardRestApi() + api.response_404 = MagicMock(return_value="404 Response") + + # Create mock cache payload that will raise exception + mock_cache_payload = MagicMock() + mock_cache_payload.get_image.side_effect = ( + ScreenshotImageNotAvailableException() + ) + + # Test the exception handling logic + try: + mock_cache_payload.get_image() + pytest.fail("Should have raised exception") + except ScreenshotImageNotAvailableException: + response = api.response_404() + assert response == "404 Response" + + def test_chart_screenshot_api_handles_exception(self): + """Test that chart screenshot API returns 404 when get_image raises exception""" + from superset.charts.api import ChartRestApi + + # Create mock API instance + api = ChartRestApi() + api.response_404 = MagicMock(return_value="404 Response") + + # Create mock cache payload that will raise exception + mock_cache_payload = MagicMock() + mock_cache_payload.get_image.side_effect = ( + ScreenshotImageNotAvailableException() + ) + + # Test the exception handling logic + try: + mock_cache_payload.get_image() + pytest.fail("Should have raised exception") + except ScreenshotImageNotAvailableException: + response = api.response_404() + assert response == "404 Response" + + def test_screenshot_cache_payload_exception_has_correct_status(self): + """Test that the ScreenshotImageNotAvailableException has status 404""" + exception = ScreenshotImageNotAvailableException() + assert exception.status == 404 + + def test_api_method_simulation_with_exception(self): + """Simulate the API method behavior with exception handling""" + + def simulate_dashboard_screenshot_method(cache_payload): + """Simulate the logic in dashboard screenshot methods""" + try: + image = cache_payload.get_image() + return {"status": "success", "image": image} + except ScreenshotImageNotAvailableException: + return {"status": "404", "message": "Not Found"} + + # Test with payload that has image + payload_with_image = ScreenshotCachePayload(image=b"test data") + result = simulate_dashboard_screenshot_method(payload_with_image) + assert result["status"] == "success" + assert result["image"] is not None + + # Test with payload that has no image (should raise exception) + payload_no_image = ScreenshotCachePayload() + result = simulate_dashboard_screenshot_method(payload_no_image) + assert result["status"] == "404" + assert result["message"] == "Not Found" + + def test_api_method_simulation_with_file_wrapper(self): + """Simulate the FileWrapper usage in API methods""" + from werkzeug.wsgi import FileWrapper + + def simulate_api_file_response(cache_payload): + """Simulate the FileWrapper logic in API methods""" + try: + image = cache_payload.get_image() + return FileWrapper(image) + except ScreenshotImageNotAvailableException: + return None + + # Test with valid image + payload_with_image = ScreenshotCachePayload(image=b"test data") + result = simulate_api_file_response(payload_with_image) + assert result is not None + assert isinstance(result, FileWrapper) + + # Test without image + payload_no_image = ScreenshotCachePayload() + result = simulate_api_file_response(payload_no_image) + assert result is None