fix(wren): bigquery and duckdb LIMIT wrap breaks on SQL ending in a trailing comment (#2734)

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 = [

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.

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"

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"

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