You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by di...@apache.org on 2024/06/25 13:11:33 UTC
(superset) branch geido/feat/dashboard-screenshot-endpoint created (now 9c869e98b7)
This is an automated email from the ASF dual-hosted git repository.
diegopucci pushed a change to branch geido/feat/dashboard-screenshot-endpoint
in repository https://gitbox.apache.org/repos/asf/superset.git
at 9c869e98b7 Add endpoint
This branch includes the following new commits:
new 9c869e98b7 Add endpoint
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 01/01: Add endpoint
Posted by di...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
diegopucci pushed a commit to branch geido/feat/dashboard-screenshot-endpoint
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 9c869e98b7006cb9e572d2630b1149ae2be14d05
Author: Diego Pucci <di...@gmail.com>
AuthorDate: Tue Jun 25 15:05:57 2024 +0200
Add endpoint
---
superset/dashboards/api.py | 245 +++++++++++++-----
superset/dashboards/schemas.py | 65 ++++-
superset/tasks/thumbnails.py | 39 +++
tests/integration_tests/dashboards/api_tests.py | 325 +++++++++++-------------
4 files changed, 418 insertions(+), 256 deletions(-)
diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index 3fe557a684..598c92c23b 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,23 @@ from superset.dashboards.schemas import (
get_fav_star_ids_schema,
GetFavStarIdsSchema,
openapi_spec_methods_override,
- TabsPayloadSchema,
+ screenshot_query_schema,
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,
@@ -143,12 +154,13 @@ class DashboardRestApi(BaseSupersetModelRestApi):
"remove_favorite",
"get_charts",
"get_datasets",
- "get_tabs",
"get_embedded",
"set_embedded",
"delete_embedded",
"thumbnail",
"copy_dash",
+ "cache_dashboard_screenshot",
+ "screenshot",
}
resource_name = "dashboard"
allow_browser_login = True
@@ -239,7 +251,6 @@ class DashboardRestApi(BaseSupersetModelRestApi):
chart_entity_response_schema = ChartEntityResponseSchema()
dashboard_get_response_schema = DashboardGetResponseSchema()
dashboard_dataset_schema = DashboardDatasetSchema()
- tab_schema = TabsPayloadSchema()
embedded_response_schema = EmbeddedDashboardResponseSchema()
embedded_config_schema = EmbeddedDashboardConfigSchema()
@@ -269,10 +280,10 @@ class DashboardRestApi(BaseSupersetModelRestApi):
""" Override the name set for this collection of endpoints """
openapi_spec_component_schemas = (
ChartEntityResponseSchema,
+ DashboardCacheScreenshotResponseSchema,
DashboardCopySchema,
DashboardGetResponseSchema,
DashboardDatasetSchema,
- TabsPayloadSchema,
GetFavStarIdsSchema,
EmbeddedDashboardResponseSchema,
)
@@ -400,64 +411,6 @@ class DashboardRestApi(BaseSupersetModelRestApi):
except DashboardNotFoundError:
return self.response_404()
- @expose("/<id_or_slug>/tabs", methods=("GET",))
- @protect()
- @safe
- @statsd_metrics
- @event_logger.log_this_with_context(
- action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.get_tabs",
- log_to_statsd=False,
- )
- def get_tabs(self, id_or_slug: str) -> Response:
- """Get dashboard's tabs.
- ---
- get:
- summary: Get dashboard's tabs
- description: >-
- Returns a list of a dashboard's tabs and dashboard's nested tree structure for associated tabs.
- parameters:
- - in: path
- schema:
- type: string
- name: id_or_slug
- description: Either the id of the dashboard, or its slug
- responses:
- 200:
- description: Dashboard tabs
- content:
- application/json:
- schema:
- type: object
- properties:
- result:
- type: object
- items:
- $ref: '#/components/schemas/TabsPayloadSchema'
- 400:
- $ref: '#/components/responses/400'
- 401:
- $ref: '#/components/responses/401'
- 403:
- $ref: '#/components/responses/403'
- 404:
- $ref: '#/components/responses/404'
- """
- try:
- tabs = DashboardDAO.get_tabs_for_dashboard(id_or_slug)
- result = self.tab_schema.dump(tabs)
- return self.response(200, result=result)
-
- except (TypeError, ValueError) as err:
- return self.response_400(
- message=gettext(
- "Tab schema is invalid, caused by: %(error)s", error=str(err)
- )
- )
- except DashboardAccessDeniedError:
- return self.response_403()
- except DashboardNotFoundError:
- return self.response_404()
-
@expose("/<id_or_slug>/charts", methods=("GET",))
@protect()
@safe
@@ -939,6 +892,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..8d0b093454 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 = (
@@ -267,18 +279,6 @@ class DashboardDatasetSchema(Schema):
return serialized
-class TabSchema(Schema):
- # pylint: disable=W0108
- children = fields.List(fields.Nested(lambda: TabSchema()))
- value = fields.Str()
- title = fields.Str()
-
-
-class TabsPayloadSchema(Schema):
- all_tabs = fields.Dict(keys=fields.String(), values=fields.String())
- tab_tree = fields.List(fields.Nested(lambda: TabSchema))
-
-
class BaseDashboardSchema(Schema):
# pylint: disable=unused-argument
@post_load
@@ -385,6 +385,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 +445,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..7431b3fbd9 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")
@@ -918,185 +919,6 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
for key, value in expected_results[idx].items():
assert response_item[key] == value
- def test_get_dashboard_tabs(self):
- """
- Dashboard API: Test get dashboard tabs
- """
- position_data = {
- "GRID_ID": {"children": [], "id": "GRID_ID", "type": "GRID"},
- "ROOT_ID": {
- "children": ["TABS-tDGEcwZ82u"],
- "id": "ROOT_ID",
- "type": "ROOT",
- },
- "TAB-0TkqQRxzg7": {
- "children": [],
- "id": "TAB-0TkqQRxzg7",
- "meta": {"text": "P2 - T1"},
- "type": "TAB",
- },
- "TAB-1iG_yOlKA2": {
- "children": [],
- "id": "TAB-1iG_yOlKA2",
- "meta": {"text": "P1 - T1"},
- "type": "TAB",
- },
- "TAB-2dgADEurF": {
- "children": ["TABS-LsyXZWG2rk"],
- "id": "TAB-2dgADEurF",
- "meta": {"text": "P1 - T2"},
- "type": "TAB",
- },
- "TAB-BJIt5SdCx3": {
- "children": [],
- "id": "TAB-BJIt5SdCx3",
- "meta": {"text": "P1 - T2 - T1"},
- "type": "TAB",
- },
- "TAB-CjZlNL5Uz": {
- "children": ["TABS-Ji_K1ZBE0M"],
- "id": "TAB-CjZlNL5Uz",
- "meta": {"text": "Parent Tab 2"},
- "type": "TAB",
- },
- "TAB-Nct5fiHtn": {
- "children": [],
- "id": "TAB-Nct5fiHtn",
- "meta": {"text": "P1 - T2 - T3"},
- "type": "TAB",
- },
- "TAB-PumuDkWKq": {
- "children": [],
- "id": "TAB-PumuDkWKq",
- "meta": {"text": "P2 - T2"},
- "type": "TAB",
- },
- "TAB-hyTv5L7zz": {
- "children": [],
- "id": "TAB-hyTv5L7zz",
- "meta": {"text": "P1 - T2 - T2"},
- "type": "TAB",
- },
- "TAB-qL7fSzr3jl": {
- "children": ["TABS-N8ODUqp2sE"],
- "id": "TAB-qL7fSzr3jl",
- "meta": {"text": "Parent Tab 1"},
- "type": "TAB",
- },
- "TABS-Ji_K1ZBE0M": {
- "children": ["TAB-0TkqQRxzg7", "TAB-PumuDkWKq"],
- "id": "TABS-Ji_K1ZBE0M",
- "meta": {},
- "type": "TABS",
- },
- "TABS-LsyXZWG2rk": {
- "children": ["TAB-BJIt5SdCx3", "TAB-hyTv5L7zz", "TAB-Nct5fiHtn"],
- "id": "TABS-LsyXZWG2rk",
- "meta": {},
- "type": "TABS",
- },
- "TABS-N8ODUqp2sE": {
- "children": ["TAB-1iG_yOlKA2", "TAB-2dgADEurF"],
- "id": "TABS-N8ODUqp2sE",
- "meta": {},
- "type": "TABS",
- },
- "TABS-tDGEcwZ82u": {
- "children": ["TAB-qL7fSzr3jl", "TAB-CjZlNL5Uz"],
- "id": "TABS-tDGEcwZ82u",
- "meta": {},
- "type": "TABS",
- },
- }
- admin_id = self.get_user("admin").id
- dashboard = self.insert_dashboard(
- "title", "slug", [admin_id], position_json=json.dumps(position_data)
- )
- self.login(ADMIN_USERNAME)
- uri = f"api/v1/dashboard/{dashboard.id}/tabs"
- rv = self.get_assert_metric(uri, "get_tabs")
- response = json.loads(rv.data.decode("utf-8"))
- expected_response = {
- "result": {
- "all_tabs": {
- "TAB-0TkqQRxzg7": "P2 - T1",
- "TAB-1iG_yOlKA2": "P1 - T1",
- "TAB-2dgADEurF": "P1 - T2",
- "TAB-BJIt5SdCx3": "P1 - T2 - T1",
- "TAB-CjZlNL5Uz": "Parent Tab 2",
- "TAB-Nct5fiHtn": "P1 - T2 - T3",
- "TAB-PumuDkWKq": "P2 - T2",
- "TAB-hyTv5L7zz": "P1 - T2 - T2",
- "TAB-qL7fSzr3jl": "Parent Tab 1",
- },
- "tab_tree": [
- {
- "children": [
- {
- "children": [],
- "title": "P1 - T1",
- "value": "TAB-1iG_yOlKA2",
- },
- {
- "children": [
- {
- "children": [],
- "title": "P1 - T2 - T1",
- "value": "TAB-BJIt5SdCx3",
- },
- {
- "children": [],
- "title": "P1 - T2 - T2",
- "value": "TAB-hyTv5L7zz",
- },
- {
- "children": [],
- "title": "P1 - T2 - T3",
- "value": "TAB-Nct5fiHtn",
- },
- ],
- "title": "P1 - T2",
- "value": "TAB-2dgADEurF",
- },
- ],
- "title": "Parent Tab 1",
- "value": "TAB-qL7fSzr3jl",
- },
- {
- "children": [
- {
- "children": [],
- "title": "P2 - T1",
- "value": "TAB-0TkqQRxzg7",
- },
- {
- "children": [],
- "title": "P2 - T2",
- "value": "TAB-PumuDkWKq",
- },
- ],
- "title": "Parent Tab 2",
- "value": "TAB-CjZlNL5Uz",
- },
- ],
- }
- }
- self.assertEqual(rv.status_code, 200)
- self.assertEqual(response, expected_response)
- db.session.delete(dashboard)
- db.session.commit()
-
- @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices")
- def test_get_dashboard_tabs_not_found(self):
- """
- Dashboard API: Test get dashboard tabs not found
- """
- bad_id = self.get_nonexistent_numeric_id(Dashboard)
- self.login(ADMIN_USERNAME)
- uri = f"api/v1/dashboard/{bad_id}/tabs"
- rv = self.get_assert_metric(uri, "get_tabs")
- self.assertEqual(rv.status_code, 404)
-
def create_dashboard_import(self):
buf = BytesIO()
with ZipFile(buf, "w") as bundle:
@@ -2715,3 +2537,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)