diff --git a/core/wren/src/wren/connector/datafusion.py b/core/wren/src/wren/connector/datafusion.py index 25c6afeafa..348a36953f 100644 --- a/core/wren/src/wren/connector/datafusion.py +++ b/core/wren/src/wren/connector/datafusion.py @@ -29,11 +29,14 @@ def __init__(self, connection_info: DataFusionConnectionInfo): self._register_tables() def query(self, sql: str, limit: int | None = None) -> pa.Table: + # Always strip terminating ``;`` / whitespace. Limited path needs it + # for the subquery wrap; unlimited ``ctx.query`` also rejects a trailing + # statement terminator on single-statement LocalRuntime execution. + cleaned = strip_trailing_semicolon(sql) if limit is not None: - sql = ( - f"SELECT * FROM ({strip_trailing_semicolon(sql)}) " - f"AS _q LIMIT {int(limit)}" - ) + sql = f"SELECT * FROM ({cleaned}) AS _q LIMIT {int(limit)}" + else: + sql = cleaned ipc_bytes = self.ctx.query(sql) reader = ipc.open_stream(io.BytesIO(bytes(ipc_bytes))) return reader.read_all() diff --git a/core/wren/tests/unit/test_datafusion_semicolon.py b/core/wren/tests/unit/test_datafusion_semicolon.py index 224e489f69..68b88fab4e 100644 --- a/core/wren/tests/unit/test_datafusion_semicolon.py +++ b/core/wren/tests/unit/test_datafusion_semicolon.py @@ -45,12 +45,21 @@ def test_query_strips_trailing_semicolon_before_subquery_wrap() -> None: assert ";)" not in sent -def test_query_without_limit_is_unwrapped() -> None: +def test_query_without_limit_strips_trailing_semicolon() -> None: connector, ctx = _make_mock_connector() connector.query("SELECT 1;") (sent,), _ = ctx.query.call_args - # No limit -> no subquery wrapping; passed through verbatim. - assert sent == "SELECT 1;" + # No limit -> no subquery wrapping, but terminating ``;`` is still stripped + # so LocalRuntime single-statement execution accepts client SQL. + assert sent == "SELECT 1" + assert not sent.endswith(";") + + +def test_query_without_limit_strips_semicolon_and_whitespace() -> None: + connector, ctx = _make_mock_connector() + connector.query("SELECT 1; \n") + (sent,), _ = ctx.query.call_args + assert sent == "SELECT 1" def test_helper_preserves_semicolon_inside_string_literal() -> None: