fix(connection): make close() safe while a query is running

Closing a connection while another thread was executing a query on it
destroyed the native connection under the running query. Depending on the
query this hung the process, aborted it with heap corruption, or segfaulted
(a write). AsyncConnection.close() hit the same path, since it closes its
connections before shutting down its executor.

Connection now counts calls in flight in every method that reaches the
native connection. close() refuses new calls, interrupts a running query
(repeating the interrupt until the call returns, because an interrupt that
arrives while the query is still compiling is cleared when execution
starts), and waits for in-flight calls to leave the engine before releasing
the native connection. Calls made after close() raise "Connection is
closed." The guard lives in the Python layer, so it covers both backends.

close() called from inside one of the connection's own calls on the same
thread raises instead of waiting for itself. A Python UDF that calls into a
connection while it runs a query (close(), execute() or any other call)
raises too: UDFs run inside the query in flight, so such calls used to wait
for themselves forever, on main as well. A per-thread stack of active calls
and UDFs tells these apart from the connection's own nested calls, and
covers UDFs folded on the calling thread while binding. UDFs are shared by
all connections of a database and the driver cannot tell which query runs
one, so a UDF may not call any connection that has a call in flight;
create_function() documents this. interrupt() also reaches the pybind
connection used as a fallback by the C-API backend, and is a no-op on a
closed connection. The scan-object lookup for LOAD FROM no longer assumes a
fixed call depth and walks up to the first frame outside this module.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: 2 people

src_py/async_connection.py

@@ -275,6 +275,8 @@ def close(self) -> None:
         """
         Close all connections and shutdown the thread pool.
 
+        Queries still running are interrupted and fail with "Interrupted.".
+
         Note: Call to this method is optional. The connections and thread pool
         will be closed automatically when the instance is garbage collected.
         """

src_py/connection.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
 import contextlib
+import functools
 import inspect
 import json
 import re
 import threading
 import uuid
 import warnings
 import weakref
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 from weakref import WeakSet
 
 from ._backend import get_capi_module, get_pybind_module
@@ -308,6 +309,82 @@ def _pybind_param_signature(parameters: dict[str, Any]) -> tuple:
     )
 
 
+_Method = TypeVar("_Method", bound="Callable[..., Any]")
+
+# Per thread: the connections whose calls are running, innermost last, and a
+# marker for each Python UDF entered on top of them.
+_active = threading.local()
+_UDF = object()
+
+
+def _active_stack() -> list[object]:
+    stack: list[object] | None = getattr(_active, "stack", None)
+    if stack is None:
+        stack = _active.stack = []
+    return stack
+
+
+def _called_from_udf(conn: Connection) -> bool:
+    """Whether this thread is in a UDF entered after its last call on conn."""
+    for entry in reversed(_active_stack()):
+        if entry is conn:
+            return False
+        if entry is _UDF:
+            return True
+    return False
+
+
+def _tracks_call(method: _Method) -> _Method:
+    """Count the call as in flight so that close() can wait for it."""
+
+    @functools.wraps(method)
+    def wrapper(self: Connection, *args: Any, **kwargs: Any) -> Any:
+        with self._calls:
+            if self._closing:
+                msg = "Connection is closed."
+                raise RuntimeError(msg)
+            if self._calls_in_flight and _called_from_udf(self):
+                # UDFs run inside the query in flight, which holds the
+                # connection, so the call could wait for itself forever.
+                msg = "A UDF cannot use a connection while it runs a query."
+                raise RuntimeError(msg)
+            self._calls_in_flight += 1
+        stack = _active_stack()
+        stack.append(self)
+        try:
+            return method(self, *args, **kwargs)
+        finally:
+            stack.pop()
+            with self._calls:
+                self._calls_in_flight -= 1
+                self._calls.notify_all()
+
+    return cast("_Method", wrapper)
+
+
+def _stacklevel_outside_module() -> int:
+    """Return the warnings.warn() stacklevel of the first caller outside this module."""
+    frame, level = inspect.currentframe(), 0
+    while frame is not None and frame.f_code.co_filename == __file__:
+        frame, level = frame.f_back, level + 1
+    return level
+
+
+def _marks_udf_call(udf: Callable[..., Any]) -> Callable[..., Any]:
+    """Mark the thread as running a UDF so that calls into a busy connection raise."""
+
+    @functools.wraps(udf)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        stack = _active_stack()
+        stack.append(_UDF)
+        try:
+            return udf(*args, **kwargs)
+        finally:
+            stack.pop()
+
+    return wrapper
+
+
 class Connection:
     """Connection to a database."""
 
@@ -329,6 +406,9 @@ def __init__(self, database: Database, num_threads: int = 0):
         self.database = database
         self.num_threads = num_threads
         self.is_closed = False
+        self._closing = False
+        self._calls = threading.Condition()
+        self._calls_in_flight = 0
         self._prefer_pybind = False
         self._query_timeout_ms = 0
         self._query_results: WeakSet[QueryResult] = WeakSet()
@@ -379,6 +459,7 @@ def _using_pybind_backend(self) -> bool:
             self.database._use_pybind_backend and get_pybind_module() is not None
         )
 
+    @_tracks_call
     def set_max_threads_for_exec(self, num_threads: int) -> None:
         """
         Set the maximum number of threads for executing queries.
@@ -404,10 +485,32 @@ def close(self) -> None:
 
         Note: Call to this method is optional. The connection will be closed
         automatically when the object goes out of scope.
+
+        A query running on this connection in another thread is interrupted,
+        and close() returns once it has left the engine. Called from a UDF
+        while the connection runs a query, close() raises instead.
         """
         if self.is_closed:
             return
 
