diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.test.tsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.test.tsx index 8d5c7679697..06d154d0808 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.test.tsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.test.tsx @@ -229,6 +229,7 @@ test('should call exportChart with row_limit TABLE_VIZ_MAX_ROW_SERVER when expor expect.objectContaining({ formData: expect.objectContaining({ row_limit: 999, + full_export: true, dashboardId: 111, }), resultType: 'full', @@ -284,6 +285,7 @@ test('should call exportChart with row_limit TABLE_VIZ_MAX_ROW_SERVER when expor expect.objectContaining({ formData: expect.objectContaining({ row_limit: 999, + full_export: true, dashboardId: 111, }), resultType: 'full', diff --git a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.tsx b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.tsx index ebd28994123..017c2f22c30 100644 --- a/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.tsx +++ b/superset-frontend/src/dashboard/components/gridComponents/Chart/Chart.tsx @@ -482,7 +482,7 @@ const Chart = (props: ChartProps) => { (formData as JsonObject).dashboardId = dashboardInfo.id; const exportTable = useCallback( - (format: string, isFullCSV: boolean, isPivot = false) => { + (format: string, isFullExport: boolean, isPivot = false) => { const logAction = format === 'csv' ? LOG_ACTIONS_EXPORT_CSV_DASHBOARD_CHART @@ -492,8 +492,11 @@ const Chart = (props: ChartProps) => { is_cached: isCached, }); - const exportFormData = isFullCSV - ? { ...formData, row_limit: fullExportMaxRows } + // For a "full" export, raise the requested row_limit and flag the + // request with full_export so the backend lifts the row-limit cap to + // TABLE_VIZ_MAX_ROW_SERVER (gated by the ALLOW_FULL_CSV_EXPORT flag). + const exportFormData = isFullExport + ? { ...formData, row_limit: fullExportMaxRows, full_export: true } : formData; const resultType = isPivot ? 'post_processed' : 'full'; diff --git a/superset/commands/chart/data/streaming_export_command.py b/superset/commands/chart/data/streaming_export_command.py index 9284eafa055..2e30c42a440 100644 --- a/superset/commands/chart/data/streaming_export_command.py +++ b/superset/commands/chart/data/streaming_export_command.py @@ -20,9 +20,6 @@ from __future__ import annotations from typing import Any, TYPE_CHECKING -from flask import current_app as app - -from superset import is_feature_enabled from superset.commands.streaming_export.base import BaseStreamingCSVExportCommand if TYPE_CHECKING: @@ -71,13 +68,6 @@ class StreamingCSVExportCommand(BaseStreamingCSVExportCommand): query_obj = self._query_context.queries[0] query_dict = query_obj.to_dict() - # When ALLOW_FULL_CSV_EXPORT is enabled, raise the row limit so a - # "full" export is not silently capped at SQL_MAX_ROW. The ceiling is - # TABLE_VIZ_MAX_ROW_SERVER (a bounded, predictable maximum) rather than - # truly unlimited. - if is_feature_enabled("ALLOW_FULL_CSV_EXPORT"): - query_dict["row_limit"] = app.config["TABLE_VIZ_MAX_ROW_SERVER"] - # Use get_query_str_extended (single, clean statement) instead of # get_query_str, which returns a multi-statement string (prequeries + # main SQL joined by ";" with a trailing ";"). The base command runs the diff --git a/superset/common/query_context_factory.py b/superset/common/query_context_factory.py index 0d1539a9a16..b475063f54a 100644 --- a/superset/common/query_context_factory.py +++ b/superset/common/query_context_factory.py @@ -74,6 +74,11 @@ class QueryContextFactory: # pylint: disable=too-few-public-methods bool(form_data.get("server_pagination")) if form_data else False ) + # A "full" CSV/Excel export raises the row-limit ceiling to + # TABLE_VIZ_MAX_ROW_SERVER (when ALLOW_FULL_CSV_EXPORT is enabled). + # The marker is set by the frontend's "Export to full ..." actions. + full_export = bool(form_data.get("full_export")) if form_data else False + queries_ = [ self._process_query_object( datasource_model_instance, @@ -82,6 +87,7 @@ class QueryContextFactory: # pylint: disable=too-few-public-methods result_type, datasource=datasource, server_pagination=server_pagination, + full_export=full_export, **query_obj, ), ) diff --git a/superset/common/query_object_factory.py b/superset/common/query_object_factory.py index 31d2843da56..736e25d43db 100644 --- a/superset/common/query_object_factory.py +++ b/superset/common/query_object_factory.py @@ -58,6 +58,7 @@ class QueryObjectFactory: # pylint: disable=too-few-public-methods time_range: str | None = None, time_shift: str | None = None, server_pagination: bool | None = None, + full_export: bool | None = None, **kwargs: Any, ) -> QueryObject: datasource_model_instance = None @@ -66,9 +67,12 @@ class QueryObjectFactory: # pylint: disable=too-few-public-methods processed_extras = self._process_extras(extras) result_type = kwargs.setdefault("result_type", parent_result_type) - # Process row limit taking server pagination into account + # Process row limit taking server pagination and full export into account row_limit = self._process_row_limit( - row_limit, result_type, server_pagination=server_pagination + row_limit, + result_type, + server_pagination=server_pagination, + full_export=full_export, ) processed_time_range = self._process_time_range( @@ -106,12 +110,14 @@ class QueryObjectFactory: # pylint: disable=too-few-public-methods row_limit: int | None, result_type: ChartDataResultType, server_pagination: bool | None = None, + full_export: bool | None = None, ) -> int: """Process row limit taking into account server pagination. :param row_limit: The requested row limit :param result_type: The type of result being processed :param server_pagination: Whether server-side pagination is enabled + :param full_export: Whether this is a "full" CSV/Excel export request :return: The processed row limit """ default_row_limit = ( @@ -122,6 +128,7 @@ class QueryObjectFactory: # pylint: disable=too-few-public-methods return apply_max_row_limit( row_limit or default_row_limit, server_pagination=server_pagination, + full_export=full_export, ) @staticmethod diff --git a/superset/config.py b/superset/config.py index 2a29ce87147..a8d8741d5d3 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1317,7 +1317,9 @@ MAPBOX_API_KEY = os.environ.get("MAPBOX_API_KEY", "") # Maximum number of rows returned for any analytical database query SQL_MAX_ROW = 100000 -# Maximum number of rows for any query with Server Pagination in Table Viz type +# Maximum number of rows for any query with Server Pagination in Table Viz type. +# This also serves as the row-count ceiling for "full" CSV/Excel exports when the +# ALLOW_FULL_CSV_EXPORT feature flag is enabled (see apply_max_row_limit). TABLE_VIZ_MAX_ROW_SERVER = 500000 diff --git a/superset/utils/core.py b/superset/utils/core.py index f361034e10b..59328b84edd 100644 --- a/superset/utils/core.py +++ b/superset/utils/core.py @@ -2051,13 +2051,17 @@ def parse_boolean_string(bool_str: str | None) -> bool: def apply_max_row_limit( limit: int, server_pagination: bool | None = None, + full_export: bool | None = None, ) -> int: """ - Override row limit based on server pagination setting + Override row limit based on server pagination / full-export settings :param limit: requested row limit :param server_pagination: whether server-side pagination is enabled, defaults to None + :param full_export: whether this is a "full" CSV/Excel export request, + which raises the ceiling to TABLE_VIZ_MAX_ROW_SERVER when the + ALLOW_FULL_CSV_EXPORT feature flag is enabled, defaults to None :return: Capped row limit >>> apply_max_row_limit(600000, server_pagination=True) # Server pagination @@ -2069,13 +2073,25 @@ def apply_max_row_limit( >>> apply_max_row_limit(0) # Zero returns default max limit 50000 """ + # Imported locally to avoid a circular import: superset.extensions pulls in + # superset.security.manager / superset.utils.cache_manager, both of which + # import superset.utils.core. # pylint: disable=import-outside-toplevel + from superset.extensions import feature_flag_manager - max_limit = ( - app.config["TABLE_VIZ_MAX_ROW_SERVER"] - if server_pagination - else app.config["SQL_MAX_ROW"] + # A "full" CSV/Excel export is allowed past the regular SQL_MAX_ROW cap, but + # only when the operator has opted in via the ALLOW_FULL_CSV_EXPORT flag. + # server_pagination is a separate, independent reason to raise the cap and is + # NOT gated by that flag. + allow_full_export = full_export and feature_flag_manager.is_feature_enabled( + "ALLOW_FULL_CSV_EXPORT" ) + # Both raised cases share the same ceiling, TABLE_VIZ_MAX_ROW_SERVER (a + # bounded, predictable maximum); see its definition in config.py. + if server_pagination or allow_full_export: + max_limit = app.config["TABLE_VIZ_MAX_ROW_SERVER"] + else: + max_limit = app.config["SQL_MAX_ROW"] if limit != 0: return min(max_limit, limit) return max_limit diff --git a/tests/unit_tests/commands/chart/streaming_export_command_test.py b/tests/unit_tests/commands/chart/streaming_export_command_test.py index 3fca88a2125..64dd641dc70 100644 --- a/tests/unit_tests/commands/chart/streaming_export_command_test.py +++ b/tests/unit_tests/commands/chart/streaming_export_command_test.py @@ -17,13 +17,11 @@ """Unit tests for Chart Streaming CSV Export Command.""" import pytest -from flask import current_app as app from pytest_mock import MockerFixture from superset.commands.chart.data.streaming_export_command import ( StreamingCSVExportCommand, ) -from tests.unit_tests.conftest import with_feature_flags def _setup_chart_mocks( @@ -311,30 +309,6 @@ def test_catalog_and_schema_passed_to_engine(mocker: MockerFixture) -> None: ) -@with_feature_flags(ALLOW_FULL_CSV_EXPORT=True) -def test_full_csv_export_raises_row_limit(mocker: MockerFixture) -> None: - """ALLOW_FULL_CSV_EXPORT raises the row limit to TABLE_VIZ_MAX_ROW_SERVER.""" - _, query_context, datasource = _setup_chart_mocks(mocker) - - command = StreamingCSVExportCommand(query_context) - command._get_sql_and_database() - - query_dict = datasource.get_query_str_extended.call_args[0][0] - assert query_dict["row_limit"] == app.config["TABLE_VIZ_MAX_ROW_SERVER"] - - -@with_feature_flags(ALLOW_FULL_CSV_EXPORT=False) -def test_row_limit_unchanged_without_flag(mocker: MockerFixture) -> None: - """Without the flag, the chart's own row limit is left untouched.""" - _, query_context, datasource = _setup_chart_mocks(mocker) - - command = StreamingCSVExportCommand(query_context) - command._get_sql_and_database() - - query_dict = datasource.get_query_str_extended.call_args[0][0] - assert query_dict["row_limit"] == 100 - - def test_uses_extended_sql_single_statement(mocker: MockerFixture) -> None: """SQL generation uses get_query_str_extended (no prequeries, no trailing ;). diff --git a/tests/unit_tests/common/test_query_object_factory.py b/tests/unit_tests/common/test_query_object_factory.py index 379137a34c4..8e72c541e7a 100644 --- a/tests/unit_tests/common/test_query_object_factory.py +++ b/tests/unit_tests/common/test_query_object_factory.py @@ -30,6 +30,7 @@ def create_app_config() -> dict[str, Any]: "DEFAULT_RELATIVE_END_TIME": "today", "SAMPLES_ROW_LIMIT": 1000, "SQL_MAX_ROW": 100000, + "TABLE_VIZ_MAX_ROW_SERVER": 500000, } @@ -48,12 +49,12 @@ def connector_registry() -> Mock: def apply_max_row_limit( limit: int, server_pagination: bool | None = None, + full_export: bool | None = None, ) -> int: - max_limit = ( - create_app_config()["TABLE_VIZ_MAX_ROW_SERVER"] - if server_pagination - else create_app_config()["SQL_MAX_ROW"] - ) + if server_pagination or full_export: + max_limit = create_app_config()["TABLE_VIZ_MAX_ROW_SERVER"] + else: + max_limit = create_app_config()["SQL_MAX_ROW"] if limit != 0: return min(max_limit, limit) return max_limit @@ -109,6 +110,36 @@ class TestQueryObjectFactory: assert query_object.row_limit == 100 assert query_object.row_offset == 200 + def test_query_context_full_export_raises_limit( + self, + query_object_factory: QueryObjectFactory, + raw_query_context: dict[str, Any], + ): + """full_export raises the row-limit ceiling to TABLE_VIZ_MAX_ROW_SERVER.""" + raw_query_object = raw_query_context["queries"][0] + raw_query_object["row_limit"] = 300000 + query_object = query_object_factory.create( + raw_query_context["result_type"], + full_export=True, + **raw_query_object, + ) + # Without full_export this would be capped at SQL_MAX_ROW (100000). + assert query_object.row_limit == 300000 + + def test_query_context_limit_capped_without_full_export( + self, + query_object_factory: QueryObjectFactory, + raw_query_context: dict[str, Any], + ): + """A regular request stays capped at SQL_MAX_ROW.""" + raw_query_object = raw_query_context["queries"][0] + raw_query_object["row_limit"] = 300000 + query_object = query_object_factory.create( + raw_query_context["result_type"], + **raw_query_object, + ) + assert query_object.row_limit == 100000 + def test_query_context_null_post_processing_op( self, query_object_factory: QueryObjectFactory, diff --git a/tests/unit_tests/utils/test_core.py b/tests/unit_tests/utils/test_core.py index 5e663dd095a..c01b8804769 100644 --- a/tests/unit_tests/utils/test_core.py +++ b/tests/unit_tests/utils/test_core.py @@ -28,6 +28,7 @@ from pytest_mock import MockerFixture from superset.exceptions import SupersetException from superset.utils.core import ( + apply_max_row_limit, cast_to_boolean, check_is_safe_zip, DateColumn, @@ -53,6 +54,7 @@ from superset.utils.core import ( sanitize_url, ) from tests.conftest import with_config +from tests.unit_tests.conftest import with_feature_flags ADHOC_FILTER: QueryObjectFilterClause = { "col": "foo", @@ -1730,3 +1732,34 @@ def test_markdown_with_markup_wrap() -> None: assert isinstance(result, Markup) assert "bold" in str(result) + + +@with_config({"SQL_MAX_ROW": 100000, "TABLE_VIZ_MAX_ROW_SERVER": 500000}) +def test_apply_max_row_limit_default_cap() -> None: + """A regular request is capped at SQL_MAX_ROW.""" + assert apply_max_row_limit(300000) == 100000 + assert apply_max_row_limit(5000) == 5000 + # 0 means "no explicit limit" -> default max + assert apply_max_row_limit(0) == 100000 + + +@with_config({"SQL_MAX_ROW": 100000, "TABLE_VIZ_MAX_ROW_SERVER": 500000}) +def test_apply_max_row_limit_server_pagination() -> None: + """server_pagination raises the cap to TABLE_VIZ_MAX_ROW_SERVER.""" + assert apply_max_row_limit(300000, server_pagination=True) == 300000 + assert apply_max_row_limit(900000, server_pagination=True) == 500000 + + +@with_config({"SQL_MAX_ROW": 100000, "TABLE_VIZ_MAX_ROW_SERVER": 500000}) +@with_feature_flags(ALLOW_FULL_CSV_EXPORT=True) +def test_apply_max_row_limit_full_export_with_flag() -> None: + """full_export raises the cap to TABLE_VIZ_MAX_ROW_SERVER when the flag is on.""" + assert apply_max_row_limit(300000, full_export=True) == 300000 + assert apply_max_row_limit(900000, full_export=True) == 500000 + + +@with_config({"SQL_MAX_ROW": 100000, "TABLE_VIZ_MAX_ROW_SERVER": 500000}) +@with_feature_flags(ALLOW_FULL_CSV_EXPORT=False) +def test_apply_max_row_limit_full_export_without_flag() -> None: + """full_export has no effect when ALLOW_FULL_CSV_EXPORT is disabled.""" + assert apply_max_row_limit(300000, full_export=True) == 100000