Skip to content

Commit 5d07d29

Browse files
zahariashclaude
authored andcommitted
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 <noreply@anthropic.com>
1 parent a04832b commit 5d07d29

2 files changed

Lines changed: 28 additions & 3 deletions

File tree

‎src_py/connection.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -526,7 +526,12 @@ def close(self) -> None:
526526
while the connection runs a query, close() raises instead.
527527
"""
528528
if any(x is self for x in _active_stack()):
529-
msg = "close() cannot be called from a call on the same connection."
529+
# The engine may run a UDF on the thread that issued the query, so a
530+
# UDF can find its own connection's call on this thread's stack.
531+
if _called_from_udf(self):
532+
msg = "close() cannot be called from a UDF while the connection runs a query."
533+
else:
534+
msg = "close() cannot be called from a call on the same connection."
530535
raise RuntimeError(msg)
531536

532537
with self._calls:

‎test/test_connection_close.py‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,8 @@ def closer(x: int) -> int:
5555
conn.close()
5656
return x
5757
58-
# UDFs run on engine threads, so this is not caught as a call on the same thread.
58+
# The UDF runs inside the query in flight, on an engine thread or on the one that
59+
# issued the query, so close() must name the UDF either way.
5960
conn.create_function("closer", closer)
6061
started = time.monotonic()
6162
try:
@@ -65,6 +66,25 @@ def closer(x: int) -> int:
6566
outcome = str(e)
6667
print("CLOSE", time.monotonic() - started)
6768
print("OUTCOME", outcome)
69+
""",
70+
"close_from_a_folded_udf": """
71+
conn = lb.Connection(db)
72+
73+
def closer(x: int) -> int:
74+
conn.close()
75+
return x
76+
77+
# A constant argument is folded while binding, so the UDF runs on the thread that
78+
# issued the query, below the connection's own call.
79+
conn.create_function("closer", closer)
80+
started = time.monotonic()
81+
try:
82+
conn.execute("RETURN closer(1)").get_all()
83+
outcome = "finished"
84+
except RuntimeError as e:
85+
outcome = str(e)
86+
print("CLOSE", time.monotonic() - started)
87+
print("OUTCOME", outcome)
6888
""",
6989
"query_without_parameters": """
7090
conn = lb.Connection(db)
@@ -163,7 +183,7 @@ def test_close_with_query_in_flight(tmp_path: Path, scenario: str) -> None:
163183
assert "OUTCOME" in proc.stdout
164184
if scenario == "close_from_inside_a_call":
165185
assert "same connection" in proc.stdout, proc.stdout
166-
if scenario == "close_from_a_udf":
186+
if scenario in ("close_from_a_udf", "close_from_a_folded_udf"):
167187
assert "from a UDF" in proc.stdout, proc.stdout
168188
if scenario != "write":
169189
# A running read is interrupted rather than waited for. A write is

0 commit comments

Comments
 (0)