+        if self in _active_stack():
+            msg = "close() cannot be called from a call on the same connection."
+            raise RuntimeError(msg)
+
+        with self._calls:
+            if self._calls_in_flight and _called_from_udf(self):
+                # UDFs run inside the query in flight, so waiting for it would
+                # never end.
+                msg = "close() cannot be called from a UDF while the connection runs a query."
+                raise RuntimeError(msg)
+            self._closing = True
+            while self._calls_in_flight:
+                # An interrupt that arrives while the query is still compiling is
+                # cleared when execution starts, so repeat it until the call returns.
+                with contextlib.suppress(RuntimeError):
+                    self.interrupt()
+                self._calls.wait(timeout=0.05)
+
         for query_result in list(self._query_results):
             query_result.close()
         self._query_results.clear()
@@ -814,11 +917,9 @@ def _maybe_raise_scan_unsupported_object(self, query: str) -> None:
             return
 
         var_name = match.group(1)
-        frame = inspect.currentframe()
-        if frame is None or frame.f_back is None:
-            return
-
-        caller = frame.f_back.f_back
+        caller = inspect.currentframe()
+        while caller is not None and caller.f_code.co_filename == __file__:
+            caller = caller.f_back
         if caller is None:
             return
 
@@ -837,6 +938,7 @@ def _maybe_raise_scan_unsupported_object(self, query: str) -> None:
         )
         raise RuntimeError(msg)
 
+    @_tracks_call
     def execute(
         self,
         query: str | PreparedStatement,
@@ -932,7 +1034,7 @@ def execute(
                     warnings.warn(
                         drop_result.getErrorMessage(),
                         RuntimeWarning,
-                        stacklevel=2,
+                        stacklevel=_stacklevel_outside_module(),
                     )
             finally:
                 self._capi_scan_tables.discard(table_name)
@@ -950,6 +1052,7 @@ def execute(
             all_query_results.append(next_query_result)
         return all_query_results
 
+    @_tracks_call
     def query_as_arrow(self, query: str, chunk_size: int) -> ArrowQueryResult:
         """
         Execute a query with the native Arrow collector path.
@@ -971,6 +1074,7 @@ def query_as_arrow(self, query: str, chunk_size: int) -> ArrowQueryResult:
         self._register_query_result(current_query_result)
         return current_query_result
 
+    @_tracks_call
     def _prepare(
         self,
         query: str,
@@ -1069,6 +1173,7 @@ def _get_rel_table_names(self) -> list[dict[str, Any]]:
                 results.append({"name": name, "src": src_node, "dst": dst_node})
         return results
 
+    @_tracks_call
     def set_query_timeout(self, timeout_in_ms: int) -> None:
         """
         Set the query timeout value in ms for executing queries.
@@ -1089,8 +1194,11 @@ def interrupt(self) -> None:
 
         If there is no currently executing query, this function does nothing.
         """
-        self._connection.interrupt()
+        for native in (self._connection, self._py_connection):
+            if native is not None:
+                native.interrupt()
 
+    @_tracks_call
     def create_function(
         self,
         name: str,
@@ -1124,12 +1232,16 @@ def create_function(
         catch_exceptions: Optional[bool]
             if true, when an exception is thrown from python, the function output will be null
             Otherwise, the exception will be rethrown
+
+        The UDF must not call a connection that is running a query: such calls raise
+        RuntimeError, because the query running the UDF may be on that connection.
         """
         if params_type is None:
             params_type = []
         parsed_params_type = [x if type(x) is str else x.value for x in params_type]
         if type(return_type) is not str:
             return_type = return_type.value
+        udf = _marks_udf_call(udf)
 
         try:
             self._connection.create_function(
@@ -1154,6 +1266,7 @@ def create_function(
                 catch_exceptions=catch_exceptions,
             )
 
+    @_tracks_call
     def remove_function(self, name: str) -> None:
         """
         Remove a User Defined Function (UDF).
@@ -1172,6 +1285,7 @@ def remove_function(self, name: str) -> None:
             self._prefer_pybind = True
             py_connection.remove_function(name)
 
+    @_tracks_call
     def create_arrow_table(
         self,
         table_name: str,
@@ -1211,6 +1325,7 @@ def create_arrow_table(
             raise RuntimeError(query_result_internal.getErrorMessage())
         return QueryResult(self, query_result_internal)
 
+    @_tracks_call
     def drop_arrow_table(self, table_name: str) -> QueryResult:
         """
         Drop an Arrow memory-backed table.
@@ -1239,6 +1354,7 @@ def drop_arrow_table(self, table_name: str) -> QueryResult:
             raise RuntimeError(query_result_internal.getErrorMessage())
         return QueryResult(self, query_result_internal)
 
+    @_tracks_call
     def create_arrow_rel_table(
         self,
         table_name: str,

test/capi_xfails.py

@@ -75,6 +75,10 @@
         "test/test_udf.py::test_udf_null",
         "test/test_udf.py::test_udf_except",
         "test/test_udf.py::test_udf_remove",
+        "test/test_connection_close.py::test_close_with_query_in_flight[close_from_a_udf]",
+        "test/test_udf_reentrancy.py::test_udf_calling_a_connection[idle_other_connection]",
+        "test/test_udf_reentrancy.py::test_udf_calling_a_connection[same_connection_on_a_worker]",
+        "test/test_udf_reentrancy.py::test_udf_calling_a_connection[same_connection_while_binding]",
         # C API Arrow export for JSON types is not yet implemented.  get_as_df()
         # tests that need the pybind NPArrayWrapper path are xfailed.
         "test/test_json.py::test_get_as_df_json_scalar",