fix(connection): don't block in set_query_timeout() while a query runs

src_cpp/py_connection.cpp

@@ -37,10 +37,11 @@ void PyConnection::initialize(py::handle& m) {
         .def("query_as_arrow", &PyConnection::queryAsArrow, py::arg("statement"),
             py::arg("chunk_size"))
         .def("set_max_threads_for_exec", &PyConnection::setMaxNumThreadForExec,
-            py::arg("num_threads"))
+            py::arg("num_threads"), py::call_guard<py::gil_scoped_release>())
         .def("prepare", &PyConnection::prepare, py::arg("query"),
             py::arg("parameters") = py::dict())
-        .def("set_query_timeout", &PyConnection::setQueryTimeout, py::arg("timeout_in_ms"))
+        .def("set_query_timeout", &PyConnection::setQueryTimeout, py::arg("timeout_in_ms"),
+            py::call_guard<py::gil_scoped_release>())
         .def("interrupt", &PyConnection::interrupt)
         .def("get_num_nodes", &PyConnection::getNumNodes, py::arg("node_name"))
         .def("get_num_rels", &PyConnection::getNumRels, py::arg("rel_name"))

src_py/connection.py

@@ -4,6 +4,7 @@
 import functools
 import inspect
 import json
+import operator
 import re
 import threading
 import uuid
@@ -385,6 +386,15 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
     return wrapper
 
 
+def _setting_value(name: str, value: int) -> int:
+    """Validate a connection setting before it is queued; the engine takes a uint64."""
+    value = operator.index(value)
+    if not 0 <= value < 2**64:
+        msg = f"{name} must be between 0 and 2**64 - 1; found {value}."
+        raise ValueError(msg)
+    return value
+
+
 class Connection:
     """Connection to a database."""
 
@@ -410,7 +420,10 @@ def __init__(self, database: Database, num_threads: int = 0):
         self._calls = threading.Condition()
         self._calls_in_flight = 0
         self._prefer_pybind = False
-        self._query_timeout_ms = 0
+        # Native setter name -> value, applied before the next query.
+        self._pending_settings: dict[str, int] = {}
+        self._pending_settings_lock = threading.Lock()
+        self._apply_settings_lock = threading.Lock()
         self._query_results: WeakSet[QueryResult] = WeakSet()
         self._capi_scan_tables: set[str] = set()
         # Implicit prepared-statement cache, shared by the pybind and
@@ -464,14 +477,36 @@ def set_max_threads_for_exec(self, num_threads: int) -> None:
         """
         Set the maximum number of threads for executing queries.
 
+        The setting applies from the next query on this connection; this call
+        does not wait for a running query.
+
         Parameters
         ----------
         num_threads : int
-            Maximum number of threads to use for executing queries.
+            Maximum number of threads to use for executing queries; 0 uses the
+            database's limit.
+
+        Raises
+        ------
+        TypeError
+            If `num_threads` is not an integer.
+        ValueError
+            If `num_threads` is negative or does not fit in 64 bits.
 
         """
+        num_threads = _setting_value("num_threads", num_threads)
         self.init_connection()
-        self._connection.set_max_threads_for_exec(num_threads)
+        with self._pending_settings_lock:
+            self._pending_settings["set_max_threads_for_exec"] = num_threads
+
+    def _apply_pending_settings(self) -> None:
+        # The engine applies settings under the lock a running query holds, so
+        # they are applied here, before a query, rather than in the setters.
+        with self._apply_settings_lock:
+            with self._pending_settings_lock:
+                pending, self._pending_settings = self._pending_settings, {}
+            for setter, value in pending.items():
+                getattr(self._connection, setter)(value)
 
     def _register_query_result(self, query_result: QueryResult) -> None:
         self._query_results.add(query_result)
