fix: make get_image() always return BytesIO

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,
         )

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,
         )

superset/exceptions.py

@@ -486,3 +486,7 @@ def __init__(self) -> None:
             level=ErrorLevel.ERROR,
         )
         super().__init__(error)
+
+
+class ScreenshotImageNotAvailableException(SupersetException):
+    status = 404

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

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

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