feat: use sqlglot to set limit (#33473)

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

View File

@@ -206,9 +206,6 @@ def test_select_star(mocker: MockerFixture) -> None:
"""
from superset.db_engine_specs.base import BaseEngineSpec
class NoLimitDBEngineSpec(BaseEngineSpec):
allow_limit_clause = False
cols: list[ResultSetColumnType] = [
{
"column_name": "a",
@@ -243,19 +240,7 @@ def test_select_star(mocker: MockerFixture) -> None:
latest_partition=False,
cols=cols,
)
assert sql == "SELECT a\nFROM my_table\nLIMIT ?\nOFFSET ?"
sql = NoLimitDBEngineSpec.select_star(
database=database,
table=Table("my_table"),
engine=engine,
limit=100,
show_cols=True,
indent=True,
latest_partition=False,
cols=cols,
)
assert sql == "SELECT a\nFROM my_table"
assert sql == "SELECT\n a\nFROM my_table\nLIMIT ?\nOFFSET ?"
def test_extra_table_metadata(mocker: MockerFixture) -> None:

View File

@@ -254,36 +254,6 @@ def test_cte_query_parsing(original: TypeEngine, expected: str) -> None:
assert actual == expected
@pytest.mark.parametrize(
"original,expected,top",
[
("SEL TOP 1000 * FROM My_table", "SEL TOP 100 * FROM My_table", 100),
("SEL TOP 1000 * FROM My_table;", "SEL TOP 100 * FROM My_table", 100),
("SEL TOP 1000 * FROM My_table;", "SEL TOP 1000 * FROM My_table", 10000),
("SEL TOP 1000 * FROM My_table;", "SEL TOP 1000 * FROM My_table", 1000),
(
"""with abc as (select * from test union select * from test1)
select TOP 100 * from currency""",
"""WITH abc as (select * from test union select * from test1)
select TOP 100 * from currency""",
1000,
),
("SELECT DISTINCT x from tbl", "SELECT DISTINCT TOP 100 x from tbl", 100),
("SELECT 1 as cnt", "SELECT TOP 10 1 as cnt", 10),
(
"select TOP 1000 * from abc where id=1",
"select TOP 10 * from abc where id=1",
10,
),
],
)
def test_top_query_parsing(original: TypeEngine, expected: str, top: int) -> None:
from superset.db_engine_specs.mssql import MssqlEngineSpec
actual = MssqlEngineSpec.apply_top_to_sql(original, top)
assert actual == expected
def test_extract_errors() -> None:
"""
Test that custom error messages are extracted correctly.

View File

@@ -1,43 +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.
# pylint: disable=unused-argument, import-outside-toplevel, protected-access
import pytest
@pytest.mark.parametrize(
"limit,original,expected",
[
(100, "SEL TOP 1000 * FROM My_table", "SEL TOP 100 * FROM My_table"),
(100, "SEL TOP 1000 * FROM My_table;", "SEL TOP 100 * FROM My_table"),
(10000, "SEL TOP 1000 * FROM My_table;", "SEL TOP 1000 * FROM My_table"),
(1000, "SEL TOP 1000 * FROM My_table;", "SEL TOP 1000 * FROM My_table"),
(100, "SELECT TOP 1000 * FROM My_table", "SELECT TOP 100 * FROM My_table"),
(100, "SEL SAMPLE 1000 * FROM My_table", "SEL SAMPLE 100 * FROM My_table"),
(10000, "SEL SAMPLE 1000 * FROM My_table", "SEL SAMPLE 1000 * FROM My_table"),
],
)
def test_apply_top_to_sql_limit(
limit: int,
original: str,
expected: str,
) -> None:
"""
Ensure limits are applied to the query correctly
"""
from superset.db_engine_specs.teradata import TeradataEngineSpec
assert TeradataEngineSpec.apply_top_to_sql(original, limit) == expected