fix(wren): postgres LIMIT pushdown breaks on SQL ending in a trailing comment (#2728)

Signed-off-by: Amir Fathi <amirfathi.me@gmail.com>

core/wren/src/wren/connector/postgres.py

@@ -317,7 +317,10 @@ class PostgresConnector(ConnectorABC):
         # client-pasted statements match dry_run / limited composition rules.
         sql = strip_trailing_semicolon(sql)
         if limit is not None:
-            sql = f"SELECT * FROM ({sql}) AS _sub LIMIT {limit}"
+            # Multiline wrap so a trailing `-- line comment` in the inner SQL
+            # is terminated by the newline instead of swallowing the closing
+            # `) AS _sub LIMIT n` (same technique as athena.py/snowflake.py).
+            sql = f"SELECT * FROM (\n{sql}\n) AS _sub LIMIT {limit}"
 
         try:
             with self.connection.cursor() as cursor:
@@ -336,7 +339,7 @@ class PostgresConnector(ConnectorABC):
             ) from e
 
     def dry_run(self, sql: str) -> None:
-        wrapped = f"SELECT * FROM ({strip_trailing_semicolon(sql)}) AS _sub LIMIT 0"
+        wrapped = f"SELECT * FROM (\n{strip_trailing_semicolon(sql)}\n) AS _sub LIMIT 0"
         try:
             with self.connection.cursor() as cursor:
                 cursor.execute(wrapped)

core/wren/tests/connectors/test_postgres.py

@@ -461,15 +461,15 @@ def test_query_strips_trailing_semicolon_before_subquery_wrap() -> None:
     connector, cursor = _make_mock_connector()
     connector.query("SELECT 1;", limit=5)
     (sent,), _ = cursor.execute.call_args
-    assert sent == "SELECT * FROM (SELECT 1) AS _sub LIMIT 5"
-    assert ";)" not in sent
+    assert sent == "SELECT * FROM (\nSELECT 1\n) AS _sub LIMIT 5"
+    assert ";\n)" not in sent
 
 
 def test_dry_run_strips_trailing_semicolon() -> None:
     connector, cursor = _make_mock_connector()
     connector.dry_run("SELECT 1;  ")
     (sent,), _ = cursor.execute.call_args
-    assert sent == "SELECT * FROM (SELECT 1) AS _sub LIMIT 0"
+    assert sent == "SELECT * FROM (\nSELECT 1\n) AS _sub LIMIT 0"
 
 
 def test_helper_preserves_semicolon_inside_string_literal() -> None:

core/wren/tests/unit/test_postgres_semicolon_unlimited.py

@@ -59,5 +59,35 @@ def test_limited_query_wraps_after_strip(monkeypatch):
     connector.query("SELECT 1 AS x;", limit=9)
 
     cursor.execute.assert_called_once_with(
-        "SELECT * FROM (SELECT 1 AS x) AS _sub LIMIT 9"
+        "SELECT * FROM (\nSELECT 1 AS x\n) AS _sub LIMIT 9"
+    )
+
+
+def test_query_limit_survives_trailing_line_comment(monkeypatch):
+    """Trailing `--` must not eat the wrap (same shape as athena.py #2457)."""
+    connector = PostgresConnector.__new__(PostgresConnector)
+    connector.connection = MagicMock()
+    cursor = MagicMock()
+    connector.connection.cursor.return_value.__enter__.return_value = cursor
+    monkeypatch.setattr(
+        postgres_mod, "_build_pg_arrow_table", lambda cur: pa.table({"x": [1]})
+    )
+
+    connector.query("SELECT 1 AS x -- pick", limit=9)
+
+    cursor.execute.assert_called_once_with(
+        "SELECT * FROM (\nSELECT 1 AS x -- pick\n) AS _sub LIMIT 9"
+    )
+
+
+def test_dry_run_limit_survives_trailing_line_comment():
+    connector = PostgresConnector.__new__(PostgresConnector)
+    connector.connection = MagicMock()
+    cursor = MagicMock()
+    connector.connection.cursor.return_value.__enter__.return_value = cursor
+
+    connector.dry_run("SELECT 1 AS x -- pick")
+
+    cursor.execute.assert_called_once_with(
+        "SELECT * FROM (\nSELECT 1 AS x -- pick\n) AS _sub LIMIT 0"
     )