From 9b00d1a4b201803aaaa296e2e235570997e07fec Mon Sep 17 00:00:00 2001
From: ofekisr <ofek1israel@gmail.com>
Date: Wed, 22 Sep 2021 19:44:05 +0300
Subject: [PATCH] refactor move apply_display_max_row_limit to sqllab package

---
 superset/sqllab/command.py |  4 ++--
 superset/sqllab/utils.py   | 47 ++++++++++++++++++++++++++++++++++++++
 superset/views/core.py     |  4 ++--
 3 files changed, 51 insertions(+), 4 deletions(-)
 create mode 100644 superset/sqllab/utils.py

diff --git a/superset/sqllab/command.py b/superset/sqllab/command.py
index 272b2f2740a9a..a377d1c883fc2 100644
--- a/superset/sqllab/command.py
+++ b/superset/sqllab/command.py
@@ -50,7 +50,7 @@
 from superset.utils import core as utils
 from superset.utils.dates import now_as_float
 from superset.utils.sqllab_execution_context import SqlJsonExecutionContext
-from superset.views.utils import apply_display_max_row_limit
+from superset.sqllab.utils import apply_display_max_row_configuration_if_require
 
 config = app.config
 logger = logging.getLogger(__name__)
@@ -397,7 +397,7 @@ def _to_payload_results_based(  # pylint: disable=no-self-use
     ) -> str:
         display_max_row = config["DISPLAY_MAX_ROW"]
         return json.dumps(
-            apply_display_max_row_limit(execution_result, display_max_row),
+            apply_display_max_row_configuration_if_require(execution_result, display_max_row),
             default=utils.pessimistic_json_iso_dttm_ser,
             ignore_nan=True,
             encoding=None,
diff --git a/superset/sqllab/utils.py b/superset/sqllab/utils.py
new file mode 100644
index 0000000000000..4259ac5c5b069
--- /dev/null
+++ b/superset/sqllab/utils.py
@@ -0,0 +1,47 @@
+# 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 typing import Any, Dict
+
+from superset.common.db_query_status import QueryStatus
+
+
+def apply_display_max_row_configuration_if_require(
+    sql_results: Dict[str, Any], max_rows_in_result: int
+) -> Dict[str, Any]:
+    """
+    Given a `sql_results` nested structure, applies a limit to the number of rows
+
+    `sql_results` here is the nested structure coming out of sql_lab.get_sql_results, it
+    contains metadata about the query, as well as the data set returned by the query.
+    This method limits the number of rows adds a `displayLimitReached: True` flag to the
+    metadata.
+
+    :param max_rows_in_result:
+    :param sql_results: The results of a sql query from sql_lab.get_sql_results
+    :returns: The mutated sql_results structure
+    """
+
+    def is_require_to_apply() -> bool:
+        return (
+            sql_results["status"] == QueryStatus.SUCCESS
+            and sql_results["query"]["rows"] > max_rows_in_result
+        )
+
+    if is_require_to_apply():
+        sql_results["data"] = sql_results["data"][:max_rows_in_result]
+        sql_results["displayLimitReached"] = True
+    return sql_results
diff --git a/superset/views/core.py b/superset/views/core.py
index 77192ba17e3b6..b1d8f4970f24b 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -98,6 +98,7 @@
 from superset.security.analytics_db_safety import check_sqlalchemy_uri
 from superset.sql_parse import ParsedQuery, Table
 from superset.sql_validators import get_validator_by_name
+from superset.sqllab.utils import apply_display_max_row_configuration_if_require
 from superset.sqllab.command import CommandResult, ExecuteSqlCommand
 from superset.sqllab.command_status import SqlJsonExecutionStatus
 from superset.tasks.async_queries import load_explore_json_into_cache
@@ -128,7 +129,6 @@
 )
 from superset.views.utils import (
     _deserialize_results_payload,
-    apply_display_max_row_limit,
     bootstrap_user_data,
     check_datasource_perms,
     check_explore_cache_perms,
@@ -2314,7 +2314,7 @@ def results_exec(key: str) -> FlaskResponse:
                     status=400,
                 ) from ex
 
-            obj = apply_display_max_row_limit(obj, rows)
+            obj = apply_display_max_row_configuration_if_require(obj, rows)
 
         return json_success(
             json.dumps(