diff --git a/core/wren/src/wren/connector/athena.py b/core/wren/src/wren/connector/athena.py index c957c017cc..a7629d4985 100644 --- a/core/wren/src/wren/connector/athena.py +++ b/core/wren/src/wren/connector/athena.py @@ -16,7 +16,7 @@ import pyarrow as pa -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError # Athena's DB-API cursor returns Trino-style type names. We delegate the @@ -298,6 +298,7 @@ def __init__(self, connection_info): self.connection = connect(**_build_connect_kwargs(connection_info)) def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) # Push LIMIT into Athena when requested so Presto/Trino-flavoured # engines can stop early instead of us downloading a full result and # slicing in Python. Subquery-wrap + trailing-semicolon strip keeps @@ -308,7 +309,7 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table: # is terminated by the newline instead of swallowing the closing # `) AS _wren_sub LIMIT n`. (Single-line sibling connectors don't # guard this.) - executed = f"SELECT * FROM (\n{executed}\n) AS _wren_sub LIMIT {int(limit)}" + executed = f"SELECT * FROM (\n{executed}\n) AS _wren_sub LIMIT {limit}" try: with contextlib.closing(self.connection.cursor()) as cursor: cursor.execute(executed) diff --git a/core/wren/src/wren/connector/base.py b/core/wren/src/wren/connector/base.py index c483f74ada..13d5731914 100644 --- a/core/wren/src/wren/connector/base.py +++ b/core/wren/src/wren/connector/base.py @@ -20,6 +20,44 @@ def strip_trailing_semicolon(sql: str) -> str: return _TRAILING_SEMICOLONS_RE.sub("", sql) +def coerce_limit(limit: int | None) -> int | None: + """Validate and coerce a user-supplied ``limit`` to a non-negative ``int``. + + ``ConnectorABC.query`` is typed as ``limit: int | None``. At runtime this + helper still defends against accidental non-ints so every connector that + interpolates LIMIT shares one contract: + + - ``None`` stays unlimited + - ``bool`` is rejected (``bool`` is an ``int`` subclass) + - non-integral numbers (e.g. ``-0.5``, ``1.5``, ``Decimal("1.5")``) are + rejected — never truncated by ``int()`` + - non-numeric / overflow values raise ``ValueError`` + - negatives raise ``ValueError`` + + Invalid limits surface as a consistent ``ValueError`` instead of a + driver-level error after SQL interpolation. + """ + if limit is None: + return None + if isinstance(limit, bool): + raise ValueError("limit must be an integer, not bool") + if isinstance(limit, float): + if not limit.is_integer(): + raise ValueError(f"limit must be an integral value, got {limit!r}") + try: + coerced = int(limit) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"limit must be an integer, got {limit!r}") from exc + # Reject truncation (Decimal/Fraction/etc.) via direct equality — no float() + # so oversized ints do not surface OverflowError. Integral floats already + # passed is_integer() above; plain int/str need no extra check. + if not isinstance(limit, (int, float, str)) and limit != coerced: + raise ValueError(f"limit must be an integral value, got {limit!r}") + if coerced < 0: + raise ValueError(f"limit must be non-negative, got {coerced}") + return coerced + + class ConnectorABC(ABC): @abstractmethod def query(self, sql: str, limit: int | None = None) -> pa.Table: diff --git a/core/wren/src/wren/connector/bigquery.py b/core/wren/src/wren/connector/bigquery.py index 5562cf5979..4dd5115d67 100644 --- a/core/wren/src/wren/connector/bigquery.py +++ b/core/wren/src/wren/connector/bigquery.py @@ -4,7 +4,7 @@ import pyarrow as pa from loguru import logger -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon def _apply_limit(sql: str, limit: int) -> str: @@ -18,7 +18,7 @@ def _apply_limit(sql: str, limit: int) -> str: Avoids comment-sensitive outer-LIMIT detection heuristics. """ cleaned = strip_trailing_semicolon(sql) - return f"SELECT * FROM ({cleaned}) AS _sub LIMIT {int(limit)}" + return f"SELECT * FROM ({cleaned}) AS _sub LIMIT {limit}" class BigQueryConnector(ConnectorABC): @@ -51,6 +51,7 @@ def __init__(self, connection_info): self.connection = client def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) if limit is not None: sql = _apply_limit(sql, limit) else: diff --git a/core/wren/src/wren/connector/canner.py b/core/wren/src/wren/connector/canner.py index 2cbd611b00..5d273c1d6e 100644 --- a/core/wren/src/wren/connector/canner.py +++ b/core/wren/src/wren/connector/canner.py @@ -17,7 +17,7 @@ import pyarrow as pa from loguru import logger -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError # Postgres OID → Arrow type. Canner publishes Trino-style values over the @@ -242,6 +242,7 @@ def __init__(self, connection_info): self._closed = False def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) import psycopg # noqa: PLC0415 # Always strip a trailing statement terminator. Unlimited queries still diff --git a/core/wren/src/wren/connector/clickhouse.py b/core/wren/src/wren/connector/clickhouse.py index 2a3964b627..7927114d95 100644 --- a/core/wren/src/wren/connector/clickhouse.py +++ b/core/wren/src/wren/connector/clickhouse.py @@ -20,7 +20,7 @@ from loguru import logger from sqlglot.expressions import DataType -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model.error import ( DIALECT_SQL, DatabaseTimeoutError, @@ -392,6 +392,7 @@ def __init__(self, connection_info: Any): self._closed = False def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) # Strip the terminating run of ``;`` / whitespace before wrapping — # ``SELECT * FROM (SELECT 1;) AS _wren_sub LIMIT N`` is invalid SQL. # Semicolons inside string literals are preserved. diff --git a/core/wren/src/wren/connector/datafusion.py b/core/wren/src/wren/connector/datafusion.py index 25c6afeafa..747f02ff38 100644 --- a/core/wren/src/wren/connector/datafusion.py +++ b/core/wren/src/wren/connector/datafusion.py @@ -7,7 +7,7 @@ import pyarrow.ipc as ipc from loguru import logger -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model import DataFusionConnectionInfo from wren.model.error import ErrorCode, WrenError @@ -29,11 +29,12 @@ def __init__(self, connection_info: DataFusionConnectionInfo): self._register_tables() def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) + stripped = strip_trailing_semicolon(sql) if limit is not None: - sql = ( - f"SELECT * FROM ({strip_trailing_semicolon(sql)}) " - f"AS _q LIMIT {int(limit)}" - ) + sql = f"SELECT * FROM ({stripped}) AS _q LIMIT {limit}" + else: + sql = stripped ipc_bytes = self.ctx.query(sql) reader = ipc.open_stream(io.BytesIO(bytes(ipc_bytes))) return reader.read_all() diff --git a/core/wren/src/wren/connector/duckdb.py b/core/wren/src/wren/connector/duckdb.py index f5769c7a33..7f3feca0bc 100644 --- a/core/wren/src/wren/connector/duckdb.py +++ b/core/wren/src/wren/connector/duckdb.py @@ -4,7 +4,7 @@ import pyarrow as pa from loguru import logger -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model import ( GcsFileConnectionInfo, MinioFileConnectionInfo, @@ -81,10 +81,11 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table: Trailing statement terminators are always stripped so client-pasted ``SELECT …;`` behaves the same on limited and unlimited paths. """ + limit = coerce_limit(limit) stripped = strip_trailing_semicolon(sql) if limit is not None: # Subquery wrap rejects an interior terminator after strip. - sql = f"SELECT * FROM ({stripped}) AS _q LIMIT {int(limit)}" + sql = f"SELECT * FROM ({stripped}) AS _q LIMIT {limit}" else: sql = stripped return self.connection.execute(sql).fetch_arrow_table() diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index 439b08387c..f6ccdc2838 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -20,6 +20,7 @@ from wren.connector.base import ( ConnectorABC, + coerce_limit, strip_trailing_semicolon, ) from wren.model.data_source import DataSource @@ -41,20 +42,6 @@ def _apply_limit(sql: str, limit: int) -> str: return f"{strip_trailing_semicolon(sql)}\nLIMIT {limit}" -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 MySqlConnector(ConnectorABC): """Native MySQLdb connector that bypasses ibis-project.""" @@ -89,7 +76,7 @@ def __init__(self, connection_info): raise def query(self, sql: str, limit: int | None = None) -> pa.Table: - limit = _coerce_limit(limit) + limit = coerce_limit(limit) if limit is not None: sql = _apply_limit(sql, limit) else: diff --git a/core/wren/src/wren/connector/oracle.py b/core/wren/src/wren/connector/oracle.py index 363c7bf59a..f120f817c8 100644 --- a/core/wren/src/wren/connector/oracle.py +++ b/core/wren/src/wren/connector/oracle.py @@ -14,7 +14,7 @@ except ImportError: # pragma: no cover oracledb = None -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError @@ -178,6 +178,7 @@ def __init__(self, connection_info): self.connection = _make_oracle_connection(connection_info) def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) # Always strip terminating `;` even on the unlimited path: a bare # trailing semicolon is rejected by some Oracle clients/drivers even # though engines accept multi-statement scripts elsewhere. diff --git a/core/wren/src/wren/connector/postgres.py b/core/wren/src/wren/connector/postgres.py index 2df0b17d61..8d952fa3a0 100644 --- a/core/wren/src/wren/connector/postgres.py +++ b/core/wren/src/wren/connector/postgres.py @@ -20,7 +20,7 @@ import pyarrow as pa from loguru import logger -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError # Map of well-known PostgreSQL OIDs to Arrow types. OIDs that we have not @@ -249,6 +249,7 @@ def __init__(self, connection_info): self._closed = False def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) # Strip terminating ``;`` even when no LIMIT wrapper is applied so # client-pasted statements match dry_run / limited composition rules. sql = strip_trailing_semicolon(sql) diff --git a/core/wren/src/wren/connector/redshift.py b/core/wren/src/wren/connector/redshift.py index 4ac76b4940..2c3d9fd4ce 100644 --- a/core/wren/src/wren/connector/redshift.py +++ b/core/wren/src/wren/connector/redshift.py @@ -4,7 +4,7 @@ import pyarrow as pa from loguru import logger -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model import ( RedshiftConnectionInfo, RedshiftConnectionUnion, @@ -44,11 +44,9 @@ def __init__(self, connection_info: RedshiftConnectionUnion): self.connection.autocommit = True def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) if limit is not None: - sql = ( - f"SELECT * FROM ({strip_trailing_semicolon(sql)}) " - f"AS _q LIMIT {int(limit)}" - ) + sql = f"SELECT * FROM ({strip_trailing_semicolon(sql)}) AS _q LIMIT {limit}" else: # Unlimited path also rejects trailing ``;`` for single statements # depending on driver/session settings — strip for consistency. diff --git a/core/wren/src/wren/connector/snowflake.py b/core/wren/src/wren/connector/snowflake.py index 37bbdfb48f..842c0d5478 100644 --- a/core/wren/src/wren/connector/snowflake.py +++ b/core/wren/src/wren/connector/snowflake.py @@ -4,7 +4,7 @@ import pyarrow as pa -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError @@ -55,6 +55,7 @@ def __init__(self, connection_info): self.connection = make_snowflake_connection(connection_info) def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) # Push LIMIT into Snowflake when requested so we do not download a # full result set only to slice it in Python. Wrap as a subquery so a # trailing semicolon in the user SQL cannot break composition, and so @@ -67,7 +68,7 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table: executed = ( "SELECT * FROM (\n" f"{strip_trailing_semicolon(sql)}\n" - f") AS _wren_sub LIMIT {int(limit)}" + f") AS _wren_sub LIMIT {limit}" ) try: with self.connection.cursor() as cursor: diff --git a/core/wren/src/wren/connector/trino.py b/core/wren/src/wren/connector/trino.py index 05f3360085..1cb08cdf18 100644 --- a/core/wren/src/wren/connector/trino.py +++ b/core/wren/src/wren/connector/trino.py @@ -19,7 +19,7 @@ from loguru import logger from sqlglot.expressions import ColumnDef, DataType -from wren.connector.base import ConnectorABC, strip_trailing_semicolon +from wren.connector.base import ConnectorABC, coerce_limit, strip_trailing_semicolon from wren.model.error import ( DIALECT_SQL, ErrorCode, @@ -482,6 +482,7 @@ def __init__(self, connection_info): self._closed = False def query(self, sql: str, limit: int | None = None) -> pa.Table: + limit = coerce_limit(limit) trino = _import_trino() if limit is not None: diff --git a/core/wren/tests/unit/test_coerce_limit.py b/core/wren/tests/unit/test_coerce_limit.py new file mode 100644 index 0000000000..8ff5bef184 --- /dev/null +++ b/core/wren/tests/unit/test_coerce_limit.py @@ -0,0 +1,71 @@ +"""Unit tests for wren.connector.base.coerce_limit.""" + +from __future__ import annotations + +from decimal import Decimal +from fractions import Fraction + +import pytest + +from wren.connector.base import coerce_limit + + +def test_none_passthrough() -> None: + assert coerce_limit(None) is None + + +def test_accepts_int() -> None: + assert coerce_limit(10) == 10 + assert coerce_limit(0) == 0 + + +def test_rejects_bool() -> None: + with pytest.raises(ValueError, match="bool"): + coerce_limit(True) # type: ignore[arg-type] + with pytest.raises(ValueError, match="bool"): + coerce_limit(False) # type: ignore[arg-type] + + +def test_rejects_non_integral_float() -> None: + with pytest.raises(ValueError, match="integral"): + coerce_limit(-0.5) # type: ignore[arg-type] + with pytest.raises(ValueError, match="integral"): + coerce_limit(1.5) # type: ignore[arg-type] + + +def test_accepts_integral_float() -> None: + assert coerce_limit(2.0) == 2 # type: ignore[arg-type] + + +def test_rejects_injection_string() -> None: + with pytest.raises(ValueError): + coerce_limit("1; DROP TABLE foo") # type: ignore[arg-type] + + +def test_rejects_negative() -> None: + with pytest.raises(ValueError, match="non-negative"): + coerce_limit(-3) + + +def test_rejects_non_numeric() -> None: + with pytest.raises(ValueError): + coerce_limit(object()) # type: ignore[arg-type] + + +def test_rejects_non_integral_decimal_and_fraction() -> None: + with pytest.raises(ValueError, match="integral"): + coerce_limit(Decimal("1.5")) # type: ignore[arg-type] + with pytest.raises(ValueError, match="integral"): + coerce_limit(Fraction(3, 2)) # type: ignore[arg-type] + with pytest.raises(ValueError): + coerce_limit(Decimal("-0.5")) # type: ignore[arg-type] + + +def test_accepts_integral_decimal() -> None: + assert coerce_limit(Decimal("2.0")) == 2 # type: ignore[arg-type] + + +def test_preserves_oversized_int_without_float_overflow() -> None: + huge = 10**400 + # Preserve arbitrary-size integers without converting through float. + assert coerce_limit(huge) == huge diff --git a/core/wren/tests/unit/test_datafusion_semicolon.py b/core/wren/tests/unit/test_datafusion_semicolon.py index 224e489f69..4dd4b90ce7 100644 --- a/core/wren/tests/unit/test_datafusion_semicolon.py +++ b/core/wren/tests/unit/test_datafusion_semicolon.py @@ -49,8 +49,9 @@ def test_query_without_limit_is_unwrapped() -> None: connector, ctx = _make_mock_connector() connector.query("SELECT 1;") (sent,), _ = ctx.query.call_args - # No limit -> no subquery wrapping; passed through verbatim. - assert sent == "SELECT 1;" + # No limit -> no subquery wrapping; still strip trailing terminator + # so DataFusion does not treat the SQL as a multi-statement batch. + assert sent == "SELECT 1" def test_helper_preserves_semicolon_inside_string_literal() -> None: diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index 633958b9d1..06120c14bc 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -10,12 +10,12 @@ import pyarrow as pa import pytest +from wren.connector.base import coerce_limit from wren.connector.mysql import ( _apply_limit, _arrow_decimal_from_mysql_field, _build_mysql_column, _build_mysql_connect_kwargs, - _coerce_limit, _mysql_blob_codes, _mysql_decimal_codes, _mysql_field_type_map, @@ -41,31 +41,43 @@ def __init__(self, url: str, kwargs: dict[str, str] | None = None) -> None: self.kwargs = kwargs -# ── _coerce_limit ───────────────────────────────────────────────────────── +# ── coerce_limit (shared base helper; mysql private removed) ───────────── def test_coerce_limit_none_passthrough() -> None: - assert _coerce_limit(None) is None + assert coerce_limit(None) is None def test_coerce_limit_accepts_int() -> None: - assert _coerce_limit(10) == 10 + assert coerce_limit(10) == 10 def test_coerce_limit_accepts_numeric_string() -> None: # ``int()`` accepts numeric strings — keep that contract. - assert _coerce_limit("25") == 25 + assert coerce_limit("25") == 25 def test_coerce_limit_rejects_injection_string() -> None: """A crafted limit value must not survive ``int()`` coercion.""" with pytest.raises(ValueError): - _coerce_limit("1; DROP TABLE foo") + coerce_limit("1; DROP TABLE foo") def test_coerce_limit_rejects_negative() -> None: with pytest.raises(ValueError): - _coerce_limit(-1) + coerce_limit(-1) + + +def test_coerce_limit_rejects_fractional_decimal_and_fraction() -> None: + from decimal import Decimal + from fractions import Fraction + + with pytest.raises(ValueError, match="integral"): + coerce_limit(Decimal("1.5")) + with pytest.raises(ValueError, match="integral"): + coerce_limit(Fraction(3, 2)) + with pytest.raises(ValueError): + coerce_limit(Decimal("-0.5")) # ── _apply_limit ──────────────────────────────────────────────────────────