This is an automated email from the ASF dual-hosted git repository. EnxDev pushed a commit to branch enxdev/feat/dashboard-excel-export-sync-fallback in repository https://gitbox.apache.org/repos/asf/superset.git
commit f64b45f2c4c245f8cdce559f5ca7ea1543c0a6ed Author: Enzo Martellucci <[email protected]> AuthorDate: Wed Sep 9 12:06:05 2026 +0200 feat(dashboard): stream the Excel export inline when no storage is configured "Export Data to Excel" was only usable where EXCEL_EXPORT_S3_BUCKET was set: everywhere else the menu item rendered, the endpoint answered 501, and the click produced nothing but an error toast. The endpoint now picks its execution path from that same config rather than refusing. With a bucket it queues the export and answers 202 with a job id, unchanged. Without one it builds the workbook during the request and returns it as the response, so the browser downloads the file directly -- no worker, bucket or email required. An export served this way has to finish inside one request, so its size is settled before any query runs: the new EXCEL_EXPORT_SYNC_MAX_ROWS (100,000) caps the combined row_limit of every query the export would run, and an export over that total -- or one where any query has no finite limit -- is refused with a 400 naming the bucket as the fix, rather than being started and left to hit the request timeout. Both paths build the same workbook from the same code: the builder moves out of the Celery task into superset.dashboards.excel_export.workbook, leaving the task to own only queueing, upload, email and its own error handling. Both acquire the same per-user+dashboard lock; the inline path releases it, and deletes its temp file, however the request ends. Guest, embedded, anonymous and no-email sessions stay blocked on both paths -- guest support is handled separately in #43805 -- and are covered by a regression test using a guest token that does grant access to the dashboard. The storage check is isolated in a single helper so it can be repointed at EXPORT_STORAGE when #43805 lands. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> --- UPDATING.md | 36 ++- .../using-superset/exporting-dashboard-data.mdx | 60 ++-- .../DownloadMenuItems/DownloadMenuItems.test.tsx | 106 +++++- .../components/menu/DownloadMenuItems/index.tsx | 102 ++++-- superset/config.py | 9 + superset/dashboards/api.py | 128 +++++++- superset/dashboards/excel_export/storage.py | 37 +++ superset/dashboards/excel_export/sync_budget.py | 98 ++++++ .../excel_export/workbook.py} | 211 +++--------- superset/tasks/export_dashboard_excel.py | 354 +-------------------- tests/integration_tests/dashboards/api_tests.py | 257 ++++++++++++++- .../dashboards/test_excel_export_storage.py | 38 +++ .../dashboards/test_excel_export_sync_budget.py | 152 +++++++++ .../tasks/test_export_dashboard_excel.py | 77 +++-- 14 files changed, 1020 insertions(+), 645 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index 786f4add82b..e944b087f43 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -509,19 +509,34 @@ Note that a retried query returns partial data with no truncation indicator (e.g. a filter dropdown may list only a subset of values on tables above the row cap). -### Dashboard "Export Data to Excel" requires a Celery worker and S3 bucket +### Dashboard "Export Data to Excel" scales with a Celery worker and S3 bucket A new dashboard action exports every chart's data to a single multi-sheet -`.xlsx` asynchronously. It is disabled by default and turns on only when -`EXCEL_EXPORT_S3_BUCKET` is set (the endpoint returns `501` otherwise). It also -requires a running Celery worker and a configured SMTP transport, since the task -emails the requesting user a pre-signed download link. New config keys: -`EXCEL_EXPORT_S3_BUCKET`, `EXCEL_EXPORT_S3_KEY_PREFIX`, +`.xlsx`. Setting `EXCEL_EXPORT_S3_BUCKET` selects how it runs: with a bucket the +export is queued to a Celery worker, uploaded, and emailed to the requesting +user as a pre-signed download link (so it also needs a running worker and a +configured SMTP transport); without one the workbook is built during the request +and returned as the response for the browser to download, needing no +configuration at all. + +Because it has to finish inside a single request, the direct-download path is +bounded by `EXCEL_EXPORT_SYNC_MAX_ROWS` (default `100_000`): the export sums the +`row_limit` of every query it would run and refuses, before running any of them, +when the total is higher or when any query has no finite limit. The refusal is a +`400` naming `EXCEL_EXPORT_S3_BUCKET` as the fix. + +API clients should note that `POST /api/v1/dashboard/<id>/export_xlsx/` now +answers either `202` with a job id (queued) or `200` with the `.xlsx` itself +(direct download), depending on this configuration. It no longer returns `501`. + +New config keys: `EXCEL_EXPORT_S3_BUCKET`, `EXCEL_EXPORT_S3_KEY_PREFIX`, `EXCEL_EXPORT_LINK_TTL_SECONDS`, `EXCEL_EXPORT_S3_CLIENT_KWARGS`, -`EXCEL_EXPORT_TABLE_VIZ_TYPES`, and `EXCEL_EXPORT_QUERY_CONTEXT_BUILDER`. +`EXCEL_EXPORT_SYNC_MAX_ROWS`, `EXCEL_EXPORT_TABLE_VIZ_TYPES`, and +`EXCEL_EXPORT_QUERY_CONTEXT_BUILDER`. -The feature depends on `boto3`, which is **not** installed by default; install it -with `pip install apache-superset[excel-export]`. +The queued path depends on `boto3`, which is **not** installed by default; install +it with `pip install apache-superset[excel-export]`. The direct-download path +does not use it. Charts store their `query_context` only once they have been (re-)saved in Explore, so older charts may have none. For a fixed, conservative set of viz @@ -529,7 +544,8 @@ types (`table`, `big_number_total`, `big_number`, `pie`) the export rebuilds a query context from the chart's saved form data so those charts still export. The rebuild is a single-query mapping and does **not** reproduce plugin post-processing (pivot, rolling, forecast) or multi-query charts, so any chart of -another type without a saved query context is skipped and listed in the email for +another type without a saved query context is skipped — listed in the email on +the queued path, and simply absent from the workbook on a direct download — for the user to re-save. To cover those types, set `EXCEL_EXPORT_QUERY_CONTEXT_BUILDER` to a callable that receives the chart's form data and returns a query-context payload (or `None` to fall back to the built-in rebuild) — for example one backed diff --git a/docs/docs/using-superset/exporting-dashboard-data.mdx b/docs/docs/using-superset/exporting-dashboard-data.mdx index 7dd2b672711..ce3a8961863 100644 --- a/docs/docs/using-superset/exporting-dashboard-data.mdx +++ b/docs/docs/using-superset/exporting-dashboard-data.mdx @@ -9,16 +9,26 @@ version: 1 Superset can export every chart on a dashboard to a single Excel workbook, with each chart's underlying data rendered as its own worksheet. The export reflects -the dashboard's currently applied filters and runs asynchronously: when it -finishes, the requesting user receives an email with a time-limited download -link. +the dashboard's currently applied filters. + +How the finished workbook reaches you depends on whether the deployment has +export storage configured: + +- **With an export bucket** the work is queued and runs in the background, so it + scales to large dashboards. When it finishes, the requesting user receives an + email with a time-limited download link. +- **Without one** the workbook is built while the request is open and downloads + straight to the browser — no worker, bucket or email needed. Because it has to + finish inside a single request, this path only accepts exports below + `EXCEL_EXPORT_SYNC_MAX_ROWS` (see [Prerequisites](#prerequisites)). ## Using the export From a dashboard, open the **... (actions) → Download** submenu and choose **Export Data to Excel**. The action appears for users who have the dashboard -`can_export` permission. You'll see a confirmation that the export is being -prepared; the workbook arrives by email when it's ready. +`can_export` permission. Where the export is queued you'll see a confirmation +that it is being prepared and the workbook arrives by email; otherwise the file +downloads directly when it is ready. A second option, **Export Images to Excel**, embeds each non-table chart as a rendered image (tables stay tabular) instead of exporting raw data. Because it @@ -36,25 +46,34 @@ Notes on the generated workbook: re-saved in Explore) still exports when it is a `table`, `big_number`, `big_number_total` or `pie`, by rebuilding the query from the chart's saved form data. Charts of other types — and charts relying on post-processing the - rebuild can't reproduce — are skipped and listed in the email; open the chart - in Explore and re-save it to include it next time, or configure - `EXCEL_EXPORT_QUERY_CONTEXT_BUILDER`. + rebuild can't reproduce — are skipped; open the chart in Explore and re-save it + to include it next time, or configure `EXCEL_EXPORT_QUERY_CONTEXT_BUILDER`. A + queued export lists the skipped charts in its email; a direct download has no + email to list them in, so they are simply absent from the workbook. - Row counts per sheet are capped the same way as the chart-level CSV/Excel export (`ROW_LIMIT`, bounded by `SQL_MAX_ROW`), and never exceed Excel's per-sheet maximum. ## Prerequisites -This feature is **disabled by default**. It requires: +Exporting dashboard data needs no configuration: with nothing set up, the +workbook is built during the request and downloaded by the browser. That path is +bounded by `EXCEL_EXPORT_SYNC_MAX_ROWS` (100,000 by default), measured as the +combined `row_limit` of every query the export would run. An export over that +total — or one where any chart has no finite row limit — is refused up front with +a message pointing here, rather than being started and left to hit the web +server's request timeout. Raise the limit only as far as that timeout allows. + +To lift the size ceiling, configure the background path. It requires: 1. **The `boto3` dependency.** It is not installed by default; install it with `pip install apache-superset[excel-export]`. Without it, exports fail and the user receives a failure email. -2. **An S3 bucket.** Set `EXCEL_EXPORT_S3_BUCKET`. Until it is set, the export - endpoint returns `501` and the menu action surfaces a "not configured" - message. -3. **A running Celery worker.** The export runs as a Celery task. If no worker - is running, the request is accepted but nothing is produced. +2. **An S3 bucket.** Set `EXCEL_EXPORT_S3_BUCKET`. This is the switch between the + two paths: with it set, exports are queued and emailed; without it, they are + built inline and downloaded. +3. **A running Celery worker.** The queued export runs as a Celery task. If no + worker is running, the request is accepted but nothing is produced. 4. **A configured SMTP transport.** The download link is delivered by email using the same settings as alerts & reports (`SMTP_*`, `EMAIL_REPORTS_SUBJECT_PREFIX`). @@ -74,7 +93,8 @@ will not register. | Key | Default | Description | | ------------------------------- | ---------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `EXCEL_EXPORT_S3_BUCKET` | `None` | Destination bucket. Required; `501` if unset. | +| `EXCEL_EXPORT_S3_BUCKET` | `None` | Destination bucket. When unset, exports are built during the request and downloaded directly instead of being queued. | +| `EXCEL_EXPORT_SYNC_MAX_ROWS` | `100000` | Largest export served during the request, as the combined `row_limit` of every query it would run. Over this (or with any query lacking a finite limit) the export is refused and asks for a bucket. Ignored once one is configured. | | `EXCEL_EXPORT_S3_KEY_PREFIX` | `"dashboard-exports/"` | Key prefix: `{prefix}{dashboard_id}/{job_id}.xlsx`. | | `EXCEL_EXPORT_LINK_TTL_SECONDS` | `86400` | Lifetime of the pre-signed download URL (24h). | | `EXCEL_EXPORT_S3_CLIENT_KWARGS` | `{}` | Extra kwargs for `boto3.client("s3", ...)` — e.g. `region_name`, or `endpoint_url` for MinIO/LocalStack. | @@ -88,7 +108,9 @@ variables, shared config, or instance role) unless overridden via ## Security considerations - The emailed link is a **pre-signed S3 URL**: anyone who holds it can download - the workbook until it expires. Keep the bucket **private**, enable + the workbook until it expires. This applies to queued exports only — a direct + download is the response to the user's own authenticated request and is never + stored or linked. Keep the bucket **private**, enable encryption, and consider a lifecycle rule to delete objects after a few days. Lower `EXCEL_EXPORT_LINK_TTL_SECONDS` if 24 hours is too long for your data. - The export runs with the requesting user's permissions; each chart's query is @@ -97,8 +119,10 @@ variables, shared config, or instance role) unless overridden via ## Limitations - **Embedded dashboards / guest tokens are not supported** in this version, - because guest users have no email address to deliver the link to. Logged-in - users viewing an embedded dashboard can still use the export. + because guest users have no email address to deliver the link to. This holds on + both paths: a guest session is refused even where the workbook would be + downloaded directly and no email is involved. Logged-in users viewing an + embedded dashboard can still use the export. - The default **Export Data to Excel** mode exports data only (no visual styling). Use **Export Images to Excel** to embed rendered chart images, which requires the webdriver infrastructure described in the prerequisites. diff --git a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadMenuItems.test.tsx b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadMenuItems.test.tsx index c06cd6a06df..861e588c17c 100644 --- a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadMenuItems.test.tsx +++ b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadMenuItems.test.tsx @@ -141,10 +141,40 @@ test('Excel export items are hidden when userCanExport is false', () => { expect(screen.getByText('Export YAML')).toBeInTheDocument(); }); -test('Export Data to Excel posts mode "data" and shows a pending toast', async () => { +/** A queued export: 202 with a job id, delivered later by email. */ +const mockQueuedResponse = ( + body: Record<string, unknown> = { job_id: 'abc' }, +) => + mockSupersetClient.post.mockResolvedValue({ + status: 202, + json: jest.fn().mockResolvedValue(body), + } as never); + +/** An inline export: the workbook itself, as the response to the request. */ +const mockWorkbookResponse = ( + filename = 'World_Health_1.xlsx', +): { blob: jest.Mock } => { + const blob = jest.fn().mockResolvedValue(new Blob(['xlsx'])) as jest.Mock; mockSupersetClient.post.mockResolvedValue({ - json: { job_id: 'abc' }, + status: 200, + blob, + headers: new Headers({ + 'Content-Disposition': `attachment; filename=${filename}`, + }), } as never); + return { blob }; +}; + +/** jsdom implements neither, and the download path needs both. */ +const stubObjectUrls = (): { createObjectURL: jest.Mock } => { + const createObjectURL = jest.fn(() => 'blob:http://localhost/fake'); + window.URL.createObjectURL = createObjectURL; + window.URL.revokeObjectURL = jest.fn(); + return { createObjectURL }; +}; + +test('Export Data to Excel posts mode "data" and shows a pending toast', async () => { + mockQueuedResponse(); render(<MenuWrapper />, { useRedux: true }); @@ -154,6 +184,7 @@ test('Export Data to Excel posts mode "data" and shows a pending toast', async ( expect(mockSupersetClient.post).toHaveBeenCalledWith({ endpoint: '/api/v1/dashboard/123/export_xlsx/', jsonPayload: { active_data_mask: {}, mode: 'data' }, + parseMethod: 'raw', }); expect(mockAddSuccessToast).toHaveBeenCalledWith( "Your export is being prepared. You'll receive an email when it's ready.", @@ -163,9 +194,7 @@ test('Export Data to Excel posts mode "data" and shows a pending toast', async ( test('Export Images to Excel posts mode "images" and shows a pending toast', async () => { enableWebDriverScreenshot(); - mockSupersetClient.post.mockResolvedValue({ - json: { job_id: 'abc' }, - } as never); + mockQueuedResponse(); render(<MenuWrapper />, { useRedux: true }); @@ -175,6 +204,7 @@ test('Export Images to Excel posts mode "images" and shows a pending toast', asy expect(mockSupersetClient.post).toHaveBeenCalledWith({ endpoint: '/api/v1/dashboard/123/export_xlsx/', jsonPayload: { active_data_mask: {}, mode: 'images' }, + parseMethod: 'raw', }); expect(mockAddSuccessToast).toHaveBeenCalledWith( "Your export is being prepared. You'll receive an email when it's ready.", @@ -182,13 +212,55 @@ test('Export Images to Excel posts mode "images" and shows a pending toast', asy }); }); +test('Export Data to Excel downloads the workbook when it arrives inline', async () => { + // Without export storage the server builds the workbook during the request and + // returns the file itself, so the browser downloads it instead of waiting on + // an email that is never sent. + const { blob } = mockWorkbookResponse(); + const { createObjectURL } = stubObjectUrls(); + + render(<MenuWrapper />, { useRedux: true }); + + await userEvent.click(screen.getByText('Export Data to Excel')); + + await waitFor(() => { + expect(blob).toHaveBeenCalled(); + expect(createObjectURL).toHaveBeenCalled(); + expect(mockAddSuccessToast).toHaveBeenCalledWith( + 'Dashboard data exported to Excel', + ); + }); + // The email copy belongs to the queued path only; nothing was queued here. + expect(mockAddSuccessToast).not.toHaveBeenCalledWith( + "Your export is being prepared. You'll receive an email when it's ready.", + ); +}); + +test('Export Data to Excel names the downloaded file from the response', async () => { + mockWorkbookResponse('Sales_Overview_7.xlsx'); + stubObjectUrls(); + // Record the name each download is offered under, and keep jsdom from trying + // to follow the link. + const downloaded: string[] = []; + const click = jest + .spyOn(HTMLAnchorElement.prototype, 'click') + .mockImplementation(function recordDownload(this: HTMLAnchorElement) { + downloaded.push(this.download); + }); + + render(<MenuWrapper />, { useRedux: true }); + + await userEvent.click(screen.getByText('Export Data to Excel')); + + await waitFor(() => expect(downloaded).toEqual(['Sales_Overview_7.xlsx'])); + click.mockRestore(); +}); + test('Export Data to Excel shows an "already in progress" toast when throttled', async () => { // The throttle response is 202 with a message but no job_id. - mockSupersetClient.post.mockResolvedValue({ - json: { - message: 'An Excel export for this dashboard is already in progress.', - }, - } as never); + mockQueuedResponse({ + message: 'An Excel export for this dashboard is already in progress.', + }); render(<MenuWrapper />, { useRedux: true }); @@ -201,18 +273,20 @@ test('Export Data to Excel shows an "already in progress" toast when throttled', }); }); -test('Export Data to Excel shows a config error toast on 501', async () => { - mockSupersetClient.post.mockRejectedValue(new Error('not configured')); - mockGetClientErrorObject.mockResolvedValue({ status: 501 }); +test('Export Data to Excel surfaces the reason an export was refused', async () => { + // An export too large to build during the request is refused with a message + // naming the fix, which is worth more to the user than a generic failure. + const message = + 'This dashboard requests too many rows to export in a single request.'; + mockSupersetClient.post.mockRejectedValue(new Error('too big')); + mockGetClientErrorObject.mockResolvedValue({ status: 400, message }); render(<MenuWrapper />, { useRedux: true }); await userEvent.click(screen.getByText('Export Data to Excel')); await waitFor(() => { - expect(mockAddDangerToast).toHaveBeenCalledWith( - 'Excel export is not configured on this server.', - ); + expect(mockAddDangerToast).toHaveBeenCalledWith(message); }); }); diff --git a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/index.tsx b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/index.tsx index 248d2795260..8109a5fffc9 100644 --- a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/index.tsx +++ b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/index.tsx @@ -120,6 +120,39 @@ export const useDownloadMenuItems = ( } }; + const fileNameFromResponse = ( + response: Response, + fallback: string, + ): string => { + const disposition = response.headers.get('Content-Disposition'); + if (!disposition) { + return fallback; + } + try { + return ( + parseContentDisposition(disposition)?.parameters?.filename ?? fallback + ); + } catch (error) { + logging.warn('Failed to parse Content-Disposition header:', error); + return fallback; + } + }; + + const downloadBlob = (blob: Blob, fileName: string) => { + const url = window.URL.createObjectURL(blob); + try { + const a = document.createElement('a'); + a.href = url; + a.download = fileName; + a.style.display = 'none'; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + } finally { + window.URL.revokeObjectURL(url); + } + }; + const onExportAsExample = async () => { try { const response = await SupersetClient.get({ @@ -130,35 +163,11 @@ export const useDownloadMenuItems = ( parseMethod: 'raw', }); - // Parse filename from Content-Disposition header - const disposition = response.headers.get('Content-Disposition'); - let fileName = `dashboard_${dashboardId}_example.zip`; - - if (disposition) { - try { - const parsed = parseContentDisposition(disposition); - if (parsed?.parameters?.filename) { - fileName = parsed.parameters.filename; - } - } catch (error) { - logging.warn('Failed to parse Content-Disposition header:', error); - } - } - - // Convert response to blob and trigger download const blob = await response.blob(); - const url = window.URL.createObjectURL(blob); - try { - const a = document.createElement('a'); - a.href = url; - a.download = fileName; - a.style.display = 'none'; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - } finally { - window.URL.revokeObjectURL(url); - } + downloadBlob( + blob, + fileNameFromResponse(response, `dashboard_${dashboardId}_example.zip`), + ); addSuccessToast(t('Dashboard exported as example successfully')); } catch (error) { @@ -169,13 +178,31 @@ export const useDownloadMenuItems = ( const onExportXlsx = async (mode: 'data' | 'images') => { try { - const { json } = await SupersetClient.post({ + const response = await SupersetClient.post({ endpoint: `/api/v1/dashboard/${dashboardId}/export_xlsx/`, jsonPayload: { active_data_mask: buildActiveDataMask(), mode }, + // The response is either JSON describing a queued export or the workbook + // itself, so it is parsed here rather than by the client. + parseMethod: 'raw', }); + + // Where the deployment has export storage the work is queued and delivered + // by email (202). Where it has none the server builds the workbook during + // the request and returns the file, which the browser downloads directly. + if (response.status !== 202) { + const blob = await response.blob(); + downloadBlob( + blob, + fileNameFromResponse(response, `dashboard_${dashboardId}.xlsx`), + ); + addSuccessToast(t('Dashboard data exported to Excel')); + return; + } + // The throttle response (an export is already running) returns 202 with a // message but no job_id; only a freshly enqueued job carries a job_id. - if ((json as { job_id?: string })?.job_id) { + const json = (await response.json()) as { job_id?: string }; + if (json?.job_id) { addSuccessToast( t( "Your export is being prepared. You'll receive an email when it's ready.", @@ -187,13 +214,18 @@ export const useDownloadMenuItems = ( ); } } catch (error) { - // status comes from the response (Partial<SupersetClientResponse>), which - // the union type does not expose uniformly; read it via a narrow cast. - const { status } = (await getClientErrorObject(error)) as { + // status/message come from the response (Partial<SupersetClientResponse>), + // which the union type does not expose uniformly; read them via a narrow + // cast. + const { status, message } = (await getClientErrorObject(error)) as { status?: number; + message?: string; }; - if (status === 501) { - addDangerToast(t('Excel export is not configured on this server.')); + // A refusal explains itself — the export is too large to build in one + // request, and says what to configure — so pass it on rather than + // replacing it with a generic failure. Server-side faults stay generic. + if (message && status && status >= 400 && status < 500) { + addDangerToast(message); } else { addDangerToast(t('Sorry, something went wrong. Try again later.')); } diff --git a/superset/config.py b/superset/config.py index 6886ee48df8..5dea16ef1d3 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1559,6 +1559,15 @@ EXCEL_EXPORT_S3_CLIENT_KWARGS: dict[str, Any] = {} # a rendered image. Set to None to fall back to the built-in default. EXCEL_EXPORT_TABLE_VIZ_TYPES: set[str] | None = None +# Ceiling for an export served inline, as the response to the request that asked +# for it — the path taken when no bucket is configured above. The export adds up +# the ``row_limit`` of every query it would run and refuses, before running any +# of them, when the total exceeds this (or when any query has no finite limit), +# pointing the user at the asynchronous path instead of risking a request that +# outlives its timeout. Raise it only as far as the deployment's own request +# timeout allows. +EXCEL_EXPORT_SYNC_MAX_ROWS = 100_000 + # Optional hook to build a query context for a chart that has no saved # ``query_context``, called before the built-in form-data rebuild. Receives the # chart's form data (its ``params`` with ``viz_type`` and the diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py index 650df5e442f..66482c43e44 100644 --- a/superset/dashboards/api.py +++ b/superset/dashboards/api.py @@ -17,6 +17,8 @@ # pylint: disable=too-many-lines import functools import logging +import os +import tempfile import uuid from datetime import datetime from io import BytesIO @@ -92,6 +94,9 @@ from superset.commands.importers.v1.utils import get_contents_from_bundle from superset.commands.purge import PurgeArchivedCommand, SoftDeleteBinding from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod from superset.daos.dashboard import DashboardDAO, EmbeddedDashboardDAO +from superset.dashboards.excel_export.storage import is_export_storage_configured +from superset.dashboards.excel_export.sync_budget import is_within_sync_row_budget +from superset.dashboards.excel_export.workbook import build_workbook from superset.dashboards.filter_scope import derive_json_metadata from superset.dashboards.filters import ( DashboardAccessFilter, @@ -180,6 +185,7 @@ from superset.versioning.api_helpers import ( ) from superset.versioning.etag import set_version_etag from superset.versioning.schemas import VersionListItemSchema +from superset.views.base import generate_download_headers, XlsxResponse from superset.views.base_api import ( BaseSupersetModelRestApi, RelatedFieldFilter, @@ -1727,14 +1733,17 @@ class DashboardRestApi( log_to_statsd=False, ) def export_xlsx(self, pk: int) -> WerkzeugResponse: - """Export all of a dashboard's chart data to an Excel workbook (async). + """Export all of a dashboard's chart data to an Excel workbook. --- post: summary: Export dashboard chart data to Excel description: >- - Enqueues an async task that writes each chart's data to its own - worksheet, uploads the .xlsx to S3, and emails the requesting user a - pre-signed download link. Returns immediately with a job id. + Writes each chart's data to its own worksheet of a single .xlsx. + Where export storage is configured the work is queued: the response + is a job id and the finished file is uploaded and emailed to the + requesting user as a download link. Where it is not, the workbook is + built during the request and returned as the response body, provided + the export is small enough to serve that way. parameters: - in: path schema: @@ -1747,6 +1756,13 @@ class DashboardRestApi( schema: $ref: '#/components/schemas/DashboardExportXlsxPostSchema' responses: + 200: + description: The exported workbook, built during this request + content: + application/vnd.openxmlformats-officedocument.spreadsheetml.sheet: + schema: + type: string + format: binary 202: description: Export task accepted content: @@ -1763,13 +1779,11 @@ class DashboardRestApi( $ref: '#/components/responses/404' 500: $ref: '#/components/responses/500' - 501: - description: Excel export is not configured on this server """ - if not current_app.config["EXCEL_EXPORT_S3_BUCKET"]: - return self.response( - 501, message="Excel export is not configured on this server." - ) + # With storage the export is queued and delivered by link; without it the + # workbook is built here and returned as the response. Resolved once, so a + # single request cannot take one path's checks and the other's delivery. + queued = is_export_storage_configured() try: # Tolerate an empty/non-JSON body (e.g. a POST with no Content-Type); # request.json would otherwise raise 415. @@ -1806,11 +1820,28 @@ class DashboardRestApi( if not dashboard.slices: return self.response_400(message="Dashboard has no charts to export.") + active_data_mask = payload.get("active_data_mask", {}) + mode = payload.get("mode", "data") + + # An export served as the response has to finish inside this request, so + # refuse an oversized one before doing any of the work rather than letting + # it run into a gateway timeout. Checked ahead of the lock so a refusal + # never leaves a lock to be released. + if not queued and not is_within_sync_row_budget(dashboard, mode): + return self.response_400( + message=( + "This dashboard requests too many rows to export in a single " + "request. Configure EXCEL_EXPORT_S3_BUCKET to export it in the " + "background, or lower the row limits of its charts." + ) + ) + # Throttle: one concurrent export per user+dashboard. Acquire a shared, # atomic distributed lock (Redis when configured, the metadata DB # otherwise) so the guard works across the web server and workers and is - # not a no-op under the default cache. The task releases it when it - # settles; the TTL is the backstop if that release is ever lost. + # not a no-op under the default cache. The queued path's task releases it + # when it settles and the inline path releases it before responding; the + # TTL is the backstop if either release is ever lost. lock_params = export_lock_params(g.user.id, dashboard.id) try: AcquireDistributedLock( @@ -1825,14 +1856,31 @@ class DashboardRestApi( ) job_id = str(uuid.uuid4()) + run_export = self._export_xlsx_queued if queued else self._export_xlsx_inline + return run_export(dashboard, active_data_mask, mode, job_id, lock_params) + + def _export_xlsx_queued( # pylint: disable=too-many-arguments + self, + dashboard: Dashboard, + active_data_mask: dict[str, Any], + mode: str, + job_id: str, + lock_params: dict[str, int], + ) -> WerkzeugResponse: + """ + Hand the export to a worker, which uploads it and emails a download link. + + Used where export storage is configured. Returns as soon as the job is + queued, so an export of any size is free to take as long as it needs. + """ try: export_dashboard_excel.apply_async( kwargs={ "dashboard_id": dashboard.id, "user_id": g.user.id, - "active_data_mask": payload.get("active_data_mask", {}), + "active_data_mask": active_data_mask, "job_id": job_id, - "mode": payload.get("mode", "data"), + "mode": mode, }, task_id=job_id, ) @@ -1844,6 +1892,58 @@ class DashboardRestApi( raise return self.response(202, job_id=job_id) + @staticmethod + def _export_xlsx_inline( # pylint: disable=too-many-arguments + dashboard: Dashboard, + active_data_mask: dict[str, Any], + mode: str, + job_id: str, + lock_params: dict[str, int], + ) -> WerkzeugResponse: + """ + Build the export during this request and return it as the response. + + Used where no export storage is configured, so there is nowhere to upload + a finished file and nothing to link to in an email. The workbook is the + same one the Celery task builds, from the same builder: only the delivery + differs. It is written to a temp file (the writer streams to disk in + constant memory) and read back once, so the response carries a complete + file and the temp file never outlives the request. + + The charts the export had to skip are not reported here. The queued path + lists them in its email, which this path has no equivalent of; the + workbook itself is identical either way. + """ + tmp_path: str | None = None + try: + file_descriptor, tmp_path = tempfile.mkstemp( + suffix=".xlsx", prefix=f"dash-export-{job_id}-" + ) + os.close(file_descriptor) + + build_workbook(tmp_path, dashboard, active_data_mask, job_id, mode, g.user) + with open(tmp_path, "rb") as workbook: + content = workbook.read() + finally: + # Both of these have to happen however the export ends: a held lock + # would keep the user from retrying until its TTL expires, and an + # abandoned temp file would sit on the web server's disk. + try: + ReleaseDistributedLock(EXPORT_LOCK_NAMESPACE, lock_params).run() + except Exception: # pylint: disable=broad-except + # Best-effort: the lock's TTL is the backstop if this fails. + logger.exception( + "Failed to release in-flight export lock for dashboard %s", + dashboard.id, + ) + if tmp_path and os.path.exists(tmp_path): + os.remove(tmp_path) + + filename = get_filename(dashboard.dashboard_title, dashboard.id, skip_id=False) + return XlsxResponse( + content, headers=generate_download_headers("xlsx", filename) + ) + def _validate_permalink_for_dashboard( self, permalink_key: str, dashboard: Dashboard ) -> WerkzeugResponse | None: diff --git a/superset/dashboards/excel_export/storage.py b/superset/dashboards/excel_export/storage.py new file mode 100644 index 00000000000..ee9a7da526b --- /dev/null +++ b/superset/dashboards/excel_export/storage.py @@ -0,0 +1,37 @@ +# 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. +""" +Availability of the object storage the asynchronous dashboard Excel export needs. + +This is the single place that answers "can this deployment run an asynchronous +export?", and it is deliberately the only thing that knows *how* that is +configured. The dashboard API branches on it to pick the export path: with +storage the export is queued and delivered by link, without it the workbook is +built inline and streamed back as the response. + +The answer is derived from configuration rather than exposed as a feature flag, +so there is no second knob that can disagree with the storage settings. +""" + +from __future__ import annotations + +from flask import current_app + + +def is_export_storage_configured() -> bool: + """Whether generated exports can be uploaded somewhere and served by link.""" + return bool(current_app.config["EXCEL_EXPORT_S3_BUCKET"]) diff --git a/superset/dashboards/excel_export/sync_budget.py b/superset/dashboards/excel_export/sync_budget.py new file mode 100644 index 00000000000..34b6f445454 --- /dev/null +++ b/superset/dashboards/excel_export/sync_budget.py @@ -0,0 +1,98 @@ +# 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. +""" +Decide whether a dashboard is small enough to export inline. + +An export served as the HTTP response has to finish inside one request, so the +size of the workbook is settled *before* any query runs, by adding up the rows +the export is allowed to ask for: the ``row_limit`` of every query it would run. +A request/server timeout is the last-resort backstop, not the criterion — a +timed-out export wastes the work already done and tells the user nothing +actionable, whereas an up-front refusal can name the fix. + +The total is deliberately the *requested* row count rather than the delivered +one. It is knowable without touching a database, and it is an upper bound: an +export that clears the budget cannot exceed it once the queries run. +""" + +from __future__ import annotations + +from typing import Any + +from flask import current_app + +from superset.dashboards.excel_export.layout import get_charts_in_layout_order +from superset.dashboards.excel_export.workbook import ( + renders_as_image, + resolve_query_context, +) + + +def _finite_row_limit(query: Any) -> int | None: + """ + A query's ``row_limit`` when it bounds the result, else ``None``. + + Anything else — absent, ``0`` (which defers to the deployment's configured + limits), negative, or not an integer — leaves the query's size unknown. + ``bool`` is rejected too: it is an ``int`` subclass, so ``True`` would + otherwise pass as a limit of one row. + """ + if not isinstance(query, dict): + return None + row_limit = query.get("row_limit") + if isinstance(row_limit, bool) or not isinstance(row_limit, int): + return None + return row_limit if row_limit > 0 else None + + +def requested_row_total(dashboard: Any, mode: str) -> int | None: + """ + Total rows every query in this export is allowed to return. + + Returns ``None`` when any query the export would run has no finite row limit, + which makes the total — and so the size of the export — indeterminate. + + Charts the export cannot run contribute nothing: one with no usable query + context is skipped by the export itself, and in image mode a non-table chart + is rendered rather than queried. + + Note that this resolves each chart's query context, which the export then + resolves again when it builds the workbook. For a saved context that is a + JSON parse; a deployment using ``EXCEL_EXPORT_QUERY_CONTEXT_BUILDER`` pays + for its hook twice on this path. + """ + total = 0 + for chart in get_charts_in_layout_order(dashboard): + if renders_as_image(chart, mode): + continue + query_context = resolve_query_context(chart) + if query_context is None: + continue + for query in query_context["queries"]: + row_limit = _finite_row_limit(query) + if row_limit is None: + return None + total += row_limit + return total + + +def is_within_sync_row_budget(dashboard: Any, mode: str) -> bool: + """Whether this export may run inline, as the response to one request.""" + total = requested_row_total(dashboard, mode) + return ( + total is not None and total <= current_app.config["EXCEL_EXPORT_SYNC_MAX_ROWS"] + ) diff --git a/superset/tasks/export_dashboard_excel.py b/superset/dashboards/excel_export/workbook.py similarity index 67% copy from superset/tasks/export_dashboard_excel.py copy to superset/dashboards/excel_export/workbook.py index 5fb95d95bd2..154c1a65b9c 100644 --- a/superset/tasks/export_dashboard_excel.py +++ b/superset/dashboards/excel_export/workbook.py @@ -15,38 +15,38 @@ # 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. - -In ``"data"`` mode 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. In ``"images"`` mode non-table charts are instead rendered to -images (through the same headless path as scheduled reports, reflecting the live -filters) and embedded, while table-like charts stay tabular. +Build the multi-sheet ``.xlsx`` workbook for a dashboard export. + +This module owns everything that turns a dashboard into a workbook on disk: +resolving each chart's query context, applying the live dashboard filter state, +running the chart queries under the requesting user's permissions, and streaming +the results row-by-row into a constant-memory workbook so large dashboards never +load all data at once. In ``"images"`` mode non-table charts are instead rendered +to images (through the same headless path as scheduled reports, reflecting the +live filters) and embedded, while table-like charts stay tabular. + +It is deliberately free of any delivery concern. Both export paths share it: +:mod:`superset.tasks.export_dashboard_excel` wraps it in a Celery task that +uploads the result to object storage and emails a link, while the dashboard API +calls it inline to stream the workbook back as the HTTP response when no export +storage is configured. """ from __future__ import annotations import copy import logging -import os -import tempfile -from datetime import datetime, timedelta, timezone 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_dashboard_filter_context, get_dashboard_filter_context, ) from superset.charts.schemas import ChartDataQueryContextSchema from superset.commands.chart.data.get_data_command import ChartDataCommand -from superset.commands.distributed_lock.release import ReleaseDistributedLock from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType from superset.common.form_data_query_context import ( build_query_context_from_form_data, @@ -55,9 +55,7 @@ from superset.common.form_data_query_context import ( from superset.dashboards.excel_export import email from superset.dashboards.excel_export.layout import get_charts_in_layout_order from superset.dashboards.excel_export.screenshot import render_chart_image -from superset.extensions import celery_app -from superset.utils import json, s3 -from superset.utils.core import override_user +from superset.utils import json from superset.utils.excel_streaming import StreamingXlsxWriter logger = logging.getLogger(__name__) @@ -78,31 +76,12 @@ TABLE_VIZ_TYPES = {"table", "pivot_table_v2", "pivot_table"} # saved query context is skipped and listed for the user to re-save in Explore. REBUILD_VIZ_TYPES = {"table", "big_number_total", "big_number", "pie"} -EXPORT_SOFT_TIME_LIMIT = 600 -EXPORT_HARD_TIME_LIMIT = 660 -# Namespace + TTL for the per-user+dashboard in-flight lock the API acquires -# before enqueue and this task releases when it settles. The lock uses the -# shared, atomic DistributedLock backend (Redis when configured, the metadata -# DB otherwise) so it actually synchronizes across the web server and workers — -# unlike a plain cache, which is a no-op under the default ``NullCache``. -# The TTL outlives the hard time limit so a worker killed at that limit (which -# skips the ``finally`` release) cannot hold the lock forever; the release in -# ``finally`` is the fast path that frees it as soon as the task settles. -EXPORT_LOCK_NAMESPACE = "excel_export" -EXPORT_LOCK_TTL_SECONDS = EXPORT_HARD_TIME_LIMIT + 60 - - -def export_lock_params(user_id: int, dashboard_id: int) -> dict[str, int]: - """Key parameters identifying the per-user+dashboard in-flight lock.""" - return {"user_id": user_id, "dashboard_id": dashboard_id} - - -class _ChartSkippedError(Exception): +class ChartSkippedError(Exception): """Signals a chart that could not be exported and should be listed as skipped.""" -def _chart_label(chart: Any) -> str: +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() @@ -169,7 +148,7 @@ def _needs_unsupported_processing(form_data: dict[str, Any]) -> bool: return form_data.get("aggregation") == "raw" -def _resolve_query_context(chart: Any) -> dict[str, Any] | None: +def resolve_query_context(chart: Any) -> dict[str, Any] | None: """ The query-context payload to run for a chart's data export, or ``None`` when none can be obtained. @@ -201,7 +180,7 @@ def _resolve_query_context(chart: Any) -> dict[str, Any] | None: built = builder(chart.form_data) except SoftTimeLimitExceeded: # A soft timeout is a task-level signal, not a builder failure: let it - # propagate to _build_workbook so the export aborts cleanly instead of + # propagate to build_workbook so the export aborts cleanly instead of # continuing on to rebuild this chart and start the next one. raise except Exception: # pylint: disable=broad-except @@ -249,7 +228,7 @@ def _table_viz_types() -> set[str]: return current_app.config.get("EXCEL_EXPORT_TABLE_VIZ_TYPES") or TABLE_VIZ_TYPES -def _renders_as_image(chart: Any, mode: str) -> bool: +def renders_as_image(chart: Any, mode: str) -> bool: """Whether this chart is embedded as an image rather than streamed as data.""" return mode == EXPORT_MODE_IMAGES and chart.viz_type not in _table_viz_types() @@ -264,12 +243,12 @@ def _write_chart_image_sheet( """ Render a single chart to an image and embed it as its own sheet. - :raises _ChartSkippedError: if the chart could not be rendered + :raises ChartSkippedError: if the chart could not be rendered """ image = render_chart_image(chart, dashboard_id, active_data_mask, user) if image is None: - raise _ChartSkippedError - writer.add_image_sheet(_chart_label(chart), image) + raise ChartSkippedError + writer.add_image_sheet(chart_label(chart), image) def _write_chart_sheets( @@ -291,7 +270,7 @@ def _write_chart_sheets( # Shallow-copy before setting our own top-level keys so the caller's payload # keeps its original result_format/result_type. (The nested ``queries`` are # mutated in place by apply_dashboard_filter_context below, which is safe - # because every payload ``_resolve_query_context`` returns is this chart's + # because every payload ``resolve_query_context`` returns is this chart's # alone: freshly parsed, freshly built, or deep-copied from the builder hook.) json_body = dict(json_body) # Override any stale saved values: we always want full JSON results. @@ -329,7 +308,7 @@ def _write_chart_sheets( ) -def _build_workbook( +def build_workbook( path: str, dashboard: Any, active_data_mask: dict[str, Any], @@ -342,14 +321,21 @@ def _build_workbook( Return the charts that could not be exported, grouped by the reason they were omitted (see the ``email.ERROR_*`` reason keys), so the notification can explain each group separately. + + :param path: Destination path for the ``.xlsx`` file + :param dashboard: The dashboard whose charts to export + :param active_data_mask: Live dashboard filter state keyed by native filter id + :param job_id: Correlation id used in log lines + :param mode: ``"data"`` or ``"images"`` + :param user: The requesting user (used to render images) """ errored: dict[str, list[str]] = {} writer = StreamingXlsxWriter(path) try: for chart in get_charts_in_layout_order(dashboard): - label = _chart_label(chart) + label = chart_label(chart) try: - if _renders_as_image(chart, mode): + if renders_as_image(chart, mode): # Image charts render from their saved params via the # webdriver and don't need a query context. _write_chart_image_sheet( @@ -359,7 +345,7 @@ def _build_workbook( # Data charts need a query context: use the saved one, or # rebuild it from form data for eligible viz types. Skip # cleanly when none is available rather than failing. - json_body = _resolve_query_context(chart) + json_body = resolve_query_context(chart) if json_body is None: errored.setdefault(email.ERROR_NO_QUERY_CONTEXT, []).append( label @@ -374,9 +360,10 @@ def _build_workbook( # cleanup, rather than continuing until the hard limit kills the # worker (which would skip cleanup, leak temp files, and hold the # in-flight lock until its TTL). ``except Exception`` below would - # otherwise swallow it, since it subclasses ``Exception``. + # otherwise swallow it, since it subclasses ``Exception``. Only the + # Celery path can raise it; the synchronous path never does. raise - except _ChartSkippedError: + except ChartSkippedError: logger.warning( "Skipping chart %s in dashboard export %s (could not render)", chart.id, @@ -398,123 +385,3 @@ def _build_workbook( finally: writer.close() return errored - - -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=EXPORT_SOFT_TIME_LIMIT, - time_limit=EXPORT_HARD_TIME_LIMIT, - 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, - mode: str = EXPORT_MODE_DATA, -) -> None: - """ - Export a dashboard's charts 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 - :param mode: ``"data"`` streams every chart's tabular result; ``"images"`` - embeds non-table charts as rendered images and keeps tables tabular - """ - # pylint: disable=import-outside-toplevel - from superset.models.dashboard import Dashboard - - requested_at = datetime.now(tz=timezone.utc) - 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) - - errored = _build_workbook( - tmp_path, dashboard, active_data_mask, job_id, mode, user - ) - - 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.now(tz=timezone.utc) + 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, - errored=errored, - ), - ) - 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: - try: - ReleaseDistributedLock( - EXPORT_LOCK_NAMESPACE, - export_lock_params(user_id, dashboard_id), - ).run() - except Exception: # pylint: disable=broad-except - # Best-effort: the lock's TTL is the backstop if this fails. - logger.exception( - "Failed to release in-flight export lock for user %s dashboard %s", - user_id, - dashboard_id, - ) - if tmp_path and os.path.exists(tmp_path): - os.remove(tmp_path) diff --git a/superset/tasks/export_dashboard_excel.py b/superset/tasks/export_dashboard_excel.py index 5fb95d95bd2..624db29c1c0 100644 --- a/superset/tasks/export_dashboard_excel.py +++ b/superset/tasks/export_dashboard_excel.py @@ -19,17 +19,15 @@ 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. -In ``"data"`` mode 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. In ``"images"`` mode non-table charts are instead rendered to -images (through the same headless path as scheduled reports, reflecting the live -filters) and embedded, while table-like charts stay tabular. +The workbook itself is built by +:func:`superset.dashboards.excel_export.workbook.build_workbook`, which the +dashboard API also calls inline when no export storage is configured. This module +owns only the asynchronous concerns: the Celery task, the storage upload, email +delivery, and releasing the in-flight lock the API acquired before enqueueing. """ from __future__ import annotations -import copy import logging import os import tempfile @@ -37,52 +35,24 @@ from datetime import datetime, timedelta, timezone from typing import Any from celery.exceptions import SoftTimeLimitExceeded -from flask import current_app, g +from flask import current_app from superset import db, security_manager -from superset.charts.data.dashboard_filter_context import ( - apply_dashboard_filter_context, - get_dashboard_filter_context, -) -from superset.charts.schemas import ChartDataQueryContextSchema -from superset.commands.chart.data.get_data_command import ChartDataCommand from superset.commands.distributed_lock.release import ReleaseDistributedLock -from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType -from superset.common.form_data_query_context import ( - build_query_context_from_form_data, - is_raw_query_mode, -) from superset.dashboards.excel_export import email -from superset.dashboards.excel_export.layout import get_charts_in_layout_order -from superset.dashboards.excel_export.screenshot import render_chart_image +from superset.dashboards.excel_export.workbook import build_workbook, EXPORT_MODE_DATA from superset.extensions import celery_app -from superset.utils import json, s3 +from superset.utils import s3 from superset.utils.core import override_user -from superset.utils.excel_streaming import StreamingXlsxWriter logger = logging.getLogger(__name__) -# Export modes: "data" streams every chart's tabular result (the default, -# unchanged behavior); "images" embeds non-table charts as rendered images and -# keeps only table-like charts tabular. -EXPORT_MODE_DATA = "data" -EXPORT_MODE_IMAGES = "images" - -# Viz types kept as tabular data in image mode; everything else is rendered as an -# image. Operators can override the set via ``EXCEL_EXPORT_TABLE_VIZ_TYPES``. -TABLE_VIZ_TYPES = {"table", "pivot_table_v2", "pivot_table"} - -# Viz types whose missing query context may be rebuilt from saved form data. -# Conservative: only charts whose data maps faithfully to a single plain query -# (no post-processing, no multi-query fan-out). Every other viz type without a -# saved query context is skipped and listed for the user to re-save in Explore. -REBUILD_VIZ_TYPES = {"table", "big_number_total", "big_number", "pie"} - EXPORT_SOFT_TIME_LIMIT = 600 EXPORT_HARD_TIME_LIMIT = 660 # Namespace + TTL for the per-user+dashboard in-flight lock the API acquires -# before enqueue and this task releases when it settles. The lock uses the +# before either export path runs, released by this task when it settles (the +# synchronous path releases it in its own ``finally``). The lock uses the # shared, atomic DistributedLock backend (Redis when configured, the metadata # DB otherwise) so it actually synchronizes across the web server and workers — # unlike a plain cache, which is a no-op under the default ``NullCache``. @@ -98,308 +68,6 @@ def export_lock_params(user_id: int, dashboard_id: int) -> dict[str, int]: return {"user_id": user_id, "dashboard_id": dashboard_id} -class _ChartSkippedError(Exception): - """Signals a chart that could not be exported and should be listed as skipped.""" - - -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 _usable_query_context(value: Any) -> dict[str, Any] | None: - """ - ``value`` when it is a usable query-context payload, else ``None``. - - A payload is usable only if it is a dict with a non-empty ``queries`` list; a - blank, query-less, mistyped, or non-object value (e.g. ``{}``, - ``{"queries": []}``, ``{"queries": "oops"}``, ``None``) is treated the same as - a missing context. Shared by the saved-context path and the builder hook so - both apply the same validity rule — and so a malformed builder return falls - through to the built-in rebuild instead of failing later in the general - error bucket. - """ - if not isinstance(value, dict) or not isinstance(value.get("queries"), list): - return None - return value if value["queries"] else None - - -def _saved_query_context(raw: Any) -> dict[str, Any] | None: - """ - The chart's saved query context parsed to a dict, or ``None`` when it is - missing or unusable. - - Returns ``None`` for a blank value, a string that does not parse as JSON, and - any value that is not a dict with a non-empty ``queries`` list. - """ - if not raw: - return None - try: - parsed = json.loads(raw) - except (TypeError, ValueError): - return None - return _usable_query_context(parsed) - - -# Form-data keys whose behavior needs plugin post-processing or extra queries -# (contribution/time comparison, rolling window, resampling, raw big-number -# aggregation) that the single-query rebuild cannot reproduce. A chart using any -# of these is skipped rather than exported with values that differ from the chart. -_UNSUPPORTED_PROCESSING_KEYS = ("time_compare", "rolling_type", "resample_rule") - - -def _needs_unsupported_processing(form_data: dict[str, Any]) -> bool: - """Whether the form data relies on processing the rebuild can't reproduce.""" - # ``percent_metrics`` are "% of total" columns produced by contribution - # post-processing the rebuild can't apply; skip so the export doesn't silently - # omit columns the user sees. - if form_data.get("percent_metrics"): - return True - # ``show_totals`` adds a totals row via a *second* query - # (``plugin-chart-table/src/buildQuery.ts``, gated on aggregate mode); the - # single-query rebuild would silently drop that row. The mode check mirrors - # the frontend so a raw-mode table carrying a stale value still exports. - if form_data.get("show_totals") and not is_raw_query_mode(form_data): - return True - for key in _UNSUPPORTED_PROCESSING_KEYS: - value = form_data.get(key) - # ``rolling_type`` is often the literal string ``"None"`` when unset. - if value and value != "None": - return True - return form_data.get("aggregation") == "raw" - - -def _resolve_query_context(chart: Any) -> dict[str, Any] | None: - """ - The query-context payload to run for a chart's data export, or ``None`` when - none can be obtained. - - Resolution order: - - 1. the chart's saved ``query_context``; - 2. an optional ``EXCEL_EXPORT_QUERY_CONTEXT_BUILDER`` hook, letting a deployment - supply a faithful context (e.g. from a service running the chart's real - frontend ``buildQuery``) for viz types the built-in rebuild can't handle; - 3. the built-in form-data rebuild, restricted to viz types whose data maps - faithfully to a single plain query (``REBUILD_VIZ_TYPES``) without - post-processing or extra queries. - - Returns ``None`` when none apply, so the caller lists the chart for re-saving - rather than exporting inaccurate data. - """ - if saved := _saved_query_context(chart.query_context): - return saved - - # The hook receives the chart's form data and must return ``None`` — not a - # partial/stub context — whenever it can't build the chart faithfully, so we - # fall through to the built-in rebuild (which handles the allowlisted viz types - # well). A hook failure falls through too, preserving "builder problem → - # rebuild, don't fail the export" — the one exception being a task-level - # timeout, which has to abort the whole export rather than this chart. - if builder := current_app.config.get("EXCEL_EXPORT_QUERY_CONTEXT_BUILDER"): - try: - built = builder(chart.form_data) - except SoftTimeLimitExceeded: - # A soft timeout is a task-level signal, not a builder failure: let it - # propagate to _build_workbook so the export aborts cleanly instead of - # continuing on to rebuild this chart and start the next one. - raise - except Exception: # pylint: disable=broad-except - logger.warning( - "EXCEL_EXPORT_QUERY_CONTEXT_BUILDER failed for chart %s; " - "falling back to the built-in rebuild", - chart.id, - exc_info=True, - ) - built = None - if (from_builder := _usable_query_context(built)) is not None: - # Copy: the payload's ``queries`` are mutated in place downstream (by - # ``apply_dashboard_filter_context``), and a builder is free to - # memoize or otherwise share its return value — which would then - # accumulate filters across charts and across exports. - return copy.deepcopy(from_builder) - - # The allowlist and ``_needs_unsupported_processing`` bound only the built-in - # rebuild below; the builder hook above is intentionally not gated by them (a - # faithful builder includes the post-processing the built-in rebuild lacks). - if chart.viz_type not in REBUILD_VIZ_TYPES or chart.datasource_id is None: - return None - try: - form_data = json.loads(chart.params) if chart.params else {} - except (TypeError, ValueError): - return None - if not isinstance(form_data, dict) or not form_data: - return None - if _needs_unsupported_processing(form_data): - return None - - return build_query_context_from_form_data( - form_data, - {"id": chart.datasource_id, "type": chart.datasource_type or "table"}, - chart.viz_type, - ) - - -def _record_to_row(record: dict[str, Any], colnames: list[str]) -> list[Any]: - return [record.get(col) for col in colnames] - - -def _table_viz_types() -> set[str]: - """Viz types kept tabular in image mode (config override or built-in default).""" - return current_app.config.get("EXCEL_EXPORT_TABLE_VIZ_TYPES") or TABLE_VIZ_TYPES - - -def _renders_as_image(chart: Any, mode: str) -> bool: - """Whether this chart is embedded as an image rather than streamed as data.""" - return mode == EXPORT_MODE_IMAGES and chart.viz_type not in _table_viz_types() - - -def _write_chart_image_sheet( - writer: StreamingXlsxWriter, - chart: Any, - dashboard_id: int, - active_data_mask: dict[str, Any], - user: Any, -) -> None: - """ - Render a single chart to an image and embed it as its own sheet. - - :raises _ChartSkippedError: if the chart could not be rendered - """ - image = render_chart_image(chart, dashboard_id, active_data_mask, user) - if image is None: - raise _ChartSkippedError - writer.add_image_sheet(_chart_label(chart), image) - - -def _write_chart_sheets( - writer: StreamingXlsxWriter, - chart: Any, - json_body: dict[str, 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. - - ``json_body`` is the resolved query-context payload (the chart's saved - context or one synthesized from its form data). 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. - """ - # Shallow-copy before setting our own top-level keys so the caller's payload - # keeps its original result_format/result_type. (The nested ``queries`` are - # mutated in place by apply_dashboard_filter_context below, which is safe - # because every payload ``_resolve_query_context`` returns is this chart's - # alone: freshly parsed, freshly built, or deep-copied from the builder hook.) - json_body = dict(json_body) - # 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, - ) - if filter_context.extra_form_data: - apply_dashboard_filter_context(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, - mode: str, - user: Any, -) -> dict[str, list[str]]: - """Build the workbook on disk. - - Return the charts that could not be exported, grouped by the reason they - were omitted (see the ``email.ERROR_*`` reason keys), so the notification - can explain each group separately. - """ - errored: dict[str, list[str]] = {} - writer = StreamingXlsxWriter(path) - try: - for chart in get_charts_in_layout_order(dashboard): - label = _chart_label(chart) - try: - if _renders_as_image(chart, mode): - # Image charts render from their saved params via the - # webdriver and don't need a query context. - _write_chart_image_sheet( - writer, chart, dashboard.id, active_data_mask, user - ) - else: - # Data charts need a query context: use the saved one, or - # rebuild it from form data for eligible viz types. Skip - # cleanly when none is available rather than failing. - json_body = _resolve_query_context(chart) - if json_body is None: - errored.setdefault(email.ERROR_NO_QUERY_CONTEXT, []).append( - label - ) - continue - _write_chart_sheets( - writer, chart, json_body, dashboard.id, active_data_mask - ) - except SoftTimeLimitExceeded: - # A soft timeout is a task-level signal, not a per-chart failure: - # let it propagate so the outer handler emails a failure and runs - # cleanup, rather than continuing until the hard limit kills the - # worker (which would skip cleanup, leak temp files, and hold the - # in-flight lock until its TTL). ``except Exception`` below would - # otherwise swallow it, since it subclasses ``Exception``. - raise - except _ChartSkippedError: - logger.warning( - "Skipping chart %s in dashboard export %s (could not render)", - chart.id, - job_id, - ) - errored.setdefault(email.ERROR_GENERAL, []).append(label) - except Exception: # pylint: disable=broad-except - logger.exception( - "Skipping chart %s in dashboard export %s", chart.id, job_id - ) - errored.setdefault(email.ERROR_GENERAL, []).append(label) - - if writer.sheet_count == 0: - flat = [label for labels in errored.values() for label in labels] - writer.add_summary_sheet( - "Export Summary", - ["No chart data could be exported.", *flat], - ) - finally: - writer.close() - return errored - - def _send_failure_email( user: Any, dashboard_title: str, requested_at: datetime ) -> None: @@ -462,7 +130,7 @@ def export_dashboard_excel( ) os.close(file_descriptor) - errored = _build_workbook( + errored = build_workbook( tmp_path, dashboard, active_data_mask, job_id, mode, user ) diff --git a/tests/integration_tests/dashboards/api_tests.py b/tests/integration_tests/dashboards/api_tests.py index ec1d02a82cf..1e16a8e6acc 100644 --- a/tests/integration_tests/dashboards/api_tests.py +++ b/tests/integration_tests/dashboards/api_tests.py @@ -28,11 +28,14 @@ import pytest import rison import yaml +from flask import current_app from freezegun import freeze_time from sqlalchemy import and_ from superset import db, security_manager # noqa: F401 from superset.commands.dashboard.permalink.create import CreateDashboardPermalinkCommand +from superset.daos.dashboard import EmbeddedDashboardDAO from superset.exceptions import LockAlreadyHeldException +from superset.security.guest_token import GuestTokenResourceType from superset.models.dashboard import Dashboard from superset.models.core import FavStar, FavStarClassName from superset.reports.models import ReportSchedule, ReportScheduleType @@ -3558,14 +3561,16 @@ class TestDashboardApi(ApiEditorsTestCaseMixin, InsertChartMixin, SupersetTestCa response = json.loads(rv.data.decode("utf-8")) assert response["count"] > 0 - def test_export_xlsx_501_when_bucket_unset(self): - """Dashboard API: export_xlsx returns 501 when the S3 bucket is unset.""" + def test_export_xlsx_400_for_empty_dashboard_without_storage(self): + """Dashboard API: with no storage configured the request is still validated + before an export runs, so a dashboard with no charts is rejected rather + than streaming an empty workbook.""" admin = self.get_user("admin") - dashboard = self.insert_dashboard("xlsx-501", None, [admin.id]) + dashboard = self.insert_dashboard("xlsx-sync-empty", None, [admin.id]) self.login(ADMIN_USERNAME) try: rv = self.client.post(f"api/v1/dashboard/{dashboard.id}/export_xlsx/") - assert rv.status_code == 501 + assert rv.status_code == 400 finally: db.session.delete(dashboard) db.session.commit() @@ -3725,6 +3730,250 @@ class TestDashboardApi(ApiEditorsTestCaseMixin, InsertChartMixin, SupersetTestCa _, kwargs = mock_task.apply_async.call_args assert kwargs["kwargs"]["mode"] == "images" + # --- Synchronous fallback (no export storage configured) ------------------ + + @staticmethod + def _write_stub_workbook(path, *args, **kwargs): + """Stand in for the shared workbook builder, writing a real .xlsx.""" + from superset.utils.excel_streaming import StreamingXlsxWriter + + writer = StreamingXlsxWriter(path) + writer.add_sheet("10 - Chart", ["a"], [[1]]) + writer.close() + return {} + + @staticmethod + def _export_temp_files(): + """Temp files the export path creates, so a leak can be detected.""" + import glob + import os + import tempfile + + return glob.glob(os.path.join(tempfile.gettempdir(), "dash-export-*")) + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.export_dashboard_excel") + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_200_streams_workbook_without_storage( + self, mock_build, mock_task + ): + """Dashboard API: with no storage configured the workbook is built inline + and returned as the response, instead of the request dead-ending.""" + mock_build.side_effect = self._write_stub_workbook + self.login(ADMIN_USERNAME) + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": {}}, + ) + + assert rv.status_code == 200 + assert rv.mimetype == ( + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + ) + assert "attachment" in rv.headers["Content-Disposition"] + assert ".xlsx" in rv.headers["Content-Disposition"] + # A real workbook (xlsx files are zip archives) reached the client. + assert rv.data.startswith(b"PK") + assert is_zipfile(BytesIO(rv.data)) + # Nothing was queued: no worker, no bucket, no email. + mock_task.apply_async.assert_not_called() + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_sync_builds_with_the_same_inputs_as_the_task(self, mock_build): + """Dashboard API: the synchronous path hands the shared builder the same + dashboard, filter state and mode the Celery task would, so both paths + produce the same workbook.""" + mock_build.side_effect = self._write_stub_workbook + data_mask = {"NATIVE_FILTER-abc": {"extraFormData": {"time_range": "No"}}} + self.login(ADMIN_USERNAME) + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": data_mask}, + ) + + assert rv.status_code == 200 + args, _ = mock_build.call_args + path, built_dashboard, active_data_mask, _job_id, mode, user = args + assert path.endswith(".xlsx") + assert built_dashboard.id == dashboard.id + assert active_data_mask == data_mask + assert mode == "data" + assert user.username == ADMIN_USERNAME + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.AcquireDistributedLock") + @patch("superset.dashboards.api.build_workbook") + @patch("superset.dashboards.api.is_within_sync_row_budget", return_value=False) + def test_export_xlsx_sync_refused_when_over_the_row_budget( + self, mock_budget, mock_build, mock_acquire + ): + """Dashboard API: an export too large to serve inline is refused up front + with a message naming the fix, rather than being started and timing out.""" + self.login(ADMIN_USERNAME) + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": {}}, + ) + + assert rv.status_code == 400 + message = rv.data.decode("utf-8") + assert "EXCEL_EXPORT_S3_BUCKET" in message + # The budget is consulted for the dashboard and mode being exported. + mock_budget.assert_called_once() + assert mock_budget.call_args.args[1] == "data" + # Refused before any work started, so no lock was taken and no rows read. + mock_build.assert_not_called() + mock_acquire.return_value.run.assert_not_called() + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.ReleaseDistributedLock") + @patch("superset.dashboards.api.AcquireDistributedLock") + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_sync_releases_the_lock_on_success( + self, mock_build, mock_acquire, mock_release + ): + """Dashboard API: the in-flight lock the synchronous path takes is released + once the response is ready, so the next export is not locked out.""" + mock_build.side_effect = self._write_stub_workbook + self.login(ADMIN_USERNAME) + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": {}}, + ) + + assert rv.status_code == 200 + mock_acquire.return_value.run.assert_called_once() + mock_release.return_value.run.assert_called_once() + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.ReleaseDistributedLock") + @patch("superset.dashboards.api.AcquireDistributedLock") + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_sync_releases_the_lock_when_building_fails( + self, mock_build, mock_acquire, mock_release + ): + """Dashboard API: a failure while building must not leave the user locked + out of their own dashboard until the lock's TTL expires.""" + mock_build.side_effect = RuntimeError("boom") + self.login(ADMIN_USERNAME) + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": {}}, + ) + + assert rv.status_code == 500 + mock_acquire.return_value.run.assert_called_once() + mock_release.return_value.run.assert_called_once() + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_sync_deletes_the_temp_file_on_success(self, mock_build): + """Dashboard API: the workbook is built through a temp file, which must not + outlive the response.""" + mock_build.side_effect = self._write_stub_workbook + before = self._export_temp_files() + self.login(ADMIN_USERNAME) + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": {}}, + ) + + assert rv.status_code == 200 + assert self._export_temp_files() == before + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_sync_deletes_the_temp_file_when_building_fails( + self, mock_build + ): + """Dashboard API: a half-written workbook is cleaned up too, so a failing + export does not fill the web server's disk.""" + mock_build.side_effect = RuntimeError("boom") + before = self._export_temp_files() + self.login(ADMIN_USERNAME) + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": {}}, + ) + + assert rv.status_code == 500 + assert self._export_temp_files() == before + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.AcquireDistributedLock") + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_sync_rejected_when_export_already_in_progress( + self, mock_build, mock_acquire + ): + """Dashboard API: the synchronous path honors the same per-user+dashboard + lock as the queued one, so one user cannot run two exports at once.""" + mock_acquire.return_value.run.side_effect = LockAlreadyHeldException("held") + self.login(ADMIN_USERNAME) + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": {}}, + ) + + assert rv.status_code == 202 + assert "already in progress" in rv.data.decode("utf-8") + mock_build.assert_not_called() + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_sync_still_blocks_guest_sessions(self, mock_build): + """Dashboard API: the synchronous path does not become a way for an + embedded guest session to export a dashboard. Guest support is deliberately + out of scope here, so a guest holding a token that *does* grant access to + this dashboard is still refused, before any workbook is built.""" + dashboard = db.session.query(Dashboard).filter_by(slug="world_health").first() + embedded = EmbeddedDashboardDAO.upsert(dashboard, ["superset.example"]) + db.session.commit() + token = security_manager.create_guest_access_token( + {"username": "xlsx_guest"}, + [{"type": GuestTokenResourceType.DASHBOARD, "id": str(embedded.uuid)}], + [], + ) + + rv = self.client.post( + f"api/v1/dashboard/{dashboard.id}/export_xlsx/", + json={"active_data_mask": {}}, + headers={ + current_app.config["GUEST_TOKEN_HEADER_NAME"]: token.decode("utf-8") + if isinstance(token, bytes) + else token + }, + ) + + assert rv.status_code == 400 + assert "email address" in rv.data.decode("utf-8") + mock_build.assert_not_called() + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") def test_embedded_dashboards(self): self.login(ADMIN_USERNAME) diff --git a/tests/unit_tests/dashboards/test_excel_export_storage.py b/tests/unit_tests/dashboards/test_excel_export_storage.py new file mode 100644 index 00000000000..45ea8393118 --- /dev/null +++ b/tests/unit_tests/dashboards/test_excel_export_storage.py @@ -0,0 +1,38 @@ +# 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 pytest + +from superset.dashboards.excel_export.storage import is_export_storage_configured + + [email protected]( + ("bucket", "configured"), + [ + ("exports-bucket", True), + (None, False), + ("", False), + ], +) +def test_storage_is_configured_only_with_a_bucket( + app: pytest.FixtureRequest, bucket: str | None, configured: bool +) -> None: + from flask import current_app + + current_app.config["EXCEL_EXPORT_S3_BUCKET"] = bucket + assert is_export_storage_configured() is configured diff --git a/tests/unit_tests/dashboards/test_excel_export_sync_budget.py b/tests/unit_tests/dashboards/test_excel_export_sync_budget.py new file mode 100644 index 00000000000..b51855c2704 --- /dev/null +++ b/tests/unit_tests/dashboards/test_excel_export_sync_budget.py @@ -0,0 +1,152 @@ +# 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 + +from collections.abc import Iterator +from typing import Any +from unittest import mock + +import pytest +from flask import current_app + +from superset.dashboards.excel_export.sync_budget import ( + is_within_sync_row_budget, + requested_row_total, +) +from superset.utils import json + +MODULE = "superset.dashboards.excel_export.sync_budget" + + +def _chart(chart_id: int, *queries: dict[str, Any]) -> mock.MagicMock: + """A chart whose saved query context holds ``queries``.""" + chart = mock.MagicMock() + chart.id = chart_id + chart.slice_name = f"Chart {chart_id}" + chart.viz_type = "table" + chart.query_context = json.dumps({"queries": list(queries)}) + return chart + + [email protected] +def charts() -> Iterator[mock.MagicMock]: + """Patch the layout walk so tests supply the dashboard's charts directly.""" + with mock.patch(f"{MODULE}.get_charts_in_layout_order") as ordered: + yield ordered + + [email protected](autouse=True) +def restore_config() -> Iterator[None]: + """Undo config edits: the app fixture is shared by every test in the module.""" + original = current_app.config["EXCEL_EXPORT_SYNC_MAX_ROWS"] + yield + current_app.config["EXCEL_EXPORT_SYNC_MAX_ROWS"] = original + + +def test_row_total_sums_the_row_limit_of_every_chart(charts: mock.MagicMock) -> None: + charts.return_value = [ + _chart(10, {"row_limit": 1000}), + _chart(20, {"row_limit": 250}), + ] + + assert requested_row_total(mock.MagicMock(), "data") == 1250 + + +def test_row_total_counts_every_query_of_a_multi_query_chart( + charts: mock.MagicMock, +) -> None: + # A mixed-series chart fans out to several queries, each of which becomes its + # own sheet and runs its own row_limit worth of rows. + charts.return_value = [_chart(10, {"row_limit": 100}, {"row_limit": 400})] + + assert requested_row_total(mock.MagicMock(), "data") == 500 + + [email protected]( + "query", + [ + {}, # no row_limit at all + {"row_limit": 0}, # 0 means "fall back to the configured limits" + {"row_limit": None}, + {"row_limit": "1000"}, # not an integer + {"row_limit": -5}, + ], +) +def test_row_total_is_indeterminate_without_a_finite_row_limit( + charts: mock.MagicMock, query: dict[str, Any] +) -> None: + # Without a finite limit on every query the export's size is unknown, so the + # budget cannot vouch for it and the caller must not run it inline. + charts.return_value = [_chart(10, {"row_limit": 100}), _chart(20, query)] + + assert requested_row_total(mock.MagicMock(), "data") is None + + +def test_row_total_ignores_charts_that_cannot_be_exported( + charts: mock.MagicMock, +) -> None: + # A chart with no usable query context is skipped by the export itself, so it + # runs no query and cannot contribute rows. + skipped = _chart(20) + skipped.query_context = None + skipped.viz_type = "mixed_timeseries" # outside the rebuild allowlist + charts.return_value = [_chart(10, {"row_limit": 100}), skipped] + + assert requested_row_total(mock.MagicMock(), "data") == 100 + + +def test_row_total_ignores_charts_rendered_as_images(charts: mock.MagicMock) -> None: + # In image mode a non-table chart is rendered through the webdriver instead of + # queried, so its row_limit is not part of the row budget. + rendered = _chart(20, {"row_limit": 999_999}) + rendered.viz_type = "pie" # not a table viz type, so it renders as an image + charts.return_value = [_chart(10, {"row_limit": 100}), rendered] + + assert requested_row_total(mock.MagicMock(), "images") == 100 + + [email protected]( + ("row_limit", "within"), + [ + (99_999, True), # below the limit + (100_000, True), # exactly at the limit + (100_001, False), # above the limit + ], +) +def test_budget_allows_totals_up_to_and_including_the_limit( + charts: mock.MagicMock, row_limit: int, within: bool +) -> None: + charts.return_value = [_chart(10, {"row_limit": row_limit})] + + assert is_within_sync_row_budget(mock.MagicMock(), "data") is within + + +def test_budget_refuses_an_indeterminate_total(charts: mock.MagicMock) -> None: + charts.return_value = [_chart(10, {})] + + assert is_within_sync_row_budget(mock.MagicMock(), "data") is False + + +def test_budget_honors_the_configured_limit(charts: mock.MagicMock) -> None: + charts.return_value = [_chart(10, {"row_limit": 5_000})] + current_app.config["EXCEL_EXPORT_SYNC_MAX_ROWS"] = 1_000 + + assert is_within_sync_row_budget(mock.MagicMock(), "data") is False + + current_app.config["EXCEL_EXPORT_SYNC_MAX_ROWS"] = 10_000 + + assert is_within_sync_row_budget(mock.MagicMock(), "data") is True diff --git a/tests/unit_tests/tasks/test_export_dashboard_excel.py b/tests/unit_tests/tasks/test_export_dashboard_excel.py index d75d5cdb60f..042f7f53868 100644 --- a/tests/unit_tests/tasks/test_export_dashboard_excel.py +++ b/tests/unit_tests/tasks/test_export_dashboard_excel.py @@ -30,6 +30,9 @@ from celery.exceptions import SoftTimeLimitExceeded from superset.utils import json MODULE = "superset.tasks.export_dashboard_excel" +# Workbook building lives in a module shared with the synchronous export path; +# the task only orchestrates storage upload and email delivery around it. +WORKBOOK_MODULE = "superset.dashboards.excel_export.workbook" # A minimal valid 1x1 transparent PNG for image-mode tests. @@ -71,21 +74,27 @@ def mocks() -> Iterator[dict[str, Any]]: # 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()) + 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", - "render_chart_image", - "s3", - "email", - "ReleaseDistributedLock", + for module, name in ( + (MODULE, "security_manager"), + (MODULE, "db"), + (MODULE, "s3"), + (MODULE, "ReleaseDistributedLock"), + (WORKBOOK_MODULE, "get_charts_in_layout_order"), + (WORKBOOK_MODULE, "get_dashboard_filter_context"), + (WORKBOOK_MODULE, "ChartDataQueryContextSchema"), + (WORKBOOK_MODULE, "ChartDataCommand"), + (WORKBOOK_MODULE, "render_chart_image"), ) } + # ``email`` is imported by both modules — the task sends through it, the + # workbook reads its ERROR_* reason keys — so both must see the same mock + # for the reason keys in an assertion to match the ones in the workbook. + shared_email = mock.MagicMock() + for module in (MODULE, WORKBOOK_MODULE): + stack.enter_context(mock.patch(f"{module}.email", new=shared_email)) + patched["email"] = shared_email user = mock.MagicMock() user.email = "[email protected]" patched["security_manager"].get_user_by_id.return_value = user @@ -287,8 +296,6 @@ def _rebuildable_chart( @contextmanager def _builder_hook(builder: Any) -> Iterator[None]: """Patch current_app so EXCEL_EXPORT_QUERY_CONTEXT_BUILDER resolves to builder.""" - from superset.tasks import export_dashboard_excel as module - fake_app = mock.MagicMock() fake_app.config.get.side_effect = lambda key, default=None: ( builder if key == "EXCEL_EXPORT_QUERY_CONTEXT_BUILDER" else default @@ -300,14 +307,18 @@ def _builder_hook(builder: Any) -> Iterator[None]: "EXCEL_EXPORT_S3_KEY_PREFIX": "dashboard-exports/", "EXCEL_EXPORT_LINK_TTL_SECONDS": 3600, }.__getitem__ - with mock.patch.object(module, "current_app", fake_app): + # Both modules read config: the workbook resolves the builder hook and the + # table-viz overrides, the task reads the storage keys. + with ExitStack() as stack: + for module in (MODULE, WORKBOOK_MODULE): + stack.enter_context(mock.patch(f"{module}.current_app", fake_app)) yield def test_builder_hook_context_is_used_for_any_viz_type() -> None: # A configured builder can supply a context for a viz type outside the # built-in allowlist (pivot_table_v2), and is called with the chart's form data. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module ctx = { "datasource": {"id": 5, "type": "table"}, @@ -317,7 +328,7 @@ def test_builder_hook_context_is_used_for_any_viz_type() -> None: chart = _rebuildable_chart(viz_type="pivot_table_v2") with _builder_hook(builder): - result = module._resolve_query_context(chart) + result = module.resolve_query_context(chart) assert result == ctx builder.assert_called_once_with(chart.form_data) @@ -326,13 +337,13 @@ def test_builder_hook_context_is_used_for_any_viz_type() -> None: def test_builder_hook_none_falls_through_to_builtin_rebuild() -> None: # When the builder returns None (can't build faithfully) the export falls # through to the built-in rebuild, so an allowlisted table is unaffected. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module builder = mock.MagicMock(return_value=None) chart = _rebuildable_chart(viz_type="table") with _builder_hook(builder): - result = module._resolve_query_context(chart) + result = module.resolve_query_context(chart) builder.assert_called_once_with(chart.form_data) assert result is not None @@ -354,13 +365,13 @@ def test_builder_hook_none_falls_through_to_builtin_rebuild() -> None: def test_builder_hook_malformed_result_falls_through(built: Any) -> None: # A stub / empty / malformed builder result is treated as "not built" and # falls through to the built-in rebuild rather than shipping an empty context. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module builder = mock.MagicMock(return_value=built) chart = _rebuildable_chart(viz_type="table") with _builder_hook(builder): - result = module._resolve_query_context(chart) + result = module.resolve_query_context(chart) builder.assert_called_once_with(chart.form_data) assert result is not None @@ -370,13 +381,13 @@ def test_builder_hook_malformed_result_falls_through(built: Any) -> None: def test_builder_hook_exception_falls_through() -> None: # A raising builder (e.g. sidecar down) must not fail the chart; the export # falls through to the built-in rebuild and no exception escapes. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module builder = mock.MagicMock(side_effect=RuntimeError("sidecar down")) chart = _rebuildable_chart(viz_type="table") with _builder_hook(builder): - result = module._resolve_query_context(chart) + result = module.resolve_query_context(chart) builder.assert_called_once_with(chart.form_data) assert result is not None @@ -387,13 +398,13 @@ def test_builder_hook_soft_time_limit_propagates() -> None: # A soft timeout raised while the builder is in flight is a task-level signal, # not a builder failure: it must escape _resolve_query_context so the export # aborts cleanly, rather than being swallowed by the broad fall-through guard. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module builder = mock.MagicMock(side_effect=SoftTimeLimitExceeded()) chart = _rebuildable_chart(viz_type="table") with _builder_hook(builder), pytest.raises(SoftTimeLimitExceeded): - module._resolve_query_context(chart) + module.resolve_query_context(chart) builder.assert_called_once_with(chart.form_data) @@ -401,11 +412,11 @@ def test_builder_hook_soft_time_limit_propagates() -> None: def test_no_builder_hook_leaves_builtin_behavior_unchanged() -> None: # With no builder configured, an allowlisted chart is rebuilt and an # ineligible one is skipped — identical to the pre-hook behavior. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module with _builder_hook(None): - table = module._resolve_query_context(_rebuildable_chart(viz_type="table")) - ineligible = module._resolve_query_context( + table = module.resolve_query_context(_rebuildable_chart(viz_type="table")) + ineligible = module.resolve_query_context( _rebuildable_chart(viz_type="mixed_timeseries") ) @@ -416,14 +427,14 @@ def test_no_builder_hook_leaves_builtin_behavior_unchanged() -> None: def test_saved_context_short_circuits_builder_hook() -> None: # A saved query context wins over the builder hook, which is never called. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module builder = mock.MagicMock(return_value={"queries": [{"from": "hook"}]}) chart = _rebuildable_chart(viz_type="table") chart.query_context = json.dumps({"queries": [{"from": "saved"}]}) with _builder_hook(builder): - result = module._resolve_query_context(chart) + result = module.resolve_query_context(chart) assert result == {"queries": [{"from": "saved"}]} builder.assert_not_called() @@ -595,7 +606,7 @@ def test_raw_mode_table_ignores_stale_show_totals() -> None: # gates the totals query on queryMode === Aggregate), and the control isn't # reset when hidden. A raw-mode table carrying a stale value must still # rebuild rather than be needlessly skipped. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module chart = _rebuildable_chart( viz_type="table", @@ -603,7 +614,7 @@ def test_raw_mode_table_ignores_stale_show_totals() -> None: ) with _builder_hook(None): - result = module._resolve_query_context(chart) + result = module.resolve_query_context(chart) assert result is not None assert result["queries"][0]["columns"] == ["a"] @@ -612,7 +623,7 @@ def test_raw_mode_table_ignores_stale_show_totals() -> None: def test_rebuild_viz_types_is_the_conservative_default() -> None: # The rebuild allow-list is a fixed fallback (no config override): only viz # types whose data maps faithfully to a single plain query. - from superset.tasks import export_dashboard_excel as module + from superset.dashboards.excel_export import workbook as module assert module.REBUILD_VIZ_TYPES == { "table",
