From c54fbe6e969fe57cddc69759796e4df1f603430e Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Thu, 22 Feb 2024 17:01:29 -0500 Subject: [PATCH] fix: no limit in SELECT * for TOP dbs (#27215) --- 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 3b8bb2bd3329..5b850913ed77 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1438,7 +1438,7 @@ def select_star( # pylint: disable=too-many-arguments,too-many-locals 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 7b1977ff14d0..e4cf4dfc1f9c 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""" + )