feat: add global max row limit (#16683)

* feat: add global max limit

* fix lint and tests

* leave SAMPLES_ROW_LIMIT unchanged

* fix sample rowcount test

* replace max global limit with existing sql max row limit

* fix test

* make max_limit optional in util

* improve comments

tests/integration_tests/charts/api_tests.py

@@ -18,7 +18,7 @@
 """Unit tests for Superset"""
 import json
 import unittest
-from datetime import datetime, timedelta
+from datetime import datetime
 from io import BytesIO
 from typing import Optional
 from unittest import mock
@@ -1203,6 +1203,7 @@ class TestChartApi(SupersetTestCase, ApiOwnersTestCaseMixin, InsertChartMixin):
         self.login(username="admin")
         request_payload = get_query_context("birth_names")
         del request_payload["queries"][0]["row_limit"]
+
         rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
         response_payload = json.loads(rv.data.decode("utf-8"))
         result = response_payload["result"][0]
@@ -1210,11 +1211,46 @@ class TestChartApi(SupersetTestCase, ApiOwnersTestCaseMixin, InsertChartMixin):
 
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     @mock.patch(
-        "superset.common.query_actions.config", {**app.config, "SAMPLES_ROW_LIMIT": 5},
+        "superset.utils.core.current_app.config", {**app.config, "SQL_MAX_ROW": 10},
     )
-    def test_chart_data_default_sample_limit(self):
+    def test_chart_data_sql_max_row_limit(self):
         """
-        Chart data API: Ensure sample response row count doesn't exceed default limit
+        Chart data API: Ensure row count doesn't exceed max global row limit
+        """
+        self.login(username="admin")
+        request_payload = get_query_context("birth_names")
+        request_payload["queries"][0]["row_limit"] = 10000000
+        rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
+        response_payload = json.loads(rv.data.decode("utf-8"))
+        result = response_payload["result"][0]
+        self.assertEqual(result["rowcount"], 10)
+
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    @mock.patch(
+        "superset.common.query_object.config", {**app.config, "SAMPLES_ROW_LIMIT": 5},
+    )
+    def test_chart_data_sample_default_limit(self):
+        """
+        Chart data API: Ensure sample response row count defaults to config defaults
+        """
+        self.login(username="admin")
+        request_payload = get_query_context("birth_names")
+        request_payload["result_type"] = utils.ChartDataResultType.SAMPLES
+        del request_payload["queries"][0]["row_limit"]
+        rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
+        response_payload = json.loads(rv.data.decode("utf-8"))
+        result = response_payload["result"][0]
+        self.assertEqual(result["rowcount"], 5)
+
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    @mock.patch(
+        "superset.common.query_actions.config",
+        {**app.config, "SAMPLES_ROW_LIMIT": 5, "SQL_MAX_ROW": 15},
+    )
+    def test_chart_data_sample_custom_limit(self):
+        """
+        Chart data API: Ensure requested sample response row count is between
+        default and SQL max row limit
         """
         self.login(username="admin")
         request_payload = get_query_context("birth_names")
@@ -1223,6 +1259,24 @@ class TestChartApi(SupersetTestCase, ApiOwnersTestCaseMixin, InsertChartMixin):
         rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
         response_payload = json.loads(rv.data.decode("utf-8"))
         result = response_payload["result"][0]
+        self.assertEqual(result["rowcount"], 10)
+
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    @mock.patch(
+        "superset.utils.core.current_app.config", {**app.config, "SQL_MAX_ROW": 5},
+    )
+    def test_chart_data_sql_max_row_sample_limit(self):
+        """
+        Chart data API: Ensure requested sample response row count doesn't
+        exceed SQL max row limit
+        """
+        self.login(username="admin")
+        request_payload = get_query_context("birth_names")
+        request_payload["result_type"] = utils.ChartDataResultType.SAMPLES
+        request_payload["queries"][0]["row_limit"] = 10000000
+        rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
+        response_payload = json.loads(rv.data.decode("utf-8"))
+        result = response_payload["result"][0]
         self.assertEqual(result["rowcount"], 5)
 
     def test_chart_data_incorrect_result_type(self):

tests/integration_tests/charts/schema_tests.py

@@ -16,17 +16,25 @@
 # under the License.
 # isort:skip_file
 """Unit tests for Superset"""
-from typing import Any, Dict, Tuple
+from unittest import mock
+
+import pytest
 
 from marshmallow import ValidationError
 from tests.integration_tests.test_app import app
 from superset.charts.schemas import ChartDataQueryContextSchema
-from superset.common.query_context import QueryContext
 from tests.integration_tests.base_tests import SupersetTestCase
+from tests.integration_tests.fixtures.birth_names_dashboard import (
+    load_birth_names_dashboard_with_slices,
+)
 from tests.integration_tests.fixtures.query_context import get_query_context
 
 
 class TestSchema(SupersetTestCase):
+    @mock.patch(
+        "superset.common.query_object.config", {**app.config, "ROW_LIMIT": 5000},
+    )
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_query_context_limit_and_offset(self):
         self.login(username="admin")
         payload = get_query_context("birth_names")
@@ -36,7 +44,7 @@ class TestSchema(SupersetTestCase):
         payload["queries"][0].pop("row_offset", None)
         query_context = ChartDataQueryContextSchema().load(payload)
         query_object = query_context.queries[0]
-        self.assertEqual(query_object.row_limit, app.config["ROW_LIMIT"])
+        self.assertEqual(query_object.row_limit, 5000)
         self.assertEqual(query_object.row_offset, 0)
 
         # Valid limit and offset
@@ -55,12 +63,14 @@ class TestSchema(SupersetTestCase):
         self.assertIn("row_limit", context.exception.messages["queries"][0])
         self.assertIn("row_offset", context.exception.messages["queries"][0])
 
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_query_context_null_timegrain(self):
         self.login(username="admin")
         payload = get_query_context("birth_names")
         payload["queries"][0]["extras"]["time_grain_sqla"] = None
         _ = ChartDataQueryContextSchema().load(payload)
 
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_query_context_series_limit(self):
         self.login(username="admin")
         payload = get_query_context("birth_names")
@@ -82,6 +92,7 @@ class TestSchema(SupersetTestCase):
         }
         _ = ChartDataQueryContextSchema().load(payload)
 
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_query_context_null_post_processing_op(self):
         self.login(username="admin")
         payload = get_query_context("birth_names")

tests/integration_tests/query_context_tests.py

@@ -90,6 +90,7 @@ class TestQueryContext(SupersetTestCase):
                 self.assertEqual(post_proc["operation"], payload_post_proc["operation"])
                 self.assertEqual(post_proc["options"], payload_post_proc["options"])
 
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_cache(self):
         table_name = "birth_names"
         table = self.get_table(name=table_name)