diff --git a/core/wren/src/wren/connector/postgres.py b/core/wren/src/wren/connector/postgres.py index 2df0b17d61..cd3c2115b0 100644 --- a/core/wren/src/wren/connector/postgres.py +++ b/core/wren/src/wren/connector/postgres.py @@ -220,6 +220,20 @@ def _coerce_decimal(value: PyDecimal | None, target_type: pa.DataType): return pa.array(values, type=arrow_type, from_pandas=True) +def _coerce_limit(limit: int | None) -> int | None: + """Validate and coerce a user-supplied ``limit`` to a non-negative ``int``. + + ``int(limit)`` rejects strings like ``"5 OR 1=1"`` so the value can be + safely interpolated into SQL. Negative limits are also rejected. + """ + if limit is None: + return None + coerced = int(limit) + if coerced < 0: + raise ValueError(f"limit must be non-negative, got {coerced}") + return coerced + + class PostgresConnector(ConnectorABC): """Native psycopg3 implementation of the Wren postgres connector.""" @@ -251,6 +265,7 @@ def __init__(self, connection_info): def query(self, sql: str, limit: int | None = None) -> pa.Table: # Strip terminating ``;`` even when no LIMIT wrapper is applied so # client-pasted statements match dry_run / limited composition rules. + limit = _coerce_limit(limit) sql = strip_trailing_semicolon(sql) if limit is not None: sql = f"SELECT * FROM ({sql}) AS _sub LIMIT {limit}" diff --git a/core/wren/tests/unit/test_postgres_coerce_limit.py b/core/wren/tests/unit/test_postgres_coerce_limit.py new file mode 100644 index 0000000000..c20e6e6d25 --- /dev/null +++ b/core/wren/tests/unit/test_postgres_coerce_limit.py @@ -0,0 +1,61 @@ +"""Postgres query must coerce LIMIT before SQL interpolation. + +``postgres`` imports ``psycopg`` at module load. Unit CI does not install the +``postgres`` extra, so stub ``psycopg`` in ``sys.modules`` before importing the +connector module (same pattern as ``test_postgres_semicolon_unlimited``). +""" + +from __future__ import annotations + +import sys +import types +from unittest.mock import MagicMock, patch + +import pytest + + +def _ensure_psycopg_stub() -> None: + if "psycopg" in sys.modules: + return + mod = types.ModuleType("psycopg") + errors = types.ModuleType("psycopg.errors") + sys.modules["psycopg"] = mod + sys.modules["psycopg.errors"] = errors + + +_ensure_psycopg_stub() + +from wren.connector import postgres as pg_mod # noqa: E402 + + +def _connector(): + c = object.__new__(pg_mod.PostgresConnector) + c.connection = MagicMock() + c._closed = False + return c + + +def test_reject_negative_limit(): + c = _connector() + with pytest.raises(ValueError, match="non-negative"): + c.query("SELECT 1", limit=-1) + + +def test_reject_injection_string(): + c = _connector() + with pytest.raises(ValueError): + c.query("SELECT 1", limit="1; DROP TABLE t") + + +def test_numeric_string_limit_interpolated(): + c = _connector() + cursor = MagicMock() + cursor.__enter__ = MagicMock(return_value=cursor) + cursor.__exit__ = MagicMock(return_value=False) + c.connection.cursor.return_value = cursor + with patch.object(pg_mod, "_build_pg_arrow_table", return_value="tbl"): + out = c.query("SELECT 1", limit="3") + assert out == "tbl" + executed = cursor.execute.call_args[0][0] + assert "LIMIT 3" in executed + assert "DROP" not in executed