From b8ca729f9fc65d3bbdf16a8c08aa076510aedb0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CA=88=E1=B5=83=E1=B5=A2?= Date: Thu, 13 Aug 2026 10:43:45 -0700 Subject: [PATCH] fix(sql): only force a LIMIT onto query expressions (#43097) --- superset/sql/parse.py | 29 +++++++---- tests/unit_tests/sql/parse_tests.py | 74 ++++++++++++++++++++++++++--- 2 files changed, 88 insertions(+), 15 deletions(-) diff --git a/superset/sql/parse.py b/superset/sql/parse.py index 6d79abe1c0e..dc6d88acbd7 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py @@ -1408,16 +1408,27 @@ class SQLStatement(BaseSQLStatement[exp.Expression]): """ Modify the `LIMIT` or `TOP` value of the SQL statement inplace. """ + # Only query expressions -- `SELECT`, `UNION`, subqueries -- can carry a + # limit. Everything else (`SHOW`, `DESCRIBE`, `SET`, `USE`, `GRANT`, and + # anything sqlglot falls back to parsing as an opaque `Command`) has no + # `LIMIT` slot, so it is left untouched. + # + # `apply_limit()` only skips *mutating* statements, so read-only metadata + # statements do reach this method. Forcing a limit into one rewrites it + # into something the engine never asked for: on MySQL/StarRocks a `SHOW` + # renders with two `LIMIT` keywords and is rejected outright ("Getting + # syntax error ... Unexpected input 'LIMIT'"), and `WRAP_SQL` buries it + # in `SELECT * FROM (SHOW DATABASES)`. Dialects that would render a + # valid `SHOW ... LIMIT` are skipped too: `SHOW` returns bounded + # metadata, so there is nothing to truncate. + # + # The guard is on the node category rather than an `exp.Show` + # special-case so it holds for every non-query statement, including ones + # whose generators may learn to render a `limit` arg in a later sqlglot. + if not isinstance(self._parsed, exp.Query): + return + if method == LimitMethod.FORCE_LIMIT: - # `SHOW` statements (`SHOW TABLES`, `SHOW DATABASES`, `SHOW CREATE - # TABLE`, etc.) have no meaningful `LIMIT` slot to force. On - # MySQL/StarRocks, writing one renders a malformed statement with - # two `LIMIT` keywords that the engine rejects outright; on dialects - # like Snowflake it would render a valid `SHOW ... LIMIT`, but SHOW - # returns bounded metadata, so we skip it uniformly rather than - # special-case per dialect. Leave them untouched. - if isinstance(self._parsed, exp.Show): - return self._parsed.args["limit"] = exp.Limit( expression=exp.Literal(this=str(limit), is_string=False) ) diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index 45eb109f660..8d2924f1c98 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py @@ -2536,6 +2536,10 @@ def test_set_limit_value( assert statement.format() == expected +@pytest.mark.parametrize( + "method", + [LimitMethod.FORCE_LIMIT, LimitMethod.WRAP_SQL], +) @pytest.mark.parametrize( "engine", [ @@ -2556,10 +2560,10 @@ def test_set_limit_value( ], ) def test_set_limit_value_leaves_show_statements_unchanged( - sql: str, engine: str + sql: str, engine: str, method: LimitMethod ) -> None: """ - Regression for #36939: FORCE_LIMIT must not touch ``SHOW`` statements. + Regression for #36939: no limit method may touch ``SHOW`` statements. ``SHOW`` statements have no `LIMIT` clause in sqlglot's expression tree, so forcing one via ``args["limit"]`` doesn't reject cleanly, it produces @@ -2570,18 +2574,76 @@ def test_set_limit_value_leaves_show_statements_unchanged( left untouched instead, matching how ``SELECT`` statements without a scannable row source aren't force-limited either. + ``WRAP_SQL`` is wrong on a ``SHOW`` for the same reason but fails more + quietly, rewriting it as ``SELECT * FROM (SHOW DATABASES)``, so both + methods are covered here. + Covers multiple engines, not just StarRocks: the fix guards on the AST - node type (``exp.Show``), not the dialect, so any engine whose sqlglot - dialect parses ``SHOW`` into a real ``Show`` node (e.g. MySQL, Snowflake) - is equally exposed and must be equally protected. + node category (``exp.Query``), not the dialect, so any engine whose + sqlglot dialect parses ``SHOW`` into a real ``Show`` node (e.g. MySQL, + Snowflake) is equally exposed and must be equally protected. """ statement = SQLStatement(sql, engine) original = statement.format() - statement.set_limit_value(1000, LimitMethod.FORCE_LIMIT) + statement.set_limit_value(1000, method) assert statement.format() == original assert "LIMIT" not in statement.format() +@pytest.mark.parametrize( + "method", + [LimitMethod.FORCE_LIMIT, LimitMethod.WRAP_SQL], +) +@pytest.mark.parametrize( + "sql", + [ + "DESCRIBE test.will_test1", + "USE test", + "SET time_zone = 'UTC'", + "GRANT SELECT ON t1 TO u1", + ], +) +def test_set_limit_value_leaves_non_query_statements_unchanged( + sql: str, method: LimitMethod +) -> None: + """ + ``SHOW`` is not the only statement with nowhere to put a `LIMIT`. + + `apply_limit()` only skips *mutating* statements, so every read-only + non-query statement reaches ``set_limit_value``. These happen to survive + a forced limit today only because their sqlglot generators ignore an + unexpected ``limit`` arg -- a silent dependency on generator internals. + Guarding on ``exp.Query`` makes leaving them alone explicit, so a future + sqlglot that starts rendering `limit` for one of these node types can't + reintroduce the ``SHOW`` bug under a different keyword. + """ + statement = SQLStatement(sql, "starrocks") + original = statement.format() + statement.set_limit_value(1000, method) + assert statement.format() == original + assert "LIMIT" not in statement.format() + + +@pytest.mark.parametrize( + "sql", + [ + # `UNION` parses as `exp.Union` and a parenthesized query as + # `exp.Subquery` -- neither is an `exp.Select`, so narrowing the guard + # to `is_select()` would silently stop limiting them. + "SELECT 1 UNION SELECT 2", + "(SELECT 1)", + "WITH t AS (SELECT 1) SELECT * FROM t", + ], +) +def test_set_limit_value_limits_non_select_query_expressions(sql: str) -> None: + """ + Query expressions that aren't `SELECT` must still be limited. + """ + statement = SQLStatement(sql, "starrocks") + statement.set_limit_value(1000, LimitMethod.FORCE_LIMIT) + assert "LIMIT 1000" in statement.format() + + @pytest.mark.parametrize( "kql, limit, expected", [