diff --git a/core/wren/src/wren/connector/clickhouse.py b/core/wren/src/wren/connector/clickhouse.py index 2a3964b627..3ce4fc1941 100644 --- a/core/wren/src/wren/connector/clickhouse.py +++ b/core/wren/src/wren/connector/clickhouse.py @@ -10,6 +10,7 @@ from __future__ import annotations import json +import numbers from decimal import Decimal as PyDecimal from typing import Any from urllib.parse import parse_qsl, unquote, urlparse @@ -383,6 +384,24 @@ def _build_clickhouse_client_kwargs(connection_info: Any) -> dict: # -------------------------------------------------------------------------- +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 + # Reject negatives *before* ``int()`` truncation so a fractional value like + # ``-0.5`` (float, Decimal, Fraction, ...) cannot slip through as ``LIMIT 0``. + if isinstance(limit, numbers.Number) and not isinstance(limit, bool) and limit < 0: + raise ValueError(f"limit must be non-negative, got {limit}") + coerced = int(limit) + if coerced < 0: + raise ValueError(f"limit must be non-negative, got {coerced}") + return coerced + + class ClickHouseConnector(ConnectorABC): """Native ``clickhouse-connect`` connector that bypasses ``ibis-project``.""" @@ -395,6 +414,7 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table: # 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. + limit = _coerce_limit(limit) stripped = strip_trailing_semicolon(sql) statement = stripped if limit is not None: diff --git a/core/wren/tests/unit/test_clickhouse_coerce_limit.py b/core/wren/tests/unit/test_clickhouse_coerce_limit.py new file mode 100644 index 0000000000..14c7db820b --- /dev/null +++ b/core/wren/tests/unit/test_clickhouse_coerce_limit.py @@ -0,0 +1,62 @@ +"""ClickHouse query must coerce LIMIT before SQL interpolation.""" + +from __future__ import annotations + +from decimal import Decimal +from fractions import Fraction +from unittest.mock import MagicMock, patch + +import pytest + +from wren.connector import clickhouse as ch_mod + + +def _connector(): + c = object.__new__(ch_mod.ClickHouseConnector) + c.connection = MagicMock() + c._closed = False + return c + + +def test_reject_negative_limit(): + c = _connector() + with pytest.raises(ValueError, match="non-negative"): + c.query("SELECT 1", limit=-1) + + +def test_reject_fractional_negative_limit(): + # ``int(-0.5)`` truncates to ``0``; ensure the fractional negative is + # rejected before truncation rather than silently becoming ``LIMIT 0``. + c = _connector() + with pytest.raises(ValueError, match="non-negative"): + c.query("SELECT 1", limit=-0.5) + + +def test_reject_negative_decimal_limit(): + # ``Decimal`` and ``Fraction`` are ``numbers.Number`` too; a negative + # fractional value must be rejected before ``int()`` truncation. + c = _connector() + with pytest.raises(ValueError, match="non-negative"): + c.query("SELECT 1", limit=Decimal("-0.5")) + + +def test_reject_negative_fraction_limit(): + c = _connector() + with pytest.raises(ValueError, match="non-negative"): + c.query("SELECT 1", limit=Fraction(-1, 2)) + + +def test_reject_injection_string(): + c = _connector() + with pytest.raises(ValueError): + c.query("SELECT 1", limit="1; DROP TABLE t") + + +def test_numeric_string_limit_interpolated(): + c = _connector() + with patch.object(ch_mod, "_build_clickhouse_arrow_table", return_value="tbl"): + out = c.query("SELECT 1", limit="7") + assert out == "tbl" + statement = c.connection.query.call_args[0][0] + assert "LIMIT 7" in statement + assert "DROP" not in statement