feat: use sqlglot to set limit by betodealmeida · Pull Request #33473 · apache/superset

betodealmeida · 2025-05-16T15:47:01Z

SUMMARY

Part of #26786, stacked on:

This PR changes the logic of setting the limit in a query to use sqlglot, adding the set_limit_value(limit: int) -> None method to the BaseSQLStatement class.

The method always sets the requested limit. It's up to the application to determine if a bigger limit should be forced or not. Before, it had some logic that determined if the limit should be updated or not depending on if the existing limit was smaller than the desired limit. This keeps the method simpler and decoupled from business logic.

One of the main advantages of using sqlglot here is that it abstracts all the different ways a query can be limited:

SELECT * FROM t LIMIT 10;  -- Postgres, Clickhouse, Bigquery, etc.
SELECT TOP 10 * FROM t;  -- Teradata, MS SQL, Vertica, etc.
SELECT * FROM t FETCH FIRST 10 ROWS ONLY;  -- Oracle, DB2, Apache Derby

These are all abstracted as a sqlglot.exp.Limit, and rendered correctly when the engine is specified (see unit tests).

Previously there was additional logic for Teradata that accounted for the SAMPLE function:

SELECT * FROM t SAMPLE n  -- returns n rows from t

I removed this logic because SAMPLE is not a limit — it limits the number of rows fetched from the table, but before the projection is applied. The projection could have a UDTF that generates millions of rows even if only 10 were read from the table. Because of that, I removed the logic. (The tests were also using invalid syntax, so I'm not even sure if this feature ever worked!)

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TESTING INSTRUCTIONS

Added tests, and moved old tests to the new pattern.

ADDITIONAL INFORMATION

Has associated issue:
Required feature flags:
Changes UI
Includes DB Migration (follow approval process in SIP-59)
- Migration is atomic, supports rollback & is backwards-compatible
- Confirm DB migration upgrade and downgrade tested
- Runtime estimates and downtime expectations provided
Introduces new feature or API
Removes existing feature or API

betodealmeida · 2025-05-16T15:48:37Z

superset/db_engine_specs/kusto.py



 class KustoKqlEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-method
-    limit_method = LimitMethod.WRAP_SQL


Kusto supports limit, no need to wrap the query.

betodealmeida · 2025-05-16T15:48:59Z

superset/db_engine_specs/mssql.py

 class MssqlEngineSpec(BaseEngineSpec):
    engine = "mssql"
    engine_name = "Microsoft SQL Server"
-    limit_method = LimitMethod.WRAP_SQL


MS SQL supports limit, no need to wrap the query.

betodealmeida · 2025-05-16T15:49:14Z

superset/db_engine_specs/teradata.py


    engine = "teradatasql"
    engine_name = "Teradata"
-    limit_method = LimitMethod.WRAP_SQL


Ditto here.

betodealmeida · 2025-05-16T15:49:43Z

superset/models/core.py

+        statement = script.statements[-1]
+        current_limit = statement.get_limit_value() or float("inf")
+
+        if limit < current_limit or force:


Moved the limit application logic here.

korbit-ai

I've completed my review and didn't find any issues.

Files scanned

File Path	Reviewed
superset/db_engine_specs/teradata.py	✅
superset/db_engine_specs/hana.py	✅
superset/db_engine_specs/firebird.py	✅
superset/db_engine_specs/oracle.py	✅
superset/db_engine_specs/db2.py	✅
superset/db_engine_specs/mssql.py	✅
superset/db_engine_specs/kusto.py	✅
superset/db_engine_specs/ocient.py	✅
superset/db_engine_specs/lib.py	✅
superset/sql_parse.py	✅
superset/sql/parse.py	✅
superset/models/core.py	✅
superset/db_engine_specs/base.py	✅

Explore our documentation to understand the languages and file types we support and the files we ignore.

Check out our docs on how you can make Korbit work best for you and your team.

Loving Korbit!? Share us on LinkedIn Reddit and X

betodealmeida · 2025-05-16T15:50:59Z

superset/sql/parse.py

-
-        return self._fallback_formatting()
-
-    @deprecated(deprecated_in="4.0")


This was supposed to be removed in 5.0 and we missed the window, but it doesn't break anything other than not being able to make queries pretty when they don't have a sqlglot dialect.

We need to remove this in order to modify the sqlglot AST inplace.

betodealmeida · 2025-05-16T15:52:10Z

superset/sql/parse.py


        return None

+    def set_limit_value(self, limit: int) -> None:


It's been fun implementing the logic in KQL, but I'm considering writing a SIP to get rid of it, since it's not SQL and doubles the work.

betodealmeida · 2025-05-16T15:52:52Z

tests/unit_tests/db_engine_specs/test_teradata.py

-        (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"),


This is not even valid syntax... SAMPLE goes at the end, not in the projection! The correct would be:

SELECT * FROM My_table SAMPLE 1000

betodealmeida · 2025-05-16T15:53:23Z

tests/unit_tests/db_engine_specs/test_mssql.py

-        ("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),


This is invalid syntax... MS SQL doesn't support the SEL shortcut for SELECT!

Vitor-Avila

LGTM! Manually tested with SQLite and Snowflake too

pull-request-size bot added the size/XL label May 16, 2025

github-actions bot added the preset-io label May 16, 2025

betodealmeida commented May 16, 2025

View reviewed changes

korbit-ai bot reviewed May 16, 2025

View reviewed changes

betodealmeida commented May 16, 2025

View reviewed changes

betodealmeida added the sip Superset Improvement Proposal label May 16, 2025

betodealmeida force-pushed the set-limit branch from 3670d7f to c5364d7 Compare May 16, 2025 15:59

betodealmeida mentioned this pull request May 16, 2025

chore: remove parse_sql #33474

Merged

9 tasks

betodealmeida force-pushed the set-limit branch from c5364d7 to 2c015c7 Compare May 16, 2025 16:46

michael-s-molina requested review from michael-s-molina and villebro May 16, 2025 16:49

pull-request-size bot added size/XXL and removed size/XL labels May 16, 2025

betodealmeida force-pushed the set-limit branch 5 times, most recently from d936b6f to 3708d38 Compare May 16, 2025 19:54

This was referenced May 19, 2025

chore: sql/parse cleanup #33515

Merged

feat: implement CTEs logic in sqlglot #33518

Merged

feat: implement RLS in sqlglot #33524

Merged

feat: implement CVAS/CTAS in sqlglot #33525

Merged

feat(sqllab): use sqlglot instead of sqlparse #33542

Merged

betodealmeida force-pushed the is-select-query-refactor branch from 5cddcea to 6f904f2 Compare May 23, 2025 00:16

Base automatically changed from is-select-query-refactor to master May 23, 2025 00:53

betodealmeida added 2 commits May 22, 2025 20:54

feat: use sqlglot to set limit

0ad538b

Update tests

8c2fef9

betodealmeida force-pushed the set-limit branch from 3708d38 to 8c2fef9 Compare May 23, 2025 00:55

Vitor-Avila approved these changes May 27, 2025

View reviewed changes

betodealmeida merged commit 8de58b9 into master May 27, 2025
49 checks passed

betodealmeida deleted the set-limit branch May 27, 2025 19:20

This was referenced May 28, 2025

feat: use sqlglot to validate adhoc subquery #33560

Merged

chore: remove sqlparse #33564

Merged

chore: 100% test coverage for SQL parsing #33568

Merged

LevisNgigi pushed a commit to LevisNgigi/superset that referenced this pull request Jun 18, 2025

feat: use sqlglot to set limit (apache#33473)

4b3ad61

github-actions bot added 🏷️ bot A label used by `supersetbot` to keep track of which PR where auto-tagged with release labels 🚢 6.0.0 First shipped in 6.0.0 labels Dec 18, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: use sqlglot to set limit#33473

feat: use sqlglot to set limit#33473
betodealmeida merged 2 commits intomasterfrom
set-limit

betodealmeida commented May 16, 2025 •

edited

Loading

Uh oh!

betodealmeida May 16, 2025

Uh oh!

betodealmeida May 16, 2025

Uh oh!

betodealmeida May 16, 2025

Uh oh!

betodealmeida May 16, 2025

Uh oh!

korbit-ai bot left a comment

Uh oh!

betodealmeida May 16, 2025

Uh oh!

betodealmeida May 16, 2025 •

edited

Loading

Uh oh!

betodealmeida May 16, 2025 •

edited

Loading

Uh oh!

betodealmeida May 16, 2025

Uh oh!

Vitor-Avila left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		class KustoKqlEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method
		limit_method = LimitMethod.WRAP_SQL


		return self._fallback_formatting()

		@deprecated(deprecated_in="4.0")

		(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"),

Conversation

betodealmeida commented May 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

SUMMARY

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TESTING INSTRUCTIONS

ADDITIONAL INFORMATION

Uh oh!

betodealmeida May 16, 2025

Choose a reason for hiding this comment

Uh oh!

betodealmeida May 16, 2025

Choose a reason for hiding this comment

Uh oh!

betodealmeida May 16, 2025

Choose a reason for hiding this comment

Uh oh!

betodealmeida May 16, 2025

Choose a reason for hiding this comment

Uh oh!

korbit-ai bot left a comment

Choose a reason for hiding this comment

Uh oh!

betodealmeida May 16, 2025

Choose a reason for hiding this comment

Uh oh!

betodealmeida May 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

betodealmeida May 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

betodealmeida May 16, 2025

Choose a reason for hiding this comment

Uh oh!

Vitor-Avila left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

betodealmeida commented May 16, 2025 •

edited

Loading

betodealmeida May 16, 2025 •

edited

Loading

betodealmeida May 16, 2025 •

edited

Loading