diff --git a/core/wren/src/wren/connector/base.py b/core/wren/src/wren/connector/base.py index c483f74ada..3e6c4f1a4d 100644 --- a/core/wren/src/wren/connector/base.py +++ b/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 diff --git a/core/wren/src/wren/connector/duckdb.py b/core/wren/src/wren/connector/duckdb.py index f5769c7a33..b15ee59dd4 100644 --- a/core/wren/src/wren/connector/duckdb.py +++ b/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() diff --git a/core/wren/tests/unit/test_duckdb_coerce_limit.py b/core/wren/tests/unit/test_duckdb_coerce_limit.py new file mode 100644 index 0000000000..30919d5e5a --- /dev/null +++ b/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