This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch set-limit
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/set-limit by this push:
new 92efb6a951 Update tests
92efb6a951 is described below
commit 92efb6a95154fe69b3c9df09a50f3a55f812c346
Author: Beto Dealmeida <[email protected]>
AuthorDate: Fri May 16 14:50:17 2025 -0400
Update tests
---
superset/models/core.py | 2 +-
superset/sql/parse.py | 15 ++-
.../db_engine_specs/base_engine_spec_tests.py | 116 -----------------
.../db_engine_specs/base_tests.py | 36 ------
tests/unit_tests/models/core_test.py | 142 +++++++++++++++++++++
5 files changed, 156 insertions(+), 155 deletions(-)
diff --git a/superset/models/core.py b/superset/models/core.py
index ee14f6c66e..54d1c961b8 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -772,7 +772,7 @@ class Database(Model, AuditMixinNullable,
ImportExportMixin): # pylint: disable
current_limit = statement.get_limit_value() or float("inf")
if limit < current_limit or force:
- statement.set_limit_value(limit)
+ statement.set_limit_value(limit, self.db_engine_spec.limit_method)
return script.format()
diff --git a/superset/sql/parse.py b/superset/sql/parse.py
index 6172c349b7..d580a7a80e 100644
--- a/superset/sql/parse.py
+++ b/superset/sql/parse.py
@@ -261,7 +261,11 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
"""
raise NotImplementedError()
- def set_limit_value(self, limit: int) -> None:
+ def set_limit_value(
+ self,
+ limit: int,
+ method: LimitMethod = LimitMethod.FORCE_LIMIT,
+ ) -> None:
"""
Add a limit to the statement.
"""
@@ -781,10 +785,17 @@ class KustoKQLStatement(BaseSQLStatement[str]):
return None
- def set_limit_value(self, limit: int) -> None:
+ def set_limit_value(
+ self,
+ limit: int,
+ method: LimitMethod = LimitMethod.FORCE_LIMIT,
+ ) -> None:
"""
Add a limit to the statement.
"""
+ if method != LimitMethod.FORCE_LIMIT:
+ raise SupersetParseError("Kusto KQL only supports the FORCE_LIMIT
method.")
+
tokens = tokenize_kql(self._sql)
found_limit_token = False
for idx, (ttype, val) in enumerate(tokens):
diff --git a/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py
b/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py
index 7dd9c5cb95..fb6591a7ea 100644
--- a/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py
+++ b/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py
@@ -25,7 +25,6 @@ from superset.db_engine_specs.base import (
BaseEngineSpec,
BasicParametersMixin,
builtin_time_grains,
- LimitMethod,
)
from superset.db_engine_specs.mysql import MySQLEngineSpec
from superset.db_engine_specs.sqlite import SqliteEngineSpec
@@ -74,124 +73,9 @@ class TestDbEngineSpecs(TestDbEngineSpec):
assert engine_spec_class.get_limit_from_sql(q10) is None
assert engine_spec_class.get_limit_from_sql(q11) is None
- def test_wrapped_semi_tabs(self):
- self.sql_limit_regex(
- "SELECT * FROM a \t \n ; \t \n ", "SELECT * FROM a\nLIMIT
1000"
- )
-
- def test_simple_limit_query(self):
- self.sql_limit_regex("SELECT * FROM a", "SELECT * FROM a\nLIMIT 1000")
-
- def test_modify_limit_query(self):
- self.sql_limit_regex("SELECT * FROM a LIMIT 9999", "SELECT * FROM a
LIMIT 1000")
-
- def test_limit_query_with_limit_subquery(self): # pylint:
disable=invalid-name
- self.sql_limit_regex(
- "SELECT * FROM (SELECT * FROM a LIMIT 10) LIMIT 9999",
- "SELECT * FROM (SELECT * FROM a LIMIT 10) LIMIT 1000",
- )
-
- def test_limit_query_without_force(self):
- self.sql_limit_regex(
- "SELECT * FROM a LIMIT 10",
- "SELECT * FROM a LIMIT 10",
- limit=11,
- )
-
- def test_limit_query_with_force(self):
- self.sql_limit_regex(
- "SELECT * FROM a LIMIT 10",
- "SELECT * FROM a LIMIT 11",
- limit=11,
- force=True,
- )
-
- def test_limit_with_expr(self):
- self.sql_limit_regex(
- """
- SELECT
- 'LIMIT 777' AS a
- , b
- FROM
- table
- LIMIT 99990""",
- """SELECT
- 'LIMIT 777' AS a
- , b
- FROM
- table
- LIMIT 1000""",
- )
-
- def test_limit_expr_and_semicolon(self):
- self.sql_limit_regex(
- """
- SELECT
- 'LIMIT 777' AS a
- , b
- FROM
- table
- LIMIT 99990 ;""",
- """SELECT
- 'LIMIT 777' AS a
- , b
- FROM
- table
- LIMIT 1000""",
- )
-
def test_get_datatype(self):
assert "VARCHAR" == BaseEngineSpec.get_datatype("VARCHAR")
- def test_limit_with_implicit_offset(self):
- self.sql_limit_regex(
- """
- SELECT
- 'LIMIT 777' AS a
- , b
- FROM
- table
- LIMIT 99990, 999999""",
- """SELECT
- 'LIMIT 777' AS a
- , b
- FROM
- table
- LIMIT 99990, 1000""",
- )
-
- def test_limit_with_explicit_offset(self):
- self.sql_limit_regex(
- """
- SELECT
- 'LIMIT 777' AS a
- , b
- FROM
- table
- LIMIT 99990
- OFFSET 999999""",
- """SELECT
- 'LIMIT 777' AS a
- , b
- FROM
- table
- LIMIT 1000
- OFFSET 999999""",
- )
-
- def test_limit_with_non_token_limit(self):
- self.sql_limit_regex(
- """SELECT 'LIMIT 777'""", """SELECT 'LIMIT 777'\nLIMIT 1000"""
- )
-
- def test_limit_with_fetch_many(self):
- class DummyEngineSpec(BaseEngineSpec):
- limit_method = LimitMethod.FETCH_MANY
-
- self.sql_limit_regex(
- "SELECT * FROM table", "SELECT * FROM table", DummyEngineSpec
- )
-
def test_engine_time_grain_validity(self):
time_grains = set(builtin_time_grains.keys())
# loop over all subclasses of BaseEngineSpec
diff --git a/tests/integration_tests/db_engine_specs/base_tests.py
b/tests/integration_tests/db_engine_specs/base_tests.py
deleted file mode 100644
index 20ffcfdbc1..0000000000
--- a/tests/integration_tests/db_engine_specs/base_tests.py
+++ /dev/null
@@ -1,36 +0,0 @@
-# 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.
-# isort:skip_file
-
-from tests.integration_tests.test_app import app # noqa: F401
-from tests.integration_tests.base_tests import SupersetTestCase
-from superset.db_engine_specs.base import BaseEngineSpec
-from superset.models.core import Database
-
-
-class TestDbEngineSpec(SupersetTestCase):
- def sql_limit_regex(
- self,
- sql,
- expected_sql,
- engine_spec_class=BaseEngineSpec,
- limit=1000,
- force=False,
- ):
- main = Database(database_name="test_database",
sqlalchemy_uri="sqlite://")
- limited = main.apply_limit_to_sql(sql, limit, force)
- assert expected_sql == limited
diff --git a/tests/unit_tests/models/core_test.py
b/tests/unit_tests/models/core_test.py
index 50ccde6605..981fab3e8b 100644
--- a/tests/unit_tests/models/core_test.py
+++ b/tests/unit_tests/models/core_test.py
@@ -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"),
}
)
+
+
[email protected](
+ "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
