sadpandajoe commented on code in PR #44082: URL: https://github.com/apache/superset/pull/44082#discussion_r3970407823
########## superset/dashboards/excel_export/sync_budget.py: ########## @@ -0,0 +1,116 @@ +# 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. +""" +Plan an export that has to be served as the response to one request. + +An inline export has to finish inside its request, so its size is settled +*before* any query runs, by adding up the rows it 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. + +Working that total out means resolving each chart's query context, which is what +the export itself runs. The plan therefore hands those contexts back and the +export reuses them, so the queries that run are exactly the ones the budget was +measured against — resolution can be expensive and, through +``EXCEL_EXPORT_QUERY_CONTEXT_BUILDER``, is not guaranteed to be deterministic. +""" + +from __future__ import annotations + +from dataclasses import dataclass +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 ( + resolve_query_context, + ResolvedQueryContexts, +) + + +@dataclass(frozen=True) +class InlineExportPlan: + """What an inline export would run, and whether it is small enough to.""" + + #: Every chart's resolved query context, keyed by chart id. A ``None`` value + #: is an answer, not a gap: that chart cannot be exported and will be listed + #: as skipped. + query_contexts: ResolvedQueryContexts + #: Rows every query is allowed to return, or ``None`` when any query has no + #: finite limit and the size of the export is therefore unknowable. + requested_rows: int | None + #: The configured ceiling this plan was measured against. + max_rows: int + + @property + def fits_row_budget(self) -> bool: + """Whether this export may run inline, as the response to one request.""" + return self.requested_rows is not None and self.requested_rows <= self.max_rows + + +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): Review Comment: A saved or rebuilt Big Number context has no `row_limit`, so this treats it as unbounded and the new no-storage path rejects the entire dashboard before it runs. The query factory supplies the configured `ROW_LIMIT` for that case, and Big Number has no row-limit control to lower. Should the planner budget the same configured default instead of returning 400? ########## tests/integration_tests/dashboards/api_tests.py: ########## @@ -3725,6 +3735,331 @@ def test_export_xlsx_images_202_when_screenshot_flags_on( _, 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.ReleaseDistributedLock") + @patch("superset.dashboards.api.AcquireDistributedLock") + @patch("superset.dashboards.api.build_workbook") + @patch("superset.dashboards.api.plan_inline_export") + def test_export_xlsx_sync_refused_when_over_the_row_budget( + self, mock_plan, mock_build, mock_acquire, mock_release + ): + """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.""" + mock_plan.return_value = InlineExportPlan( + query_contexts={}, requested_rows=250_000, max_rows=100_000 + ) + 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 lock prevents a duplicate request from paying the planning cost; + # a refusal releases it immediately and never reads chart rows. + mock_plan.assert_called_once() + mock_build.assert_not_called() + 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.plan_inline_export") + def test_export_xlsx_sync_releases_the_lock_when_planning_fails( + self, mock_plan, mock_acquire, mock_release + ): + """Dashboard API: a context-builder failure while planning must not keep + the user locked out until the lock's TTL expires.""" + mock_plan.side_effect = RuntimeError("builder failed") + 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") + @patch("superset.dashboards.api.plan_inline_export") + def test_export_xlsx_sync_runs_the_contexts_the_budget_measured( + self, mock_plan, mock_build + ): + """Dashboard API: the export runs the query contexts the row budget was + measured against. Resolving them a second time would risk vouching for one + set of queries and running another, since a deployment's context builder + need not be deterministic.""" + measured = {10: {"queries": [{"row_limit": 5}]}, 20: None} + mock_plan.return_value = InlineExportPlan( + query_contexts=measured, requested_rows=5, max_rows=100_000 + ) + 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 mock_build.call_args.kwargs["query_contexts"] is measured + + @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices") + @with_config({"EXCEL_EXPORT_S3_BUCKET": None}) + @with_feature_flags( + ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True, + ENABLE_DASHBOARD_DOWNLOAD_WEBDRIVER_SCREENSHOT=True, + ) + @patch("superset.dashboards.api.AcquireDistributedLock") + @patch("superset.dashboards.api.build_workbook") + def test_export_xlsx_images_refused_without_storage(self, mock_build, mock_acquire): + """Dashboard API: image export renders every chart through the headless + webdriver, which no row budget bounds and no request should wait on, so it + is refused rather than served inline -- even with the webdriver enabled.""" + 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": {}, "mode": "images"}, + ) + + assert rv.status_code == 400 + assert "EXCEL_EXPORT_S3_BUCKET" in rv.data.decode("utf-8") + 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") + @patch("superset.dashboards.api.plan_inline_export") + def test_export_xlsx_sync_rejected_when_export_already_in_progress( + self, mock_plan, 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_plan.assert_not_called() + 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( Review Comment: This test never enables `EMBEDDED_SUPERSET`, so the guest-token request loader ignores the header and the assertion can only exercise an anonymous/no-email rejection (or an earlier auth failure). Could this enable the embedded flag and assert a real `GuestUser` request, so removing the guest guard fails the regression test? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
