fix: preventing sql lab None limit value (#17155)

* fix: preventing sql lab None limit value

* test: create a test for the fix

* pylint (#17172)

* add test (#17173)

Co-authored-by: Hugh A. Miles II <hughmil3s@gmail.com>

superset/sql_lab.py

@@ -38,6 +38,7 @@
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import SupersetErrorException, SupersetErrorsException
 from superset.extensions import celery_app
+from superset.models.core import Database
 from superset.models.sql_lab import Query
 from superset.result_set import SupersetResultSet
 from superset.sql_parse import CtasMethod, ParsedQuery
@@ -221,12 +222,7 @@ def execute_sql_statement(  # pylint: disable=too-many-arguments,too-many-locals
     ):
         if SQL_MAX_ROW and (not query.limit or query.limit > SQL_MAX_ROW):
             query.limit = SQL_MAX_ROW
-        if query.limit:
-            # We are fetching one more than the requested limit in order
-            # to test whether there are more rows than the limit.
-            # Later, the extra row will be dropped before sending
-            # the results back to the user.
-            sql = database.apply_limit_to_sql(sql, increased_limit, force=True)
+        sql = apply_limit_if_exists(database, increased_limit, query, sql)
 
     # Hook to allow environment-specific mutation (usually comments) to the SQL
     sql = SQL_QUERY_MUTATOR(sql, user_name, security_manager, database)
@@ -291,6 +287,18 @@ def execute_sql_statement(  # pylint: disable=too-many-arguments,too-many-locals
     return SupersetResultSet(data, cursor_description, db_engine_spec)
 
 
+def apply_limit_if_exists(
+    database: Database, increased_limit: Optional[int], query: Query, sql: str
+) -> str:
+    if query.limit and increased_limit:
+        # We are fetching one more than the requested limit in order
+        # to test whether there are more rows than the limit.
+        # Later, the extra row will be dropped before sending
+        # the results back to the user.
+        sql = database.apply_limit_to_sql(sql, increased_limit, force=True)
+    return sql
+
+
 def _serialize_payload(
     payload: Dict[Any, Any], use_msgpack: Optional[bool] = False
 ) -> Union[bytes, str]:

tests/integration_tests/sqllab_tests.py

@@ -44,6 +44,7 @@
     execute_sql_statement,
     get_sql_results,
     SqlLabException,
+    apply_limit_if_exists,
 )
 from superset.sql_parse import CtasMethod
 from superset.utils.core import (
@@ -990,6 +991,29 @@ def test_sql_json_soft_timeout(self):
             ]
         }
 
+    def test_apply_limit_if_exists_when_incremented_limit_is_none(self):
+        sql = """
+                   SET @value = 42;
+                   SELECT @value AS foo;
+               """
+        database = get_example_database()
+        mock_query = mock.MagicMock()
+        mock_query.limit = 300
+        final_sql = apply_limit_if_exists(database, None, mock_query, sql)
+
+        assert final_sql == sql
+
+    def test_apply_limit_if_exists_when_increased_limit(self):
+        sql = """
+                   SET @value = 42;
+                   SELECT @value AS foo;
+               """
+        database = get_example_database()
+        mock_query = mock.MagicMock()
+        mock_query.limit = 300
+        final_sql = apply_limit_if_exists(database, 1000, mock_query, sql)
+        assert "LIMIT 1000" in final_sql
+
 
 @pytest.mark.parametrize("spec", [HiveEngineSpec, PrestoEngineSpec])
 def test_cancel_query_implicit(spec: BaseEngineSpec) -> None: