diff --git a/core/wren/src/wren/connector/trino.py b/core/wren/src/wren/connector/trino.py index 05f3360085..661d574c17 100644 --- a/core/wren/src/wren/connector/trino.py +++ b/core/wren/src/wren/connector/trino.py @@ -456,6 +456,20 @@ def _import_trino(): ) from e +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 TrinoConnector(ConnectorABC): """Native trino DB-API connector that bypasses ibis-project.""" @@ -482,6 +496,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: diff --git a/core/wren/tests/unit/test_trino_coerce_limit.py b/core/wren/tests/unit/test_trino_coerce_limit.py new file mode 100644 index 0000000000..53d0cf3c72 --- /dev/null +++ b/core/wren/tests/unit/test_trino_coerce_limit.py @@ -0,0 +1,45 @@ +"""Trino query must coerce LIMIT before SQL interpolation.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from wren.connector import trino as trino_mod + + +def _connector(): + c = object.__new__(trino_mod.TrinoConnector) + 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_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() + cursor = MagicMock() + c.connection.cursor.return_value = cursor + with patch.object(trino_mod, "_import_trino") as imp, patch.object( + trino_mod, "_build_trino_arrow_table", return_value="tbl" + ): + trino_pkg = MagicMock() + trino_pkg.exceptions.TrinoQueryError = type("TrinoQueryError", (Exception,), {}) + imp.return_value = trino_pkg + out = c.query("SELECT 1", limit="4") + assert out == "tbl" + executed = cursor.execute.call_args[0][0] + assert "LIMIT 4" in executed + assert "DROP" not in executed