refactor(connector): centralize LIMIT coercion

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, strip_trailing_semicolon, coerce_limit
 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)

core/wren/src/wren/connector/base.py

@@ -20,6 +20,48 @@ 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``) are rejected — never truncated
+    - 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}")
+        # Still route through int() below for consistency / overflow.
+    try:
+        coerced = int(limit)
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"limit must be an integer, got {limit!r}") from exc
+    # Reject values whose int() truncation would change the number (e.g. Decimal)
+    # when the original compares unequal as a number.
+    if isinstance(limit, (int, float)):
+        if float(limit) != float(coerced):
+            raise ValueError(f"limit must be an integral value, got {limit!r}")
+    else:
+        # Strings / other: require exact round-trip for numeric strings only.
+        # int("1.5") already failed; int("01") == 1 is fine.
+        pass
+    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:

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, strip_trailing_semicolon, coerce_limit
 
 
 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:

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, strip_trailing_semicolon, coerce_limit
 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

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, strip_trailing_semicolon, coerce_limit
 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.

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, strip_trailing_semicolon, coerce_limit
 from wren.model import DataFusionConnectionInfo
 from wren.model.error import ErrorCode, WrenError
 
@@ -29,11 +29,9 @@ def __init__(self, connection_info: DataFusionConnectionInfo):
         self._register_tables()
 
     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}"
         ipc_bytes = self.ctx.query(sql)
         reader = ipc.open_stream(io.BytesIO(bytes(ipc_bytes)))
         return reader.read_all()

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, strip_trailing_semicolon, coerce_limit
 from wren.model import (
     GcsFileConnectionInfo,
     MinioFileConnectionInfo,
@@ -73,6 +73,7 @@ def __init__(self, connection_info):
             raise
 
     def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        limit = coerce_limit(limit)
         """Execute ``sql`` and return the result as an Arrow table.
 
         When ``limit`` is provided the query is wrapped in a ``LIMIT`` clause
@@ -84,7 +85,7 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table:
         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()

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:

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, strip_trailing_semicolon, coerce_limit
 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.

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, strip_trailing_semicolon, coerce_limit
 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)

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, strip_trailing_semicolon, coerce_limit
 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.

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, strip_trailing_semicolon, coerce_limit
 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:

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, strip_trailing_semicolon, coerce_limit
 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:

core/wren/tests/unit/test_coerce_limit.py

@@ -0,0 +1,49 @@
+"""Unit tests for wren.connector.base.coerce_limit."""
+
+from __future__ import annotations
+
+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]