fix(duckdb): coerce LIMIT before SQL interpolation

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

@@ -22,7 +22,7 @@ def strip_trailing_semicolon(sql: str) -> str:
 
 class ConnectorABC(ABC):
     @abstractmethod
-    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+    def query(self, sql: str, limit: int | str | None = None) -> pa.Table:
         pass
 
     @abstractmethod

core/wren/src/wren/connector/duckdb.py

@@ -13,6 +13,28 @@
 from wren.model.error import ErrorCode, WrenError
 
 
+def _coerce_limit(limit: int | str | None) -> int | None:
+    """Validate a limit before SQL ``LIMIT`` interpolation.
+
+    Accepts ``None`` (unlimited), integers, and integer-valued strings.
+    Rejects negative and non-integral values (e.g. ``-0.5``) so no
+    injection-like or fractional input can reach the interpolated SQL.
+    """
+    if limit is None:
+        return None
+    if isinstance(limit, bool):
+        raise ValueError(f"limit must be an integer, got {limit!r}")
+    if isinstance(limit, float):
+        if not limit.is_integer():
+            raise ValueError(f"limit must be an integer, got {limit!r}")
+        coerced = int(limit)
+    else:
+        coerced = int(limit)
+    if coerced < 0:
+        raise ValueError(f"limit must be non-negative, got {coerced}")
+    return coerced
+
+
 def _escape_sql(value: str) -> str:
     return value.replace("'", "''")
 
@@ -72,7 +94,7 @@ def __init__(self, connection_info):
             self.connection.close()
             raise
 
-    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+    def query(self, sql: str, limit: int | str | None = None) -> pa.Table:
         """Execute ``sql`` and return the result as an Arrow table.
 
         When ``limit`` is provided the query is wrapped in a ``LIMIT`` clause
@@ -81,10 +103,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()

core/wren/tests/unit/test_duckdb_coerce_limit.py

@@ -0,0 +1,55 @@
+"""DuckDBConnector limit coercion."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from wren.connector.duckdb import DuckDBConnector, _coerce_limit
+
+pytestmark = pytest.mark.unit
+
+
+def test_coerce_limit_rejects_injection() -> None:
+    with pytest.raises(ValueError):
+        _coerce_limit("1 OR 1=1")
+
+
+def test_coerce_limit_rejects_negative() -> None:
+    with pytest.raises(ValueError):
+        _coerce_limit(-2)
+
+
+def test_coerce_limit_rejects_fractional_negative() -> None:
+    with pytest.raises(ValueError):
+        _coerce_limit(-0.5)
+
+
+def test_coerce_limit_preserves_none() -> None:
+    assert _coerce_limit(None) is None
+
+
+def test_query_interpolates_coerced_limit() -> None:
+    c = DuckDBConnector.__new__(DuckDBConnector)
+    conn = MagicMock()
+    result = MagicMock()
+    conn.execute.return_value = result
+    result.fetch_arrow_table.return_value = MagicMock()
+    c.connection = conn
+    c.query("SELECT 1;", limit="4")
+    (sent,), _ = conn.execute.call_args
+    assert sent == "SELECT * FROM (SELECT 1) AS _q LIMIT 4"
+
+
+def test_query_unlimited_strips_semicolon_without_limit() -> None:
+    c = DuckDBConnector.__new__(DuckDBConnector)
+    conn = MagicMock()
+    result = MagicMock()
+    conn.execute.return_value = result
+    result.fetch_arrow_table.return_value = MagicMock()
+    c.connection = conn
+    c.query("SELECT 1;", limit=None)
+    (sent,), _ = conn.execute.call_args
+    assert sent == "SELECT 1"
+    assert "LIMIT" not in sent