fix(connection): make close() and UDF callbacks safe while a query runs

src_py/async_connection.py

@@ -225,6 +225,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,9 +485,39 @@ 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 any(x is self for x in _active_stack()):
+            msg = "close() cannot be called from a call on the same connection."
+            raise RuntimeError(msg)
+
+        with self._calls:
+            # Idempotent under concurrent close(): check-and-set inside the
+            # lock so two threads cannot both run the teardown tail.
+            if self.is_closed or self._closing:
+                return
+            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
+            try:
+                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)
+            except BaseException:
+                # Don't brick the connection if the wait is aborted
+                # (KeyboardInterrupt, ...): teardown never ran, so allow reuse.
+                self._closing = False
+                self._calls.notify_all()
+                raise
 
         for query_result in list(self._query_results):
             query_result.close()
@@ -814,11 +925,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 +946,7 @@ def _maybe_raise_scan_unsupported_object(self, query: str) -> None:
         )
         raise RuntimeError(msg)
 
+    @_tracks_call
     def execute(
         self,
         query: str | PreparedStatement,
@@ -932,7 +1042,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 +1060,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 +1082,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 +1181,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 +1202,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 +1240,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 +1274,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 +1293,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 +1333,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 +1362,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",