diff --git a/core/wren/src/wren/connector/postgres.py b/core/wren/src/wren/connector/postgres.py index c206b2949..42d595d75 100644 --- a/core/wren/src/wren/connector/postgres.py +++ b/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) diff --git a/core/wren/tests/connectors/test_postgres.py b/core/wren/tests/connectors/test_postgres.py index 49ce0574a..48a9bd3ac 100644 --- a/core/wren/tests/connectors/test_postgres.py +++ b/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"