From 9ff9649f91c8a84f721fcbecb33cc9f71ca2d1a5 Mon Sep 17 00:00:00 2001 From: Amir Fathi Date: Tue, 8 Sep 2026 21:17:50 -0600 Subject: [PATCH] fix(wren): postgres LIMIT pushdown breaks on SQL ending in a trailing comment (#2728) Signed-off-by: Amir Fathi --- core/wren/src/wren/connector/postgres.py | 7 ++-- core/wren/tests/connectors/test_postgres.py | 6 ++-- .../unit/test_postgres_semicolon_unlimited.py | 32 ++++++++++++++++++- 3 files changed, 39 insertions(+), 6 deletions(-) diff --git a/core/wren/src/wren/connector/postgres.py b/core/wren/src/wren/connector/postgres.py index 9b355b18d..a59e141df 100644 --- a/core/wren/src/wren/connector/postgres.py +++ b/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) diff --git a/core/wren/tests/connectors/test_postgres.py b/core/wren/tests/connectors/test_postgres.py index e5fd38bc0..9bfa2ea51 100644 --- a/core/wren/tests/connectors/test_postgres.py +++ b/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: diff --git a/core/wren/tests/unit/test_postgres_semicolon_unlimited.py b/core/wren/tests/unit/test_postgres_semicolon_unlimited.py index d866b26f5..df3e8e5cb 100644 --- a/core/wren/tests/unit/test_postgres_semicolon_unlimited.py +++ b/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" )