This is an automated email from the ASF dual-hosted git repository.
diegopucci 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 d89648147f feat(dashboard): add API endpoints for generating and
downloading screenshots (#29187)
d89648147f is described below
commit d89648147f40750a1207bb11d73047a2887b54a7
Author: Edgar Ulloa <[email protected]>
AuthorDate: Mon Jul 8 02:52:23 2024 -0700
feat(dashboard): add API endpoints for generating and downloading
screenshots (#29187)
Co-authored-by: Diego Pucci <[email protected]>
---
superset/dashboards/api.py | 184 +++++++++++++++++++++++-
superset/dashboards/schemas.py | 53 ++++++-
superset/tasks/thumbnails.py | 39 +++++
tests/integration_tests/dashboards/api_tests.py | 146 +++++++++++++++++++
4 files changed, 419 insertions(+), 3 deletions(-)
diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index 823bfdfa8c..2967fd1abd 100644
--- a/superset/dashboards/api.py
+++ b/superset/dashboards/api.py
@@ -47,6 +47,7 @@ from superset.commands.dashboard.exceptions import (
)
from superset.commands.dashboard.export import ExportDashboardsCommand
from superset.commands.dashboard.importers.dispatcher import
ImportDashboardsCommand
+from superset.commands.dashboard.permalink.create import
CreateDashboardPermalinkCommand
from superset.commands.dashboard.update import UpdateDashboardCommand
from superset.commands.exceptions import TagForbiddenError
from superset.commands.importers.exceptions import NoValidFilesFoundError
@@ -63,7 +64,10 @@ from superset.dashboards.filters import (
DashboardTitleOrSlugFilter,
FilterRelatedRoles,
)
+from superset.dashboards.permalink.types import DashboardPermalinkState
from superset.dashboards.schemas import (
+ CacheScreenshotSchema,
+ DashboardCacheScreenshotResponseSchema,
DashboardCopySchema,
DashboardDatasetSchema,
DashboardGetResponseSchema,
@@ -76,16 +80,24 @@ from superset.dashboards.schemas import (
get_fav_star_ids_schema,
GetFavStarIdsSchema,
openapi_spec_methods_override,
+ screenshot_query_schema,
TabsPayloadSchema,
thumbnail_query_schema,
)
from superset.extensions import event_logger
from superset.models.dashboard import Dashboard
from superset.models.embedded_dashboard import EmbeddedDashboard
-from superset.tasks.thumbnails import cache_dashboard_thumbnail
+from superset.tasks.thumbnails import (
+ cache_dashboard_screenshot,
+ cache_dashboard_thumbnail,
+)
from superset.tasks.utils import get_current_user
from superset.utils import json
-from superset.utils.screenshots import DashboardScreenshot
+from superset.utils.pdf import build_pdf_from_screenshots
+from superset.utils.screenshots import (
+ DashboardScreenshot,
+ DEFAULT_DASHBOARD_WINDOW_SIZE,
+)
from superset.utils.urls import get_url_path
from superset.views.base_api import (
BaseSupersetModelRestApi,
@@ -124,6 +136,7 @@ def with_dashboard(
return functools.update_wrapper(wraps, f)
+# pylint: disable=too-many-public-methods
class DashboardRestApi(BaseSupersetModelRestApi):
datamodel = SQLAInterface(Dashboard)
@@ -149,6 +162,8 @@ class DashboardRestApi(BaseSupersetModelRestApi):
"delete_embedded",
"thumbnail",
"copy_dash",
+ "cache_dashboard_screenshot",
+ "screenshot",
}
resource_name = "dashboard"
allow_browser_login = True
@@ -269,6 +284,7 @@ class DashboardRestApi(BaseSupersetModelRestApi):
""" Override the name set for this collection of endpoints """
openapi_spec_component_schemas = (
ChartEntityResponseSchema,
+ DashboardCacheScreenshotResponseSchema,
DashboardCopySchema,
DashboardGetResponseSchema,
DashboardDatasetSchema,
@@ -939,6 +955,170 @@ class DashboardRestApi(BaseSupersetModelRestApi):
FileWrapper(screenshot), mimetype="image/png",
direct_passthrough=True
)
+ @expose("/<pk>/cache_dashboard_screenshot/", methods=("POST",))
+ @protect()
+ @rison(screenshot_query_schema)
+ @safe
+ @statsd_metrics
+ @event_logger.log_this_with_context(
+ action=lambda self, *args, **kwargs: f"{self.__class__.__name__}"
+ f".cache_dashboard_screenshot",
+ log_to_statsd=False,
+ )
+ def cache_dashboard_screenshot(self, pk: int, **kwargs: Any) ->
WerkzeugResponse:
+ """Compute and cache a screenshot.
+ ---
+ post:
+ summary: Compute and cache a screenshot
+ parameters:
+ - in: path
+ schema:
+ type: integer
+ name: pk
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/DashboardScreenshotPostSchema'
+ responses:
+ 202:
+ description: Dashboard async result
+ content:
+ application/json:
+ schema:
+ $ref:
"#/components/schemas/DashboardCacheScreenshotResponseSchema"
+ 400:
+ $ref: '#/components/responses/400'
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ 500:
+ $ref: '#/components/responses/500'
+ """
+ try:
+ payload = CacheScreenshotSchema().load(request.json)
+ except ValidationError as error:
+ return self.response_400(message=error.messages)
+
+ dashboard = cast(Dashboard, self.datamodel.get(pk, self._base_filters))
+ if not dashboard:
+ return self.response_404()
+
+ window_size = (
+ kwargs["rison"].get("window_size") or DEFAULT_DASHBOARD_WINDOW_SIZE
+ )
+ # Don't shrink the image if thumb_size is not specified
+ thumb_size = kwargs["rison"].get("thumb_size") or window_size
+
+ dashboard_state: DashboardPermalinkState = {
+ "dataMask": payload.get("dataMask", {}),
+ "activeTabs": payload.get("activeTabs", []),
+ "anchor": payload.get("anchor", ""),
+ "urlParams": payload.get("urlParams", []),
+ }
+
+ permalink_key = CreateDashboardPermalinkCommand(
+ dashboard_id=str(dashboard.id),
+ state=dashboard_state,
+ ).run()
+
+ dashboard_url = get_url_path("Superset.dashboard_permalink",
key=permalink_key)
+ screenshot_obj = DashboardScreenshot(dashboard_url, dashboard.digest)
+ cache_key = screenshot_obj.cache_key(window_size, thumb_size)
+ image_url = get_url_path(
+ "DashboardRestApi.screenshot", pk=dashboard.id, digest=cache_key
+ )
+
+ def trigger_celery() -> WerkzeugResponse:
+ logger.info("Triggering screenshot ASYNC")
+ cache_dashboard_screenshot.delay(
+ current_user=get_current_user(),
+ dashboard_id=dashboard.id,
+ dashboard_url=dashboard_url,
+ force=True,
+ thumb_size=thumb_size,
+ window_size=window_size,
+ )
+ return self.response(
+ 202,
+ cache_key=cache_key,
+ dashboard_url=dashboard_url,
+ image_url=image_url,
+ )
+
+ return trigger_celery()
+
+ @expose("/<pk>/screenshot/<digest>/", methods=("GET",))
+ @protect()
+ @safe
+ @statsd_metrics
+ @event_logger.log_this_with_context(
+ action=lambda self, *args, **kwargs:
f"{self.__class__.__name__}.screenshot",
+ log_to_statsd=False,
+ )
+ def screenshot(self, pk: int, digest: str) -> WerkzeugResponse:
+ """Get a computed dashboard screenshot from cache.
+ ---
+ get:
+ summary: Get a computed screenshot from cache
+ parameters:
+ - in: path
+ schema:
+ type: integer
+ name: pk
+ - in: path
+ schema:
+ type: string
+ name: digest
+ responses:
+ 200:
+ description: Dashboard thumbnail image
+ content:
+ image/*:
+ schema:
+ type: string
+ format: binary
+ 400:
+ $ref: '#/components/responses/400'
+ 401:
+ $ref: '#/components/responses/401'
+ 404:
+ $ref: '#/components/responses/404'
+ 500:
+ $ref: '#/components/responses/500'
+ """
+ dashboard = self.datamodel.get(pk, self._base_filters)
+
+ # Making sure the dashboard still exists
+ if not dashboard:
+ return self.response_404()
+
+ download_format = request.args.get("download_format", "png")
+
+ # fetch the dashboard screenshot using the current user and cache if
set
+
+ if img := DashboardScreenshot.get_from_cache_key(thumbnail_cache,
digest):
+ if download_format == "pdf":
+ pdf_img = img.getvalue()
+ # Convert the screenshot to PDF
+ pdf_data = build_pdf_from_screenshots([pdf_img])
+
+ return Response(
+ pdf_data,
+ mimetype="application/pdf",
+ headers={"Content-Disposition": "inline;
filename=dashboard.pdf"},
+ direct_passthrough=True,
+ )
+ if download_format == "png":
+ return Response(
+ FileWrapper(img),
+ mimetype="image/png",
+ direct_passthrough=True,
+ )
+
+ return self.response_404()
+
@expose("/favorite_status/", methods=("GET",))
@protect()
@safe
diff --git a/superset/dashboards/schemas.py b/superset/dashboards/schemas.py
index 3d4d131b9e..1f78a22358 100644
--- a/superset/dashboards/schemas.py
+++ b/superset/dashboards/schemas.py
@@ -31,7 +31,19 @@ thumbnail_query_schema = {
"type": "object",
"properties": {"force": {"type": "boolean"}},
}
-
+width_height_schema = {
+ "type": "array",
+ "items": {"type": "integer"},
+}
+screenshot_query_schema = {
+ "type": "object",
+ "properties": {
+ "force": {"type": "boolean"},
+ "permalink": {"type": "string"},
+ "window_size": width_height_schema,
+ "thumb_size": width_height_schema,
+ },
+}
dashboard_title_description = "A title for the dashboard."
slug_description = "Unique identifying part for the web address of the
dashboard."
owners_description = (
@@ -385,6 +397,26 @@ class DashboardPutSchema(BaseDashboardSchema):
)
+class DashboardScreenshotPostSchema(Schema):
+ dataMask = fields.Dict(
+ keys=fields.Str(),
+ values=fields.Raw(),
+ metadata={"description": "An object representing the data mask."},
+ )
+ activeTabs = fields.List(
+ fields.Str(), metadata={"description": "A list representing active
tabs."}
+ )
+ anchor = fields.String(
+ metadata={"description": "A string representing the anchor."}
+ )
+ urlParams = fields.List(
+ fields.Tuple(
+ (fields.Str(), fields.Str()),
+ ),
+ metadata={"description": "A list of tuples, each containing two
strings."},
+ )
+
+
class ChartFavStarResponseResult(Schema):
id = fields.Integer(metadata={"description": "The Chart id"})
value = fields.Boolean(metadata={"description": "The FaveStar value"})
@@ -425,3 +457,22 @@ class EmbeddedDashboardResponseSchema(Schema):
dashboard_id = fields.String()
changed_on = fields.DateTime()
changed_by = fields.Nested(UserSchema)
+
+
+class DashboardCacheScreenshotResponseSchema(Schema):
+ cache_key = fields.String(metadata={"description": "The cache key"})
+ dashboard_url = fields.String(
+ metadata={"description": "The url to render the dashboard"}
+ )
+ image_url = fields.String(
+ metadata={"description": "The url to fetch the screenshot"}
+ )
+
+
+class CacheScreenshotSchema(Schema):
+ dataMask = fields.Dict(keys=fields.Str(), values=fields.Raw(),
required=False)
+ activeTabs = fields.List(fields.Str(), required=False)
+ anchor = fields.Str(required=False)
+ urlParams = fields.List(
+ fields.List(fields.Str(), validate=lambda x: len(x) == 2),
required=False
+ )
diff --git a/superset/tasks/thumbnails.py b/superset/tasks/thumbnails.py
index 03b3999dce..da42b1b146 100644
--- a/superset/tasks/thumbnails.py
+++ b/superset/tasks/thumbnails.py
@@ -77,6 +77,7 @@ def cache_dashboard_thumbnail(
dashboard_id: int,
force: bool = False,
thumb_size: Optional[WindowSize] = None,
+ window_size: Optional[WindowSize] = None,
) -> None:
# pylint: disable=import-outside-toplevel
from superset.models.dashboard import Dashboard
@@ -100,5 +101,43 @@ def cache_dashboard_thumbnail(
user=user,
cache=thumbnail_cache,
force=force,
+ window_size=window_size,
+ thumb_size=thumb_size,
+ )
+
+
+# pylint: disable=too-many-arguments
+@celery_app.task(name="cache_dashboard_screenshot", soft_time_limit=60)
+def cache_dashboard_screenshot(
+ current_user: Optional[str],
+ dashboard_id: int,
+ dashboard_url: str,
+ force: bool = False,
+ thumb_size: Optional[WindowSize] = None,
+ window_size: Optional[WindowSize] = None,
+) -> None:
+ # pylint: disable=import-outside-toplevel
+ from superset.models.dashboard import Dashboard
+
+ if not thumbnail_cache:
+ logging.warning("No cache set, refusing to compute")
+ return
+
+ dashboard = Dashboard.get(dashboard_id)
+
+ logger.info("Caching dashboard: %s", dashboard_url)
+ _, username = get_executor(
+ executor_types=current_app.config["THUMBNAIL_EXECUTE_AS"],
+ model=dashboard,
+ current_user=current_user,
+ )
+ user = security_manager.find_user(username)
+ with override_user(user):
+ screenshot = DashboardScreenshot(dashboard_url, dashboard.digest)
+ screenshot.compute_and_cache(
+ user=user,
+ cache=thumbnail_cache,
+ force=force,
+ window_size=window_size,
thumb_size=thumb_size,
)
diff --git a/tests/integration_tests/dashboards/api_tests.py
b/tests/integration_tests/dashboards/api_tests.py
index b4213039c4..99a784e95f 100644
--- a/tests/integration_tests/dashboards/api_tests.py
+++ b/tests/integration_tests/dashboards/api_tests.py
@@ -532,6 +532,7 @@ class TestDashboardApi(ApiOwnersTestCaseMixin,
InsertChartMixin, SupersetTestCas
"can_get_embedded",
"can_delete_embedded",
"can_set_embedded",
+ "can_cache_dashboard_screenshot",
}
@pytest.mark.usefixtures("load_world_bank_dashboard_with_slices")
@@ -2715,3 +2716,148 @@ class TestDashboardApi(ApiOwnersTestCaseMixin,
InsertChartMixin, SupersetTestCas
security_manager.add_permission_role(gamma_role, write_tags_perm)
security_manager.add_permission_role(gamma_role, tag_dashboards_perm)
+
+ def _cache_screenshot(self, dashboard_id, payload=None):
+ if payload is None:
+ payload = {"dataMask": {}, "activeTabs": [], "anchor": "",
"urlParams": []}
+ uri = f"/api/v1/dashboard/{dashboard_id}/cache_dashboard_screenshot/"
+ return self.client.post(uri, json=payload)
+
+ def _get_screenshot(self, dashboard_id, cache_key, download_format):
+ uri =
f"/api/v1/dashboard/{dashboard_id}/screenshot/{cache_key}/?download_format={download_format}"
+ return self.client.get(uri)
+
+ @pytest.mark.usefixtures("create_dashboard_with_tag")
+ def test_cache_dashboard_screenshot_success(self):
+ self.login(ADMIN_USERNAME)
+ dashboard = (
+ db.session.query(Dashboard)
+ .filter(Dashboard.dashboard_title == "dash with tag")
+ .first()
+ )
+ response = self._cache_screenshot(dashboard.id)
+ self.assertEqual(response.status_code, 202)
+
+ @pytest.mark.usefixtures("create_dashboard_with_tag")
+ def test_cache_dashboard_screenshot_dashboard_validation(self):
+ self.login(ADMIN_USERNAME)
+ dashboard = (
+ db.session.query(Dashboard)
+ .filter(Dashboard.dashboard_title == "dash with tag")
+ .first()
+ )
+ invalid_payload = {
+ "dataMask": ["should be a dict"],
+ "activeTabs": "should be a list",
+ "anchor": 1,
+ "urlParams": "should be a list",
+ }
+ response = self._cache_screenshot(dashboard.id, invalid_payload)
+ self.assertEqual(response.status_code, 400)
+
+ def test_cache_dashboard_screenshot_dashboard_not_found(self):
+ self.login(ADMIN_USERNAME)
+ non_existent_id = 999
+ response = self._cache_screenshot(non_existent_id)
+ self.assertEqual(response.status_code, 404)
+
+ @pytest.mark.usefixtures("create_dashboard_with_tag")
+ @patch("superset.dashboards.api.cache_dashboard_screenshot")
+ @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
+ def test_screenshot_success_png(self, mock_get_cache, mock_cache_task):
+ """
+ Validate screenshot returns png
+ """
+ self.login(ADMIN_USERNAME)
+ mock_cache_task.return_value = None
+ mock_get_cache.return_value = BytesIO(b"fake image data")
+
+ dashboard = (
+ db.session.query(Dashboard)
+ .filter(Dashboard.dashboard_title == "dash with tag")
+ .first()
+ )
+ cache_resp = self._cache_screenshot(dashboard.id)
+ self.assertEqual(cache_resp.status_code, 202)
+ cache_key = json.loads(cache_resp.data.decode("utf-8"))["cache_key"]
+
+ response = self._get_screenshot(dashboard.id, cache_key, "png")
+ self.assertEqual(response.status_code, 200)
+ self.assertEqual(response.mimetype, "image/png")
+ self.assertEqual(response.data, b"fake image data")
+
+ @pytest.mark.usefixtures("create_dashboard_with_tag")
+ @patch("superset.dashboards.api.cache_dashboard_screenshot")
+ @patch("superset.dashboards.api.build_pdf_from_screenshots")
+ @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
+ def test_screenshot_success_pdf(
+ self, mock_get_from_cache, mock_build_pdf, mock_cache_task
+ ):
+ """
+ Validate screenshot can return pdf.
+ """
+ self.login(ADMIN_USERNAME)
+ mock_cache_task.return_value = None
+ mock_get_from_cache.return_value = BytesIO(b"fake image data")
+ mock_build_pdf.return_value = b"fake pdf data"
+
+ dashboard = (
+ db.session.query(Dashboard)
+ .filter(Dashboard.dashboard_title == "dash with tag")
+ .first()
+ )
+ cache_resp = self._cache_screenshot(dashboard.id)
+ self.assertEqual(cache_resp.status_code, 202)
+ cache_key = json.loads(cache_resp.data.decode("utf-8"))["cache_key"]
+
+ response = self._get_screenshot(dashboard.id, cache_key, "pdf")
+ self.assertEqual(response.status_code, 200)
+ self.assertEqual(response.mimetype, "application/pdf")
+ self.assertEqual(response.data, b"fake pdf data")
+
+ @pytest.mark.usefixtures("create_dashboard_with_tag")
+ @patch("superset.dashboards.api.cache_dashboard_screenshot")
+ @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
+ def test_screenshot_not_in_cache(self, mock_get_cache, mock_cache_task):
+ self.login(ADMIN_USERNAME)
+ mock_cache_task.return_value = None
+ mock_get_cache.return_value = None
+
+ dashboard = (
+ db.session.query(Dashboard)
+ .filter(Dashboard.dashboard_title == "dash with tag")
+ .first()
+ )
+ cache_resp = self._cache_screenshot(dashboard.id)
+ self.assertEqual(cache_resp.status_code, 202)
+ cache_key = json.loads(cache_resp.data.decode("utf-8"))["cache_key"]
+
+ response = self._get_screenshot(dashboard.id, cache_key, "pdf")
+ self.assertEqual(response.status_code, 404)
+
+ def test_screenshot_dashboard_not_found(self):
+ self.login(ADMIN_USERNAME)
+ non_existent_id = 999
+ response = self._get_screenshot(non_existent_id, "some_cache_key",
"png")
+ self.assertEqual(response.status_code, 404)
+
+ @pytest.mark.usefixtures("create_dashboard_with_tag")
+ @patch("superset.dashboards.api.cache_dashboard_screenshot")
+ @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
+ def test_screenshot_invalid_download_format(self, mock_get_cache,
mock_cache_task):
+ self.login(ADMIN_USERNAME)
+ mock_cache_task.return_value = None
+ mock_get_cache.return_value = BytesIO(b"fake png data")
+
+ dashboard = (
+ db.session.query(Dashboard)
+ .filter(Dashboard.dashboard_title == "dash with tag")
+ .first()
+ )
+
+ cache_resp = self._cache_screenshot(dashboard.id)
+ self.assertEqual(cache_resp.status_code, 202)
+ cache_key = json.loads(cache_resp.data.decode("utf-8"))["cache_key"]
+
+ response = self._get_screenshot(dashboard.id, cache_key, "invalid")
+ self.assertEqual(response.status_code, 404)
