fix(postgres): strip trailing semicolon before subquery-wrapping (#2407)

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

@@ -13,6 +13,7 @@ It avoids the ibis-framework dependency entirely, which gives us:
 from __future__ import annotations
 
 import json
+import re
 from decimal import ROUND_HALF_EVEN
 from decimal import Decimal as PyDecimal
 
@@ -23,6 +24,22 @@ from loguru import logger
 from wren.connector.base import ConnectorABC
 from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError
 
+_TRAILING_SEMICOLONS_RE = re.compile(r"[;\s]+\Z")
+
+
+def _strip_trailing_semicolon(sql: str) -> str:
+    """Strip any trailing ``;`` characters and surrounding whitespace.
+
+    Wrapping user SQL as ``SELECT * FROM ({sql}) AS _sub LIMIT N`` breaks when
+    ``sql`` ends in a semicolon — Postgres rejects ``SELECT 1;`` inside a
+    subquery (``syntax error at or near ";"``). Only the *terminating* run of
+    semicolons/whitespace is stripped, so semicolons inside string literals
+    (e.g. ``SELECT 'a;b' FROM t``) are preserved. Mirrors the canner/trino/
+    clickhouse connectors, which already strip before subquery-wrapping.
+    """
+    return _TRAILING_SEMICOLONS_RE.sub("", sql)
+
+
 # Map of well-known PostgreSQL OIDs to Arrow types. OIDs that we have not
 # explicitly mapped fall back to ``pa.string()`` (see ``_get_pg_arrow_type``).
 _PG_OID_TO_ARROW: dict[int, pa.DataType] = {
@@ -250,7 +267,7 @@ class PostgresConnector(ConnectorABC):
 
     def query(self, sql: str, limit: int | None = None) -> pa.Table:
         if limit is not None:
-            sql = f"SELECT * FROM ({sql}) AS _sub LIMIT {limit}"
+            sql = f"SELECT * FROM ({_strip_trailing_semicolon(sql)}) AS _sub LIMIT {limit}"
 
         try:
             with self.connection.cursor() as cursor:
@@ -269,7 +286,7 @@ class PostgresConnector(ConnectorABC):
             ) from e
 
     def dry_run(self, sql: str) -> None:
-        wrapped = f"SELECT * FROM ({sql}) AS _sub LIMIT 0"
+        wrapped = f"SELECT * FROM ({_strip_trailing_semicolon(sql)}) AS _sub LIMIT 0"
         try:
             with self.connection.cursor() as cursor:
                 cursor.execute(wrapped)

core/wren/tests/connectors/test_postgres.py

@@ -250,3 +250,59 @@ class TestPostgresConnectorTypes:
         assert [field.name for field in result.schema] == ["id", "id"]
         assert result.column(0).to_pylist() == [42]
         assert result.column(1).to_pylist() == [42]
+
+
+# ---------------------------------------------------------------------------
+# Trailing-semicolon stripping (mocked — no live database).
+#
+# These live in the postgres-marked module so they actually run in the
+# ``postgres tests`` CI job (which installs psycopg). They use a mocked
+# connection and assert on the SQL the connector executes, so no container
+# is required.
+# ---------------------------------------------------------------------------
+
+from contextlib import contextmanager
+from unittest.mock import MagicMock
+
+from wren.connector.postgres import _strip_trailing_semicolon
+
+
+def _make_mock_connector() -> tuple[PostgresConnector, MagicMock]:
+    """Build a PostgresConnector bypassing __init__ (no real connection)."""
+    connector = PostgresConnector.__new__(PostgresConnector)
+    connector._closed = False
+    cursor = MagicMock()
+    cursor.description = None  # _build_pg_arrow_table returns an empty table
+
+    @contextmanager
+    def _cursor_cm():
+        yield cursor
+
+    conn = MagicMock()
+    conn.cursor.side_effect = _cursor_cm
+    connector.connection = conn
+    return connector, cursor
+
+
+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
+
+
+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"
+
+
+def test_helper_preserves_semicolon_inside_string_literal() -> None:
+    sql = "SELECT 'a;b' AS x"
+    assert _strip_trailing_semicolon(sql) == sql
+
+
+def test_helper_no_trailing_semicolon_unchanged() -> None:
+    assert _strip_trailing_semicolon("SELECT 1") == "SELECT 1"