(superset) 01/02: fix: no limit in SELECT * for TOP dbs (#27215)

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

michaelsmolina pushed a commit to branch 3.1
in repository https://gitbox.apache.org/repos/asf/superset.git

commit ee7eec60ea9854773400361088f5f792e7135664
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Thu Feb 22 17:01:29 2024 -0500

    fix: no limit in SELECT * for TOP dbs (#27215)
    
    (cherry picked from commit c54fbe6e969fe57cddc69759796e4df1f603430e)
---
 superset/db_engine_specs/base.py              |  2 +-
 tests/unit_tests/db_engine_specs/test_base.py | 75 ++++++++++++++++++++++++++-
 2 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 5edd4b7c06..66293ccf52 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -1441,7 +1441,7 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
 
         qry = select(fields).select_from(text(full_table_name))
 
-        if limit:
+        if limit and cls.allow_limit_clause:
             qry = qry.limit(limit)
         if latest_partition:
             partition_query = cls.where_latest_partition(
diff --git a/tests/unit_tests/db_engine_specs/test_base.py b/tests/unit_tests/db_engine_specs/test_base.py
index 7b1977ff14..e4cf4dfc1f 100644
--- a/tests/unit_tests/db_engine_specs/test_base.py
+++ b/tests/unit_tests/db_engine_specs/test_base.py
@@ -14,13 +14,16 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=unused-argument, import-outside-toplevel, protected-access
+# pylint: disable=import-outside-toplevel, protected-access
 
 from textwrap import dedent
 from typing import Any, Optional
 
 import pytest
+from pytest_mock import MockFixture
 from sqlalchemy import types
+from sqlalchemy.dialects import sqlite
+from sqlalchemy.sql import sqltypes
 
 from superset.superset_typing import ResultSetColumnType, SQLAColumnType
 from superset.utils.core import GenericDataType
@@ -168,3 +171,73 @@ def test_convert_inspector_columns(
     from superset.db_engine_specs.base import convert_inspector_columns
 
     assert convert_inspector_columns(cols) == expected_result
+
+
+def test_select_star(mocker: MockFixture) -> None:
+    """
+    Test the ``select_star`` method.
+    """
+    from superset.db_engine_specs.base import BaseEngineSpec
+
+    class NoLimitDBEngineSpec(BaseEngineSpec):
+        allow_limit_clause = False
+
+    cols: list[ResultSetColumnType] = [
+        {
+            "column_name": "a",
+            "name": "a",
+            "type": sqltypes.String(),
+            "nullable": True,
+            "comment": None,
+            "default": None,
+            "precision": None,
+            "scale": None,
+            "max_length": None,
+            "is_dttm": False,
+        },
+    ]
+
+    # mock the database so we can compile the query
+    database = mocker.MagicMock()
+    database.compile_sqla_query = lambda query: str(
+        query.compile(dialect=sqlite.dialect())
+    )
+
+    engine = mocker.MagicMock()
+    engine.dialect = sqlite.dialect()
+
+    sql = BaseEngineSpec.select_star(
+        database=database,
+        table_name="my_table",
+        engine=engine,
+        schema=None,
+        limit=100,
+        show_cols=True,
+        indent=True,
+        latest_partition=False,
+        cols=cols,
+    )
+    assert (
+        sql
+        == """SELECT a
+FROM my_table
+LIMIT ?
+OFFSET ?"""
+    )
+
+    sql = NoLimitDBEngineSpec.select_star(
+        database=database,
+        table_name="my_table",
+        engine=engine,
+        schema=None,
+        limit=100,
+        show_cols=True,
+        indent=True,
+        latest_partition=False,
+        cols=cols,
+    )
+    assert (
+        sql
+        == """SELECT a
+FROM my_table"""
+    )