From 0be34c041205bef6592dfcd6d7ff8cbdb12ce015 Mon Sep 17 00:00:00 2001 From: Amir Fathi Date: Sun, 13 Sep 2026 20:03:39 -0600 Subject: [PATCH] fix(wren): bigquery and duckdb LIMIT wrap breaks on SQL ending in a trailing comment (#2734) --- core/wren/src/wren/connector/bigquery.py | 7 +++-- core/wren/src/wren/connector/duckdb.py | 7 +++-- .../tests/unit/test_bigquery_semicolon.py | 16 ++++++++--- .../tests/unit/test_duckdb_file_listing.py | 8 +++--- .../unit/test_duckdb_semicolon_unlimited.py | 28 ++++++++++++++++++- 5 files changed, 53 insertions(+), 13 deletions(-) diff --git a/core/wren/src/wren/connector/bigquery.py b/core/wren/src/wren/connector/bigquery.py index 0b145ee09..daa9558b0 100644 --- a/core/wren/src/wren/connector/bigquery.py +++ b/core/wren/src/wren/connector/bigquery.py @@ -15,10 +15,13 @@ def _apply_limit(sql: str, limit: int) -> str: (after stripping a trailing semicolon) short-circuits the engine and correctly enforces the caller's limit even when *sql* already contains an inner ``LIMIT`` (the outer limit always wins / can only reduce rows). - Avoids comment-sensitive outer-LIMIT detection heuristics. + Avoids comment-sensitive outer-LIMIT detection heuristics. The wrap + itself is multiline so a trailing ``-- line comment`` in *sql* is + terminated by the newline instead of swallowing the closing paren, + alias and LIMIT clause (same technique as postgres.py/athena.py). """ cleaned = strip_trailing_semicolon(sql) - return f"SELECT * FROM ({cleaned}) AS _sub LIMIT {limit}" + return f"SELECT * FROM (\n{cleaned}\n) AS _sub LIMIT {limit}" _SCOPES = [ diff --git a/core/wren/src/wren/connector/duckdb.py b/core/wren/src/wren/connector/duckdb.py index 7f3feca0b..2cee54765 100644 --- a/core/wren/src/wren/connector/duckdb.py +++ b/core/wren/src/wren/connector/duckdb.py @@ -85,7 +85,10 @@ class DuckDBConnector(ConnectorABC): stripped = strip_trailing_semicolon(sql) if limit is not None: # Subquery wrap rejects an interior terminator after strip. - sql = f"SELECT * FROM ({stripped}) AS _q LIMIT {limit}" + # Multiline so a trailing `-- line comment` in `stripped` is + # terminated by the newline instead of swallowing the closing + # `) AS _q LIMIT n` (same technique as postgres.py/athena.py). + sql = f"SELECT * FROM (\n{stripped}\n) AS _q LIMIT {limit}" else: sql = stripped return self.connection.execute(sql).fetch_arrow_table() @@ -102,7 +105,7 @@ class DuckDBConnector(ConnectorABC): no rows are materialized. """ stripped = strip_trailing_semicolon(sql) - self.connection.execute(f"SELECT * FROM ({stripped}) AS _q LIMIT 0") + self.connection.execute(f"SELECT * FROM (\n{stripped}\n) AS _q LIMIT 0") def _attach_database(self, connection_info) -> None: """Attach every discovered DuckDB file as a read-only database. diff --git a/core/wren/tests/unit/test_bigquery_semicolon.py b/core/wren/tests/unit/test_bigquery_semicolon.py index a15ccaa7d..85a80f8cc 100644 --- a/core/wren/tests/unit/test_bigquery_semicolon.py +++ b/core/wren/tests/unit/test_bigquery_semicolon.py @@ -29,8 +29,16 @@ def test_query_pushes_limit_and_strips_semicolon() -> None: connector, client = _make_mock_connector() connector.query("SELECT 1;", limit=5) (sent,), _ = client.query.call_args - assert sent == "SELECT * FROM (SELECT 1) AS _sub LIMIT 5" - assert ";)" not in sent and not sent.rstrip().endswith(";") + assert sent == "SELECT * FROM (\nSELECT 1\n) AS _sub LIMIT 5" + assert ";\n)" not in sent and not sent.rstrip().endswith(";") + + +def test_query_limit_survives_trailing_line_comment() -> None: + """Trailing `--` must not eat the wrap (same shape as postgres.py #2728).""" + connector, client = _make_mock_connector() + connector.query("SELECT 1 AS x -- pick", limit=5) + (sent,), _ = client.query.call_args + assert sent == "SELECT * FROM (\nSELECT 1 AS x -- pick\n) AS _sub LIMIT 5" def test_query_without_limit_still_strips_semicolon() -> None: @@ -45,7 +53,7 @@ def test_query_with_inner_limit_outer_wrap_still_enforces_caller_limit() -> None connector, client = _make_mock_connector() connector.query("SELECT 1 LIMIT 100", limit=5) (sent,), _ = client.query.call_args - assert sent == "SELECT * FROM (SELECT 1 LIMIT 100) AS _sub LIMIT 5" + assert sent == "SELECT * FROM (\nSELECT 1 LIMIT 100\n) AS _sub LIMIT 5" def test_dry_run_strips_trailing_semicolon() -> None: @@ -80,4 +88,4 @@ def test_dry_run_strips_trailing_semicolon() -> None: def test_helper_preserves_literal_semicolon() -> None: assert strip_trailing_semicolon("SELECT ';' AS x") == "SELECT ';' AS x" - assert _apply_limit("SELECT 1;", 3) == "SELECT * FROM (SELECT 1) AS _sub LIMIT 3" + assert _apply_limit("SELECT 1;", 3) == "SELECT * FROM (\nSELECT 1\n) AS _sub LIMIT 3" diff --git a/core/wren/tests/unit/test_duckdb_file_listing.py b/core/wren/tests/unit/test_duckdb_file_listing.py index 214d1e604..fe9f99f12 100644 --- a/core/wren/tests/unit/test_duckdb_file_listing.py +++ b/core/wren/tests/unit/test_duckdb_file_listing.py @@ -95,7 +95,7 @@ def test_query_strips_trailing_semicolon_before_limit_wrap(): result = connector.query("SELECT 1;", limit=5) executed = connector.connection.execute.call_args.args[0] - assert executed == "SELECT * FROM (SELECT 1) AS _q LIMIT 5" + assert executed == "SELECT * FROM (\nSELECT 1\n) AS _q LIMIT 5" assert result == "tbl" @@ -112,7 +112,7 @@ def test_dry_run_wraps_in_limit_zero_subquery(): executed = connector.connection.execute.call_args.args[0] # The trailing terminator is stripped; the interior ``;`` stays inside the # subquery where DuckDB rejects it as a syntax error (no side effects). - assert executed == "SELECT * FROM (SELECT 1; DROP TABLE t) AS _q LIMIT 0" + assert executed == "SELECT * FROM (\nSELECT 1; DROP TABLE t\n) AS _q LIMIT 0" def test_dry_run_strips_trailing_semicolon(): @@ -122,7 +122,7 @@ def test_dry_run_strips_trailing_semicolon(): connector.dry_run("SELECT 1;") executed = connector.connection.execute.call_args.args[0] - assert executed == "SELECT * FROM (SELECT 1) AS _q LIMIT 0" + assert executed == "SELECT * FROM (\nSELECT 1\n) AS _q LIMIT 0" def test_dry_run_preserves_semicolon_in_string_literal(): @@ -134,4 +134,4 @@ def test_dry_run_preserves_semicolon_in_string_literal(): connector.dry_run("SELECT ';' AS x") executed = connector.connection.execute.call_args.args[0] - assert executed == "SELECT * FROM (SELECT ';' AS x) AS _q LIMIT 0" + assert executed == "SELECT * FROM (\nSELECT ';' AS x\n) AS _q LIMIT 0" diff --git a/core/wren/tests/unit/test_duckdb_semicolon_unlimited.py b/core/wren/tests/unit/test_duckdb_semicolon_unlimited.py index fb2083aa7..1436f1a35 100644 --- a/core/wren/tests/unit/test_duckdb_semicolon_unlimited.py +++ b/core/wren/tests/unit/test_duckdb_semicolon_unlimited.py @@ -34,5 +34,31 @@ def test_limited_query_still_wraps_after_strip(): connector.query("SELECT 1 AS x;", limit=2) connector.connection.execute.assert_called_once_with( - "SELECT * FROM (SELECT 1 AS x) AS _q LIMIT 2" + "SELECT * FROM (\nSELECT 1 AS x\n) AS _q LIMIT 2" + ) + + +def test_query_limit_survives_trailing_line_comment(): + """Trailing `--` must not eat the wrap (same shape as postgres.py #2728).""" + connector = DuckDBConnector.__new__(DuckDBConnector) + connector.connection = MagicMock() + connector.connection.execute.return_value.fetch_arrow_table.return_value = ( + pa.table({"x": [1]}) + ) + + connector.query("SELECT 1 AS x -- pick", limit=2) + + connector.connection.execute.assert_called_once_with( + "SELECT * FROM (\nSELECT 1 AS x -- pick\n) AS _q LIMIT 2" + ) + + +def test_dry_run_survives_trailing_line_comment(): + connector = DuckDBConnector.__new__(DuckDBConnector) + connector.connection = MagicMock() + + connector.dry_run("SELECT 1 AS x -- pick") + + connector.connection.execute.assert_called_once_with( + "SELECT * FROM (\nSELECT 1 AS x -- pick\n) AS _q LIMIT 0" )