From c8e9703b28a0efa670716f4671f89621e9796e23 Mon Sep 17 00:00:00 2001 From: Darek Date: Sun, 18 Oct 2026 11:14:35 +0200 Subject: [PATCH] fix(connection): keep the UDF close() error when the UDF runs on the query thread close() called from a UDF while its connection runs a query raised "close() cannot be called from a call on the same connection" instead of the UDF message whenever the UDF ran on the thread that issued the query: that thread's stack holds the connection's own call below the UDF marker. A constant argument folded while binding already runs the UDF there, and an engine that runs a one-thread task on the calling thread does so for every such task. The same-thread check now uses the UDF message when a UDF was entered after the connection's call, as the worker path does. The new close_from_a_folded_udf scenario covers it with RETURN closer(1). Co-Authored-By: Claude Opus 5.5 --- src_py/connection.py | 7 ++++++- test/test_connection_close.py | 24 ++++++++++++++++++++++-- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src_py/connection.py b/src_py/connection.py index c881922..d298387 100644 --- a/src_py/connection.py +++ b/src_py/connection.py @@ -491,7 +491,12 @@ def close(self) -> None: while the connection runs a query, close() raises instead. """ if any(x is self for x in _active_stack()): - msg = "close() cannot be called from a call on the same connection." + # The engine may run a UDF on the thread that issued the query, so a + # UDF can find its own connection's call on this thread's stack. + if _called_from_udf(self): + msg = "close() cannot be called from a UDF while the connection runs a query." + else: + msg = "close() cannot be called from a call on the same connection." raise RuntimeError(msg) with self._calls: diff --git a/test/test_connection_close.py b/test/test_connection_close.py index 4c62685..e8be3e0 100644 --- a/test/test_connection_close.py +++ b/test/test_connection_close.py @@ -55,7 +55,8 @@ def closer(x: int) -> int: conn.close() return x -# UDFs run on engine threads, so this is not caught as a call on the same thread. +# The UDF runs inside the query in flight, on an engine thread or on the one that +# issued the query, so close() must name the UDF either way. conn.create_function("closer", closer) started = time.monotonic() try: @@ -65,6 +66,25 @@ def closer(x: int) -> int: outcome = str(e) print("CLOSE", time.monotonic() - started) print("OUTCOME", outcome) +""", + "close_from_a_folded_udf": """ +conn = lb.Connection(db) + +def closer(x: int) -> int: + conn.close() + return x + +# A constant argument is folded while binding, so the UDF runs on the thread that +# issued the query, below the connection's own call. +conn.create_function("closer", closer) +started = time.monotonic() +try: + conn.execute("RETURN closer(1)").get_all() + outcome = "finished" +except RuntimeError as e: + outcome = str(e) +print("CLOSE", time.monotonic() - started) +print("OUTCOME", outcome) """, "query_without_parameters": """ conn = lb.Connection(db) @@ -163,7 +183,7 @@ def test_close_with_query_in_flight(tmp_path: Path, scenario: str) -> None: assert "OUTCOME" in proc.stdout if scenario == "close_from_inside_a_call": assert "same connection" in proc.stdout, proc.stdout - if scenario == "close_from_a_udf": + if scenario in ("close_from_a_udf", "close_from_a_folded_udf"): assert "from a UDF" in proc.stdout, proc.stdout if scenario != "write": # A running read is interrupted rather than waited for. A write is