fix(wren): serialize shared DuckDB connection access

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}")
 

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")

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")
 

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")