diff --git a/core/wren/src/wren/connector/duckdb.py b/core/wren/src/wren/connector/duckdb.py index 4819323304..ea22916cd7 100644 --- a/core/wren/src/wren/connector/duckdb.py +++ b/core/wren/src/wren/connector/duckdb.py @@ -1,4 +1,5 @@ import os +from threading import Lock import opendal import pyarrow as pa @@ -51,11 +52,15 @@ def _init_duckdb_gcs(connection, info: GcsFileConnectionInfo): class DuckDBConnector(ConnectorABC): def __init__(self, connection_info): + """Open a DuckDB session and serialize operations sharing its results.""" import duckdb # noqa: PLC0415 from duckdb import HTTPException, IOException # noqa: PLC0415 self._HTTPException = HTTPException self._IOException = IOException + # One in-flight call per connector; use pooled connections + # if parallel query throughput becomes necessary. + self._connection_lock = Lock() self.connection = duckdb.connect() try: @@ -91,7 +96,10 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table: sql = f"SELECT * FROM (\n{stripped}\n) AS _q LIMIT {limit}" else: sql = stripped - return self.connection.execute(sql).fetch_arrow_table() + # Pending results belong to the connection. Keep execute and fetch + # together so another call cannot overwrite or consume this result. + with self._connection_lock: + return self.connection.execute(sql).fetch_arrow_table() def dry_run(self, sql: str) -> None: """Validate ``sql`` without returning rows or side effects. @@ -105,7 +113,8 @@ def dry_run(self, sql: str) -> None: no rows are materialized. """ stripped = self._strip(sql) - self.connection.execute(f"SELECT * FROM (\n{stripped}\n) AS _q LIMIT 0") + with self._connection_lock: + self.connection.execute(f"SELECT * FROM (\n{stripped}\n) AS _q LIMIT 0") def _attach_database(self, connection_info) -> None: """Attach every discovered DuckDB file as a read-only database. @@ -170,7 +179,8 @@ def _list_duckdb_files(self, connection_info) -> list[str]: def close(self) -> None: """Close the underlying DuckDB connection, logging any error.""" try: - self.connection.close() + with self._connection_lock: + self.connection.close() except Exception as e: logger.warning(f"Error closing DuckDB connection: {e}") diff --git a/core/wren/tests/unit/test_duckdb_concurrency.py b/core/wren/tests/unit/test_duckdb_concurrency.py new file mode 100644 index 0000000000..a4bb6ba826 --- /dev/null +++ b/core/wren/tests/unit/test_duckdb_concurrency.py @@ -0,0 +1,145 @@ +"""Concurrent calls must not overwrite another call's DuckDB result.""" + +from concurrent.futures import ThreadPoolExecutor +from threading import Event + +import duckdb +import pytest + +from wren.connector.duckdb import DuckDBConnector +from wren.model import LocalFileConnectionInfo + +pytestmark = pytest.mark.unit + + +@pytest.fixture +def connector(tmp_path): + """Yield a real in-memory connector and close it after each test.""" + instance = DuckDBConnector( + LocalFileConnectionInfo(url=str(tmp_path), format="parquet") + ) + yield instance + instance.close() + + +class _InterleavedConnection: + """Pause after real execution so the second call runs before the fetch.""" + + def __init__(self, connection, executed, resume, second_executed): + """Record events controlling execution before the first result fetch.""" + self.connection = connection + self.executed = executed + self.resume = resume + self.second_executed = second_executed + + def execute(self, sql): + """Execute real SQL and pause the first query before returning its result.""" + result = self.connection.execute(sql) + if "AS first_result" in sql: + self.executed.set() + assert self.resume.wait(10), "second call did not finish" + if "AS second_result" in sql: + self.second_executed.set() + return result + + def close(self): + """Forward closing to the real DuckDB connection.""" + self.connection.close() + + +@pytest.mark.parametrize("second_operation", ["query", "dry_run"]) +def test_execute_and_fetch_are_isolated(connector, second_operation): + """Keep a competing query or dry run blocked until the first result is fetched.""" + executed, resume, second_executed = Event(), Event(), Event() + connector.connection = _InterleavedConnection( + connector.connection, executed, resume, second_executed + ) + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(connector.query, "SELECT 11 AS first_result") + try: + assert executed.wait(10), "first query did not execute" + second = pool.submit( + getattr(connector, second_operation), "SELECT 22 AS second_result" + ) + # On the broken shared connection this lets the second operation + # replace the first result. With serialization it waits for resume. + assert not second_executed.wait(1), "second operation executed before fetch" + finally: + resume.set() + assert first.result().to_pylist() == [{"first_result": 11}] + if second_operation == "query": + assert second.result().to_pylist() == [{"second_result": 22}] + else: + assert second.result() is None + + +def test_parallel_queries_share_memory_tables_and_views(connector): + """Preserve each query's ID while sharing in-memory and temporary objects.""" + connector.connection.execute("CREATE TABLE numbers AS SELECT 7 AS value") + connector.connection.execute( + "CREATE TEMP TABLE temp_numbers AS SELECT * FROM numbers" + ) + connector.connection.execute( + "CREATE TEMP VIEW number_view AS SELECT * FROM temp_numbers" + ) + with ThreadPoolExecutor(max_workers=4) as pool: + futures = [ + pool.submit( + connector.query, + f"SELECT {i} AS query_id, value FROM number_view", + ) + for i in range(16) + ] + for i, future in enumerate(futures): + assert future.result().to_pylist() == [{"query_id": i, "value": 7}] + + +def test_queries_preserve_attached_tables_and_settings(connector, tmp_path): + """Keep session state and release the lock after execution or validation errors.""" + db_path = tmp_path / "sample.duckdb" + with duckdb.connect(str(db_path)) as source: + source.execute("CREATE TABLE numbers AS SELECT 42 AS value") + connector._attach_database( + LocalFileConnectionInfo(url=str(tmp_path), format="duckdb") + ) + connector.connection.execute("SET TimeZone='Asia/Tokyo'") + assert connector.query( + "SELECT value, current_setting('TimeZone') AS timezone FROM sample.numbers" + ).to_pylist() == [{"value": 42, "timezone": "Asia/Tokyo"}] + connector.dry_run("SELECT * FROM sample.numbers") + with pytest.raises(duckdb.Error): + connector.query("SELECT * FROM missing_table") + with pytest.raises(duckdb.Error): + connector.dry_run("SELECT * FROM missing_table") + assert connector.query("SELECT 1 AS value").to_pylist() == [{"value": 1}] + connector.close() + with pytest.raises(duckdb.Error): + connector.query("SELECT 1") + + +def test_close_waits_for_query_to_fetch(connector): + """Block closing until a paused query has fetched its result.""" + executed, resume, close_started, close_finished = Event(), Event(), Event(), Event() + connector.connection = _InterleavedConnection( + connector.connection, executed, resume, Event() + ) + + def close(): + """Signal the start and completion of a competing close operation.""" + close_started.set() + connector.close() + close_finished.set() + + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(connector.query, "SELECT 11 AS first_result") + try: + assert executed.wait(10), "first query did not execute" + closing = pool.submit(close) + assert close_started.wait(10), "close did not start" + assert not close_finished.wait(1), "connection closed before fetch" + finally: + resume.set() + assert first.result().to_pylist() == [{"first_result": 11}] + closing.result() + with pytest.raises(duckdb.Error): + connector.query("SELECT 1") diff --git a/core/wren/tests/unit/test_duckdb_file_listing.py b/core/wren/tests/unit/test_duckdb_file_listing.py index fe9f99f12c..9e50a31df8 100644 --- a/core/wren/tests/unit/test_duckdb_file_listing.py +++ b/core/wren/tests/unit/test_duckdb_file_listing.py @@ -1,5 +1,6 @@ """Regression test: DuckDB file discovery is case-insensitive on extension.""" +from threading import Lock from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -86,10 +87,12 @@ def test_attach_database_uniquifies_case_colliding_aliases(): def test_query_strips_trailing_semicolon_before_limit_wrap(): + """Remove statement terminators before wrapping a limited query.""" # A semicolon-terminated statement must not produce invalid SQL such as # ``SELECT * FROM (SELECT 1;) AS _q LIMIT 5`` when a limit is applied. connector = DuckDBConnector.__new__(DuckDBConnector) connector.connection = MagicMock() + connector._connection_lock = Lock() connector.connection.execute.return_value.fetch_arrow_table.return_value = "tbl" result = connector.query("SELECT 1;", limit=5) @@ -100,12 +103,14 @@ def test_query_strips_trailing_semicolon_before_limit_wrap(): def test_dry_run_wraps_in_limit_zero_subquery(): + """Keep every statement inside the validation subquery.""" # dry_run neutralizes multi-statement input by wrapping in a LIMIT 0 # subquery (matching the other connectors) rather than pre-rejecting it, # so no rows materialize and any trailing statement becomes a natural # syntax error inside the subquery. connector = DuckDBConnector.__new__(DuckDBConnector) connector.connection = MagicMock() + connector._connection_lock = Lock() connector.dry_run("SELECT 1; DROP TABLE t;") @@ -116,8 +121,10 @@ def test_dry_run_wraps_in_limit_zero_subquery(): def test_dry_run_strips_trailing_semicolon(): + """Remove the final terminator before wrapping SQL for validation.""" connector = DuckDBConnector.__new__(DuckDBConnector) connector.connection = MagicMock() + connector._connection_lock = Lock() connector.dry_run("SELECT 1;") @@ -126,10 +133,12 @@ def test_dry_run_strips_trailing_semicolon(): def test_dry_run_preserves_semicolon_in_string_literal(): + """Preserve literal semicolons when building the validation subquery.""" # A single valid statement with a semicolon inside a string literal must # not be mangled or falsely rejected. connector = DuckDBConnector.__new__(DuckDBConnector) connector.connection = MagicMock() + connector._connection_lock = Lock() connector.dry_run("SELECT ';' AS x") diff --git a/core/wren/tests/unit/test_duckdb_semicolon_unlimited.py b/core/wren/tests/unit/test_duckdb_semicolon_unlimited.py index 1436f1a35f..a74f109018 100644 --- a/core/wren/tests/unit/test_duckdb_semicolon_unlimited.py +++ b/core/wren/tests/unit/test_duckdb_semicolon_unlimited.py @@ -2,6 +2,7 @@ from __future__ import annotations +from threading import Lock from unittest.mock import MagicMock import pyarrow as pa @@ -13,10 +14,12 @@ def test_unlimited_query_strips_trailing_semicolon(): + """Strip final statement terminators even when no row limit is requested.""" connector = DuckDBConnector.__new__(DuckDBConnector) connector.connection = MagicMock() - connector.connection.execute.return_value.fetch_arrow_table.return_value = ( - pa.table({"x": [1]}) + connector._connection_lock = Lock() + connector.connection.execute.return_value.fetch_arrow_table.return_value = pa.table( + {"x": [1]} ) connector.query("SELECT 1 AS x; ") @@ -25,10 +28,12 @@ def test_unlimited_query_strips_trailing_semicolon(): def test_limited_query_still_wraps_after_strip(): + """Apply the row limit after removing the trailing terminator.""" connector = DuckDBConnector.__new__(DuckDBConnector) connector.connection = MagicMock() - connector.connection.execute.return_value.fetch_arrow_table.return_value = ( - pa.table({"x": [1]}) + connector._connection_lock = Lock() + connector.connection.execute.return_value.fetch_arrow_table.return_value = pa.table( + {"x": [1]} ) connector.query("SELECT 1 AS x;", limit=2) @@ -42,8 +47,9 @@ def test_query_limit_survives_trailing_line_comment(): """Trailing `--` must not eat the wrap (same shape as postgres.py #2728).""" connector = DuckDBConnector.__new__(DuckDBConnector) connector.connection = MagicMock() - connector.connection.execute.return_value.fetch_arrow_table.return_value = ( - pa.table({"x": [1]}) + connector._connection_lock = Lock() + connector.connection.execute.return_value.fetch_arrow_table.return_value = pa.table( + {"x": [1]} ) connector.query("SELECT 1 AS x -- pick", limit=2) @@ -54,8 +60,10 @@ def test_query_limit_survives_trailing_line_comment(): def test_dry_run_survives_trailing_line_comment(): + """Keep the validation wrapper outside a trailing SQL line comment.""" connector = DuckDBConnector.__new__(DuckDBConnector) connector.connection = MagicMock() + connector._connection_lock = Lock() connector.dry_run("SELECT 1 AS x -- pick")