@@ -975,6 +1010,7 @@ def execute(
             parameters = {}
 
         self.init_connection()
+        self._apply_pending_settings()
         if not isinstance(parameters, dict):
             msg = f"Parameters must be a dict; found {type(parameters)}."
             raise RuntimeError(msg)  # noqa: TRY004
@@ -987,7 +1023,7 @@ def execute(
 
         if (
             not self._using_pybind_backend()
-            and self._query_timeout_ms > 0
+            and self._connection._query_timeout_ms > 0
             and isinstance(query, str)
             and len(re.findall(r"(?i)\bUNWIND\s+RANGE\s*\(", query)) >= 2
         ):
@@ -1068,6 +1104,7 @@ def query_as_arrow(self, query: str, chunk_size: int) -> ArrowQueryResult:
         This is the efficient path for CSR-aware Arrow export.
         """
         self.init_connection()
+        self._apply_pending_settings()
         if not self._using_pybind_backend():
             query_result_internal = self._connection.query(query)
         else:
@@ -1186,15 +1223,26 @@ def set_query_timeout(self, timeout_in_ms: int) -> None:
         """
         Set the query timeout value in ms for executing queries.
 
+        The timeout applies from the next query on this connection; this call
+        does not wait for a running query.
+
         Parameters
         ----------
         timeout_in_ms : int
             query timeout value in ms for executing queries.
 
+        Raises
+        ------
+        TypeError
+            If `timeout_in_ms` is not an integer.
+        ValueError
+            If `timeout_in_ms` is negative or does not fit in 64 bits.
+
         """
+        timeout_in_ms = _setting_value("timeout_in_ms", timeout_in_ms)
         self.init_connection()
-        self._query_timeout_ms = int(timeout_in_ms)
-        self._connection.set_query_timeout(timeout_in_ms)
+        with self._pending_settings_lock:
+            self._pending_settings["set_query_timeout"] = timeout_in_ms
 
     def interrupt(self) -> None:
         """
@@ -1317,6 +1365,7 @@ def create_arrow_table(
 
         """
         self.init_connection()
+        self._apply_pending_settings()
         try:
             query_result_internal = self._connection.create_arrow_table(
                 table_name, dataframe
@@ -1350,6 +1399,7 @@ def drop_arrow_table(self, table_name: str) -> QueryResult:
 
         """
         self.init_connection()
+        self._apply_pending_settings()
         try:
             query_result_internal = self._connection.drop_arrow_table(table_name)
         except NotImplementedError:
@@ -1411,6 +1461,7 @@ def create_arrow_rel_table(
 
         """
         self.init_connection()
+        self._apply_pending_settings()
         layout_value = (
             layout.value if isinstance(layout, ArrowRelTableLayout) else str(layout)
         ).upper()

test/test_timeout.py

@@ -1,8 +1,17 @@
 from __future__ import annotations
 
+import contextlib
+import threading
+import time
+from typing import TYPE_CHECKING
+
+import ladybug as lb
 import pytest
 from type_aliases import ConnDB
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 
 def test_timeout(conn_db_readonly: ConnDB) -> None:
     conn, _ = conn_db_readonly
@@ -11,3 +20,85 @@ def test_timeout(conn_db_readonly: ConnDB) -> None:
         conn.execute(
             "UNWIND RANGE(1,100000) AS x UNWIND RANGE(1, 100000) AS y RETURN COUNT(x + y);"
         )
+
+
+# Runs for seconds without materializing large lists.
+LONG_QUERY = "WITH range(1, 30000) AS r UNWIND r AS x UNWIND r AS y RETURN sum(x * y)"
+
+
+def _execute_until_interrupted(conn: lb.Connection, query: str) -> None:
+    with contextlib.suppress(RuntimeError):
+        conn.execute(query)
+
+
+def _interrupt_until_done(conn: lb.Connection, *threads: threading.Thread) -> None:
+    # An interrupt sent before execution starts is cleared, so repeat it.
+    for thread in threads:
+        while thread.is_alive():
+            conn.interrupt()
+            thread.join(0.05)
+
+
+@pytest.fixture
+def conn() -> Iterator[lb.Connection]:
+    # A fresh connection: the shared read-only one carries other tests' settings.
+    db = lb.Database(":memory:")
+    conn = lb.Connection(db)
+    yield conn
+    conn.close()
+    db.close()
+
+
+def test_settings_do_not_wait_for_running_query(conn: lb.Connection) -> None:
+    worker = threading.Thread(
+        target=_execute_until_interrupted, args=(conn, LONG_QUERY)
+    )
+    worker.start()
+    time.sleep(0.3)
+    started = time.monotonic()
+    conn.set_query_timeout(60_000)
+    conn.set_max_threads_for_exec(2)
+    elapsed = time.monotonic() - started
+    _interrupt_until_done(conn, worker)
+    assert elapsed < 1, elapsed
+
+
+def test_applying_settings_lets_other_threads_run(conn: lb.Connection) -> None:
+    running = threading.Thread(
+        target=_execute_until_interrupted, args=(conn, LONG_QUERY)
+    )
+    running.start()
+    time.sleep(0.3)
+    conn.set_query_timeout(60_000)
+    # This call applies the new timeout, which waits for the running query.
+    waiting = threading.Thread(
+        target=_execute_until_interrupted, args=(conn, "RETURN 1")
+    )
+    started = time.monotonic()
+    waiting.start()
+    time.sleep(0.4)
+    elapsed = time.monotonic() - started
+    _interrupt_until_done(conn, running, waiting)
+    assert elapsed < 1, elapsed
+
+
+@pytest.mark.parametrize("setter", ["set_query_timeout", "set_max_threads_for_exec"])
+@pytest.mark.parametrize(
+    ("value", "error"),
+    [("1", TypeError), (1.5, TypeError), (-1, ValueError), (2**64, ValueError)],
+)
+def test_settings_reject_invalid_values_at_once(
+    conn: lb.Connection, setter: str, value: object, error: type[Exception]
+) -> None:
+    with pytest.raises(error):
+        getattr(conn, setter)(value)
+    assert conn.execute("RETURN 1").get_all() == [[1]]
+
+
+def test_cleared_timeout_applies_to_next_query(conn: lb.Connection) -> None:
+    query = "UNWIND range(1, 10) AS x UNWIND range(1, 10) AS y RETURN count(*)"
+    conn.set_query_timeout(1000)
+    with contextlib.suppress(RuntimeError):
+        conn.execute(query)
+    conn.set_query_timeout(0)
+    assert conn.execute(query).get_all() == [[100]]