From 81dc499b3afe82eb3eb269b9e8af6866f07e465f Mon Sep 17 00:00:00 2001 From: Bartok9 Date: Fri, 31 Jul 2026 02:08:52 -0400 Subject: [PATCH 1/4] fix(duckdb): coerce LIMIT before SQL interpolation Validate limit before embedding it in the subquery LIMIT wrap; reject negatives and non-numeric strings. --- core/wren/src/wren/connector/duckdb.py | 14 +++++++- .../tests/unit/test_duckdb_coerce_limit.py | 33 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) create mode 100644 core/wren/tests/unit/test_duckdb_coerce_limit.py diff --git a/core/wren/src/wren/connector/duckdb.py b/core/wren/src/wren/connector/duckdb.py index f5769c7a33..9b54e9f1d9 100644 --- a/core/wren/src/wren/connector/duckdb.py +++ b/core/wren/src/wren/connector/duckdb.py @@ -13,6 +13,17 @@ from wren.model.error import ErrorCode, WrenError +def _coerce_limit(limit: int | None) -> int | None: + """Validate limit before SQL LIMIT interpolation.""" + if limit is None: + return None + 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("'", "''") @@ -81,10 +92,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..dac07d23f9 --- /dev/null +++ b/core/wren/tests/unit/test_duckdb_coerce_limit.py @@ -0,0 +1,33 @@ +"""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_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" From 72c65f3f286698890b53c982abf4f1f53af14cc3 Mon Sep 17 00:00:00 2001 From: Bartok9 Date: Fri, 31 Jul 2026 02:14:26 -0400 Subject: [PATCH 2/4] style(duckdb): apply ruff format --- core/wren/src/wren/connector/duckdb.py | 1 - 1 file changed, 1 deletion(-) diff --git a/core/wren/src/wren/connector/duckdb.py b/core/wren/src/wren/connector/duckdb.py index 9b54e9f1d9..f12544c554 100644 --- a/core/wren/src/wren/connector/duckdb.py +++ b/core/wren/src/wren/connector/duckdb.py @@ -23,7 +23,6 @@ def _coerce_limit(limit: int | None) -> int | None: return coerced - def _escape_sql(value: str) -> str: return value.replace("'", "''") From 6507f03884b3be9af75fb0c61df21f4473777c27 Mon Sep 17 00:00:00 2001 From: Bartok9 Date: Fri, 31 Jul 2026 02:16:51 -0400 Subject: [PATCH 3/4] fix(duckdb): reject fractional limits, align limit type, cover unlimited path - _coerce_limit rejects non-integral floats (e.g. -0.5) and bools - annotate limit as int|str|None across helper, query, and ConnectorABC - add tests for None passthrough and unlimited query semicolon strip --- core/wren/src/wren/connector/base.py | 2 +- core/wren/src/wren/connector/duckdb.py | 20 +++++++++++++---- .../tests/unit/test_duckdb_coerce_limit.py | 22 +++++++++++++++++++ 3 files changed, 39 insertions(+), 5 deletions(-) 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 f12544c554..b15ee59dd4 100644 --- a/core/wren/src/wren/connector/duckdb.py +++ b/core/wren/src/wren/connector/duckdb.py @@ -13,11 +13,23 @@ from wren.model.error import ErrorCode, WrenError -def _coerce_limit(limit: int | None) -> int | None: - """Validate limit before SQL LIMIT interpolation.""" +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 - coerced = int(limit) + 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 @@ -82,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 diff --git a/core/wren/tests/unit/test_duckdb_coerce_limit.py b/core/wren/tests/unit/test_duckdb_coerce_limit.py index dac07d23f9..30919d5e5a 100644 --- a/core/wren/tests/unit/test_duckdb_coerce_limit.py +++ b/core/wren/tests/unit/test_duckdb_coerce_limit.py @@ -21,6 +21,15 @@ def test_coerce_limit_rejects_negative() -> None: _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() @@ -31,3 +40,16 @@ def test_query_interpolates_coerced_limit() -> None: 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 From 6c9b541fb76c0ce50903990bda35f9842819b27f Mon Sep 17 00:00:00 2001 From: Bartok9 Date: Fri, 31 Jul 2026 02:17:07 -0400 Subject: [PATCH 4/4] style: ruff format