fix(connection): keep the UDF close() error when the UDF runs on the query thread

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:

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