QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-17 07:05:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: use sqlglot to set limit (#33473)

Browse Source
This commit is contained in:
Beto Dealmeida
2025-05-27 15:20:02 -04:00
committed by GitHub
parent cc8ab2c556
commit 8de58b9848
34 changed files with 573 additions and 557 deletions
Download Patch File Download Diff File Expand all files Collapse all files

142
tests/unit_tests/models/core_test.py
View File

@@ -38,6 +38,7 @@ from superset.connectors.sqla.models import SqlaTable, TableColumn
from superset.errors import SupersetErrorType
from superset.exceptions import OAuth2Error, OAuth2RedirectError
from superset.models.core import Database
from superset.sql.parse import LimitMethod
from superset.sql_parse import Table
from superset.utils import json
from tests.unit_tests.conftest import with_feature_flags
@@ -910,3 +911,144 @@ def test_get_all_view_names_in_schema(mocker: MockerFixture) -> None:
("third_view", "public", "examples"),
}
)
@pytest.mark.parametrize(
"sql, limit, force, method, expected",
[
(
"SELECT * FROM table",
100,
False,
LimitMethod.FORCE_LIMIT,
"SELECT\n *\nFROM table\nLIMIT 100",
),
(
"SELECT * FROM table LIMIT 100",
10,
False,
LimitMethod.FORCE_LIMIT,
"SELECT\n *\nFROM table\nLIMIT 10",
),
(
"SELECT * FROM table LIMIT 10",
100,
False,
LimitMethod.FORCE_LIMIT,
"SELECT\n *\nFROM table\nLIMIT 10",
),
(
"SELECT * FROM table LIMIT 10",
100,
True,
LimitMethod.FORCE_LIMIT,
"SELECT\n *\nFROM table\nLIMIT 100",
),
(
"SELECT * FROM a \t \n ; \t \n ",
1000,
False,
LimitMethod.FORCE_LIMIT,
"SELECT\n *\nFROM a\nLIMIT 1000",
),
(
"SELECT 'LIMIT 777'",
1000,
False,
LimitMethod.FORCE_LIMIT,
"SELECT\n 'LIMIT 777'\nLIMIT 1000",
),
(
"SELECT * FROM table",
1000,
False,
LimitMethod.FETCH_MANY,
"SELECT\n *\nFROM table",
),
(
"SELECT * FROM (SELECT * FROM a LIMIT 10) LIMIT 9999",
1000,
False,
LimitMethod.FORCE_LIMIT,
"""SELECT
*
FROM (
SELECT
*
FROM a
LIMIT 10
)
LIMIT 1000""",
),
(
"""
SELECT
'LIMIT 777' AS a
, b
FROM
table
LIMIT 99990""",
1000,
None,
LimitMethod.FORCE_LIMIT,
"SELECT\n 'LIMIT 777' AS a,\n b\nFROM table\nLIMIT 1000",
),
(
"""
SELECT
'LIMIT 777' AS a
, b
FROM
table
LIMIT 99990 ;""",
1000,
None,
LimitMethod.FORCE_LIMIT,
"SELECT\n 'LIMIT 777' AS a,\n b\nFROM table\nLIMIT 1000",
),
(
"""
SELECT
'LIMIT 777' AS a
, b
FROM
table
LIMIT 99990, 999999""",
1000,
None,
LimitMethod.FORCE_LIMIT,
"SELECT\n 'LIMIT 777' AS a,\n b\nFROM table\nLIMIT 1000\nOFFSET 99990",
),
(
"""
SELECT
'LIMIT 777' AS a
, b
FROM
table
LIMIT 99990
OFFSET 999999""",
1000,
None,
LimitMethod.FORCE_LIMIT,
"SELECT\n 'LIMIT 777' AS a,\n b\nFROM table\nLIMIT 1000\nOFFSET 999999",
),
],
)
def test_apply_limit_to_sql(
sql: str,
limit: int,
force: bool,
method: LimitMethod,
expected: str,
mocker: MockerFixture,
) -> None:
"""
Test the `apply_limit_to_sql` method.
"""
db = Database(database_name="test_database", sqlalchemy_uri="sqlite://")
db_engine_spec = mocker.MagicMock(limit_method=method)
db.get_db_engine_spec = mocker.MagicMock(return_value=db_engine_spec)
limited = db.apply_limit_to_sql(sql, limit, force)
assert limited == expected