diff --git a/core/wren/src/wren/connector/datafusion.py b/core/wren/src/wren/connector/datafusion.py index 85c1952e0b..9daed46233 100644 --- a/core/wren/src/wren/connector/datafusion.py +++ b/core/wren/src/wren/connector/datafusion.py @@ -1,6 +1,7 @@ from __future__ import annotations import io +import re from pathlib import Path import pyarrow as pa @@ -11,6 +12,22 @@ from wren.model import DataFusionConnectionInfo from wren.model.error import ErrorCode, WrenError +_TRAILING_SEMICOLONS_RE = re.compile(r"[;\s]+\Z") + + +def _strip_trailing_semicolon(sql: str) -> str: + """Strip the terminating run of ``;`` characters and surrounding whitespace. + + Wrapping user SQL as ``SELECT * FROM ({sql}) AS _q LIMIT N`` breaks when + ``sql`` ends in a semicolon — DataFusion rejects ``SELECT 1;`` inside a + subquery (``sql parser error: Expected: an expression, found: ;``). Only + the *terminating* run of semicolons/whitespace is stripped, so semicolons + inside string literals (e.g. ``SELECT 'a;b' FROM t``) are preserved. + Mirrors the postgres/redshift/duckdb connectors, which already strip + before subquery-wrapping. + """ + return _TRAILING_SEMICOLONS_RE.sub("", sql) + class DataFusionConnector(ConnectorABC): """DataFusion-native connector for local file analysis. @@ -30,7 +47,10 @@ def __init__(self, connection_info: DataFusionConnectionInfo): def query(self, sql: str, limit: int | None = None) -> pa.Table: if limit is not None: - sql = f"SELECT * FROM ({sql}) AS _q LIMIT {int(limit)}" + sql = ( + f"SELECT * FROM ({_strip_trailing_semicolon(sql)}) " + f"AS _q LIMIT {int(limit)}" + ) 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 new file mode 100644 index 0000000000..48d82f1b66 --- /dev/null +++ b/core/wren/tests/unit/test_datafusion_semicolon.py @@ -0,0 +1,61 @@ +"""Trailing-semicolon stripping for the DataFusion connector (mocked ctx). + +The DataFusion connector wraps user SQL as ``SELECT * FROM (...) AS _q LIMIT N`` +when a limit is supplied. If the user SQL ends in a semicolon, DataFusion +rejects ``SELECT 1;`` inside a subquery +(``sql parser error: Expected: an expression, found: ;``). These tests use a +mocked ``ctx`` and assert on the SQL string the connector builds, so no native +runtime file registration is required. Mirrors the postgres/redshift/duckdb +connector semicolon tests. +""" + +from unittest.mock import MagicMock + +import pyarrow as pa +import pyarrow.ipc as ipc + +from wren.connector.datafusion import ( + DataFusionConnector, + _strip_trailing_semicolon, +) + + +def _make_mock_connector() -> tuple[DataFusionConnector, MagicMock]: + """Build a DataFusionConnector bypassing __init__ (no real runtime).""" + connector = DataFusionConnector.__new__(DataFusionConnector) + ctx = MagicMock() + + # ctx.query must return IPC-stream bytes that read back into a table. + empty = pa.table({"x": pa.array([], type=pa.int64())}) + sink = pa.BufferOutputStream() + with ipc.new_stream(sink, empty.schema) as writer: + writer.write_table(empty) + ctx.query.return_value = sink.getvalue().to_pybytes() + + connector.ctx = ctx + return connector, ctx + + +def test_query_strips_trailing_semicolon_before_subquery_wrap() -> None: + connector, ctx = _make_mock_connector() + connector.query("SELECT 1;", limit=5) + (sent,), _ = ctx.query.call_args + assert sent == "SELECT * FROM (SELECT 1) AS _q LIMIT 5" + assert ";)" not in sent + + +def test_query_without_limit_is_unwrapped() -> 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;" + + +def test_helper_preserves_semicolon_inside_string_literal() -> None: + sql = "SELECT 'a;b' AS x" + assert _strip_trailing_semicolon(sql) == sql + + +def test_helper_no_trailing_semicolon_unchanged() -> None: + assert _strip_trailing_semicolon("SELECT 1") == "SELECT 1"