diff --git a/superset/config.py b/superset/config.py index 0ba239a509c..6c8a4073a14 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1448,6 +1448,7 @@ class CeleryConfig: # pylint: disable=too-few-public-methods "superset.tasks.thumbnails", "superset.tasks.cache", "superset.tasks.slack", + "superset.tasks.export_dashboard_excel", ) result_backend = "db+sqlite:///celery_results.sqlite" worker_prefetch_multiplier = 1 diff --git a/superset/tasks/export_dashboard_excel.py b/superset/tasks/export_dashboard_excel.py new file mode 100644 index 00000000000..6022b84b28f --- /dev/null +++ b/superset/tasks/export_dashboard_excel.py @@ -0,0 +1,247 @@ +# 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. +""" +Celery task that exports every chart on a dashboard to a single multi-sheet +``.xlsx`` file, uploads it to S3, and emails the requesting user a pre-signed +download link. + +The task re-runs each chart's saved query context under the requesting user, +applies the live dashboard filter state, and streams the results row-by-row into +a constant-memory workbook so large dashboards never load all data at once. +""" + +from __future__ import annotations + +import logging +import os +import tempfile +from datetime import datetime, timedelta +from typing import Any + +from celery.exceptions import SoftTimeLimitExceeded +from flask import current_app, g + +from superset import db, security_manager +from superset.charts.data.dashboard_filter_context import ( + apply_extra_form_data_to_query_context_json, + get_dashboard_filter_context, +) +from superset.charts.schemas import ChartDataQueryContextSchema +from superset.commands.chart.data.get_data_command import ChartDataCommand +from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType +from superset.dashboards.excel_export import email +from superset.dashboards.excel_export.layout import get_charts_in_layout_order +from superset.extensions import celery_app +from superset.utils import json, s3 +from superset.utils.core import override_user +from superset.utils.excel_streaming import StreamingXlsxWriter + +logger = logging.getLogger(__name__) + + +def _chart_label(chart: Any) -> str: + """Human-readable label for a chart in the skipped-charts list.""" + return f"{chart.id} - {chart.slice_name or ''}".strip() + + +def _record_to_row(record: dict[str, Any], colnames: list[str]) -> list[Any]: + return [record.get(col) for col in colnames] + + +def _write_chart_sheets( + writer: StreamingXlsxWriter, + chart: Any, + dashboard_id: int, + active_data_mask: dict[str, Any], +) -> None: + """ + Run a single chart's query and stream its result(s) into the workbook. + + Charts may yield more than one query (e.g. mixed-series charts); each becomes + its own sheet. Raises if the chart cannot be exported, so the caller can skip + it and note it in the email. + """ + json_body = json.loads(chart.query_context) + # Override any stale saved values: we always want full JSON results. + json_body["result_format"] = ChartDataResultFormat.JSON + json_body["result_type"] = ChartDataResultType.FULL + json_body.pop("force", None) + + filter_context = get_dashboard_filter_context( + dashboard_id=dashboard_id, + chart_id=chart.id, + active_data_mask=active_data_mask, + ) + apply_extra_form_data_to_query_context_json( + json_body, filter_context.extra_form_data + ) + + # Jinja macros resolve form data from g.form_data; expose the saved context. + g.form_data = json_body + + query_context = ChartDataQueryContextSchema().load(json_body) + command = ChartDataCommand(query_context) + command.validate() + result = command.run() + + for index, query in enumerate(result["queries"]): + colnames = query.get("colnames") or [] + data = query.get("data") or [] + if index == 0: + name = f"{chart.id} - {chart.slice_name or ''}" + else: + name = f"{chart.id}.{index} - {chart.slice_name or ''}" + writer.add_sheet( + name, + colnames, + (_record_to_row(record, colnames) for record in data), + ) + + +def _build_workbook( + path: str, + dashboard: Any, + active_data_mask: dict[str, Any], + job_id: str, +) -> list[str]: + """Build the workbook on disk; return the list of skipped chart labels.""" + skipped: list[str] = [] + writer = StreamingXlsxWriter(path) + try: + for chart in get_charts_in_layout_order(dashboard): + if not chart.query_context: + skipped.append(_chart_label(chart)) + continue + try: + _write_chart_sheets(writer, chart, dashboard.id, active_data_mask) + except Exception: # pylint: disable=broad-except + logger.exception( + "Skipping chart %s in dashboard export %s", chart.id, job_id + ) + skipped.append(_chart_label(chart)) + + if writer.sheet_count == 0: + writer.add_summary_sheet( + "Export Summary", + ["No chart data could be exported.", *skipped], + ) + finally: + writer.close() + return skipped + + +def _send_failure_email( + user: Any, dashboard_title: str, requested_at: datetime +) -> None: + if not (user and getattr(user, "email", None)): + return + try: + email.send_export_email( + user.email, + email.build_subject(dashboard_title, success=False), + email.build_failure_email(dashboard_title, requested_at), + ) + except Exception: # pylint: disable=broad-except + logger.exception("Failed to send export failure email") + + +@celery_app.task( + name="export_dashboard_excel", + bind=True, + soft_time_limit=600, + time_limit=660, + max_retries=0, +) +def export_dashboard_excel( + self: Any, # pylint: disable=unused-argument + dashboard_id: int, + user_id: int, + active_data_mask: dict[str, Any], + job_id: str, +) -> None: + """ + Export a dashboard's chart data to an ``.xlsx`` and email a download link. + + :param dashboard_id: The dashboard to export + :param user_id: The requesting user (the task runs with their permissions) + :param active_data_mask: Live dashboard filter state keyed by native filter id + :param job_id: Correlation id, also the Celery task id and S3 object name + """ + # pylint: disable=import-outside-toplevel + from superset.models.dashboard import Dashboard + + requested_at = datetime.utcnow() + user = security_manager.get_user_by_id(user_id) + dashboard_title = "" + tmp_path: str | None = None + + try: + with override_user(user, force=False): + dashboard = ( + db.session.query(Dashboard).filter_by(id=dashboard_id).one_or_none() + ) + if dashboard is None: + raise ValueError(f"Dashboard {dashboard_id} not found") + dashboard_title = dashboard.dashboard_title or f"Dashboard {dashboard_id}" + + file_descriptor, tmp_path = tempfile.mkstemp( + suffix=".xlsx", prefix=f"dash-export-{job_id}-" + ) + os.close(file_descriptor) + + skipped = _build_workbook(tmp_path, dashboard, active_data_mask, job_id) + + bucket = current_app.config["EXCEL_EXPORT_S3_BUCKET"] + key = ( + f"{current_app.config['EXCEL_EXPORT_S3_KEY_PREFIX']}" + f"{dashboard_id}/{job_id}.xlsx" + ) + ttl = current_app.config["EXCEL_EXPORT_LINK_TTL_SECONDS"] + + s3.upload_file_to_s3(tmp_path, bucket, key) + download_url = s3.generate_presigned_url(bucket, key, ttl) + expires_at = datetime.utcnow() + timedelta(seconds=ttl) + + if user and getattr(user, "email", None): + try: + email.send_export_email( + user.email, + email.build_subject(dashboard_title, success=True), + email.build_success_email( + dashboard_title=dashboard_title, + download_url=download_url, + requested_at=requested_at, + expires_at=expires_at, + ttl_seconds=ttl, + skipped_charts=skipped, + ), + ) + except Exception: # pylint: disable=broad-except + # The file is already in S3; a send failure should not trigger + # a misleading failure email. + logger.exception("Failed to send export success email") + except SoftTimeLimitExceeded: + logger.warning("Dashboard excel export %s timed out", job_id) + _send_failure_email(user, dashboard_title, requested_at) + raise + except Exception: + logger.exception("Dashboard excel export %s failed", job_id) + _send_failure_email(user, dashboard_title, requested_at) + raise + finally: + if tmp_path and os.path.exists(tmp_path): + os.remove(tmp_path) diff --git a/tests/unit_tests/tasks/test_export_dashboard_excel.py b/tests/unit_tests/tasks/test_export_dashboard_excel.py new file mode 100644 index 00000000000..0ef3aa97ade --- /dev/null +++ b/tests/unit_tests/tasks/test_export_dashboard_excel.py @@ -0,0 +1,212 @@ +# 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. +from __future__ import annotations + +import glob +import os +import tempfile +from collections.abc import Iterator +from contextlib import ExitStack +from typing import Any +from unittest import mock + +import pytest +from celery.exceptions import SoftTimeLimitExceeded + +from superset.utils import json + +MODULE = "superset.tasks.export_dashboard_excel" + + +def _chart(chart_id: int, name: str, has_context: bool = True) -> mock.MagicMock: + chart = mock.MagicMock() + chart.id = chart_id + chart.slice_name = name + chart.query_context = json.dumps({"queries": [{}]}) if has_context else None + return chart + + +@pytest.fixture +def mocks() -> Iterator[dict[str, Any]]: + """Patch every external dependency of the task; keep the real xlsx writer.""" + with ExitStack() as stack: + # Use explicit MagicMock instances: patch() auto-creates async-flavored + # mocks for these targets (their real objects expose async members), which + # would make calls like security_manager.get_user_by_id() return coroutines. + patched = { + name: stack.enter_context( + mock.patch(f"{MODULE}.{name}", new=mock.MagicMock()) + ) + for name in ( + "security_manager", + "db", + "get_charts_in_layout_order", + "get_dashboard_filter_context", + "ChartDataQueryContextSchema", + "ChartDataCommand", + "s3", + "email", + ) + } + user = mock.MagicMock() + user.email = "user@example.com" + patched["security_manager"].get_user_by_id.return_value = user + + dashboard = mock.MagicMock() + dashboard.id = 1 + dashboard.dashboard_title = "Sales" + patched[ + "db" + ].session.query.return_value.filter_by.return_value.one_or_none.return_value = ( # noqa: E501 + dashboard + ) + + patched["get_dashboard_filter_context"].return_value.extra_form_data = {} + patched["s3"].generate_presigned_url.return_value = "https://signed/file.xlsx" + + patched["user"] = user + patched["dashboard"] = dashboard + yield patched + + +def _run(job_id: str = "job-1") -> None: + from superset.tasks.export_dashboard_excel import export_dashboard_excel + + export_dashboard_excel( + dashboard_id=1, user_id=2, active_data_mask={}, job_id=job_id + ) + + +def _no_temp_files_left(job_id: str) -> bool: + pattern = os.path.join(tempfile.gettempdir(), f"dash-export-{job_id}-*") + return glob.glob(pattern) == [] + + +def _read_sheets(path: str) -> dict[str, list[list[object]]]: + openpyxl = pytest.importorskip("openpyxl") + workbook = openpyxl.load_workbook(path, read_only=True) + sheets = { + ws.title: [list(r) for r in ws.iter_rows(values_only=True)] + for ws in workbook.worksheets + } + workbook.close() + return sheets + + +def test_happy_path_uploads_and_emails(mocks: dict[str, Any]) -> None: + mocks["get_charts_in_layout_order"].return_value = [ + _chart(10, "First"), + _chart(20, "Second"), + ] + mocks["ChartDataCommand"].return_value.run.side_effect = [ + {"queries": [{"colnames": ["a", "b"], "data": [{"a": 1, "b": 2}]}]}, + {"queries": [{"colnames": ["c"], "data": [{"c": "x"}]}]}, + ] + + # Capture the workbook before the task deletes it. + uploaded: dict[str, Any] = {} + + def _capture(path: str, bucket: str, key: str) -> None: + uploaded["sheets"] = _read_sheets(path) + + mocks["s3"].upload_file_to_s3.side_effect = _capture + + _run() + + mocks["s3"].upload_file_to_s3.assert_called_once() + assert list(uploaded["sheets"].keys()) == ["10 - First", "20 - Second"] + mocks["email"].send_export_email.assert_called_once() + mocks["email"].build_success_email.assert_called_once() + assert _no_temp_files_left("job-1") + + +def test_chart_without_query_context_is_skipped(mocks: dict[str, Any]) -> None: + mocks["get_charts_in_layout_order"].return_value = [ + _chart(10, "Good"), + _chart(20, "NoContext", has_context=False), + ] + mocks["ChartDataCommand"].return_value.run.return_value = { + "queries": [{"colnames": ["a"], "data": [{"a": 1}]}] + } + + _run() + + _, kwargs = mocks["email"].build_success_email.call_args + assert kwargs["skipped_charts"] == ["20 - NoContext"] + + +def test_chart_query_error_is_skipped_export_continues( + mocks: dict[str, Any], +) -> None: + mocks["get_charts_in_layout_order"].return_value = [ + _chart(10, "Boom"), + _chart(20, "Ok"), + ] + mocks["ChartDataCommand"].return_value.run.side_effect = [ + RuntimeError("query failed"), + {"queries": [{"colnames": ["a"], "data": [{"a": 1}]}]}, + ] + + _run() + + mocks["s3"].upload_file_to_s3.assert_called_once() + _, kwargs = mocks["email"].build_success_email.call_args + assert kwargs["skipped_charts"] == ["10 - Boom"] + + +def test_all_charts_skipped_writes_summary(mocks: dict[str, Any]) -> None: + mocks["get_charts_in_layout_order"].return_value = [ + _chart(10, "NoContext", has_context=False), + ] + uploaded: dict[str, Any] = {} + + def _capture(path: str, bucket: str, key: str) -> None: + uploaded["sheets"] = _read_sheets(path) + + mocks["s3"].upload_file_to_s3.side_effect = _capture + + _run() + + assert "Export Summary" in uploaded["sheets"] + mocks["email"].build_success_email.assert_called_once() + + +def test_upload_failure_sends_failure_email_and_cleans_up( + mocks: dict[str, Any], +) -> None: + mocks["get_charts_in_layout_order"].return_value = [_chart(10, "Good")] + mocks["ChartDataCommand"].return_value.run.return_value = { + "queries": [{"colnames": ["a"], "data": [{"a": 1}]}] + } + mocks["s3"].upload_file_to_s3.side_effect = RuntimeError("s3 down") + + with pytest.raises(RuntimeError): + _run("job-fail") + + mocks["email"].build_failure_email.assert_called_once() + mocks["email"].send_export_email.assert_called_once() + assert _no_temp_files_left("job-fail") + + +def test_soft_time_limit_sends_failure_email(mocks: dict[str, Any]) -> None: + mocks["get_charts_in_layout_order"].side_effect = SoftTimeLimitExceeded() + + with pytest.raises(SoftTimeLimitExceeded): + _run("job-timeout") + + mocks["email"].build_failure_email.assert_called_once() + assert _no_temp_files_left("job-timeout")