diff --git a/src_cpp/py_connection.cpp b/src_cpp/py_connection.cpp index de7c09d..d7f4323 100644 --- a/src_cpp/py_connection.cpp +++ b/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()) .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()) .def("interrupt", &PyConnection::interrupt) .def("get_num_nodes", &PyConnection::getNumNodes, py::arg("node_name")) .def("get_num_rels", &PyConnection::getNumRels, py::arg("rel_name")) diff --git a/src_py/connection.py b/src_py/connection.py index c881922..7f6c34b 100644 --- a/src_py/connection.py +++ b/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() diff --git a/test/test_timeout.py b/test/test_timeout.py index f5f4ae8..bf8f091 100644 --- a/test/test_timeout.py +++ b/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]]