fix(pybind): keep the connection alive across a Ctrl-C guarded call

query(), execute() and query_as_arrow() now hold the connection state for
the whole call, so the ClientContext the SIGINT handler interrupts outlives
the guard even if another thread closes the connection. Before, a close()
from another thread destroyed the Connection while the main thread was
still inside the query, after waiting for it with the GIL held.

The handler's signal safety is pinned: static_asserts that both atomics it
touches are lock-free, and a comment that ClientContext::interrupt() must
stay a store to an atomic flag. Comments also state that prepare() is not
guarded and why the handler swap is not wrapped in a signal mask. A failure
of the main-thread/handler check is reported as unraisable, except during
interpreter shutdown, instead of being swallowed; the query still runs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: 2 people

src_cpp/py_connection.cpp

@@ -9,6 +9,10 @@
 #include <csignal>
 #endif
 
+#if PY_VERSION_HEX < 0x030D0000
+#define Py_IsFinalizing _Py_IsFinalizing
+#endif
+
 #include "cached_import/py_cached_import.h"
 #include "common/constants.h"
 #include "common/exception/not_implemented.h"
@@ -193,7 +197,14 @@ namespace {
 // Python runs its SIGINT handler only between bytecodes, so not while the main thread waits in
 // the engine. During such a wait, SIGINT interrupts the connection's query and then calls
 // Python's handler. This is only done on the main thread while Python's default handler
-// (KeyboardInterrupt) is active, so custom handlers keep their behaviour.
+// (KeyboardInterrupt) is active, so custom handlers keep their behaviour. Only query(),
+// execute() and query_as_arrow() install it.
+//
+// target is set before the handler is installed and cleared after the previous one is
+// restored, so a SIGINT in between reaches Python's handler, as without this guard. A signal
+// mask would not help: masks are per thread, and a process-directed SIGINT runs the handler on
+// whichever thread does not block it. The statics are per process; this single-phase-init
+// module cannot be imported into subinterpreters.
 class SigintInterruptsQuery {
 public:
     explicit SigintInterruptsQuery([[maybe_unused]] main::ClientContext* context) {
@@ -246,11 +257,17 @@ class SigintInterruptsQuery {
             }
             auto& signal = importCache->signal;
             return signal.getsignal()(signal.sigint()).is(signal.default_int_handler());
-        } catch (py::error_already_set&) {
+        } catch (py::error_already_set& e) {
+            if (!Py_IsFinalizing()) {
+                e.discard_as_unraisable("checking the SIGINT handler before a query");
+            }
             return false;
         }
     }
 
+    // Runs in signal context, so it may only use async-signal-safe operations: lock-free atomics
+    // and Python's own handler, which is written for signal context. ClientContext::interrupt()
+    // is a store to an std::atomic<bool> in the engine and must stay one.
     static void handle(int signum, siginfo_t* info, void* ucontext) {
         if (auto* context = target.load()) {
             context->interrupt();
@@ -266,6 +283,8 @@ class SigintInterruptsQuery {
 
     static inline std::atomic<main::ClientContext*> target{nullptr};
     static inline std::atomic<bool> received{false};
+    static_assert(std::atomic<main::ClientContext*>::is_always_lock_free);
+    static_assert(std::atomic<bool>::is_always_lock_free);
     static inline struct sigaction previous {};
     bool installed = false;
 #endif
@@ -280,40 +299,47 @@ std::unique_ptr<PyQueryResult> PyConnection::execute(PyPreparedStatement* prepar
         throw RuntimeException("Prepared statement is closed.");
     }
     auto parameters = transformPythonParameters(params, &stateRef.ref());
+    // Keeps the ClientContext alive past the guard, even if another thread closes this
+    // connection while the GIL is released.
+    auto keepAlive = state;
     SigintInterruptsQuery sigint(stateRef.ref().getClientContext());
     py::gil_scoped_release release;
     auto queryResult =
         stateRef.ref().executeWithParams(&preparedStatement->state->ref(), std::move(parameters));
     py::gil_scoped_acquire acquire;
     sigint.raiseIfInterrupted();
-    return checkAndWrapQueryResult(queryResult, state);
+    return checkAndWrapQueryResult(queryResult, keepAlive);
 }
 
 std::unique_ptr<PyQueryResult> PyConnection::query(const std::string& statement) {
     auto& stateRef = refState();
+    auto keepAlive = state;
     SigintInterruptsQuery sigint(stateRef.ref().getClientContext());
     py::gil_scoped_release release;
     auto queryResult = stateRef.ref().query(statement);
     py::gil_scoped_acquire acquire;
     sigint.raiseIfInterrupted();
-    return checkAndWrapQueryResult(queryResult, state);
+    return checkAndWrapQueryResult(queryResult, keepAlive);
 }
 
 std::unique_ptr<PyQueryResult> PyConnection::queryAsArrow(const std::string& statement,
     int64_t chunkSize) {
     auto& stateRef = refState();
+    auto keepAlive = state;
     SigintInterruptsQuery sigint(stateRef.ref().getClientContext());
     py::gil_scoped_release release;
     auto queryResult = stateRef.ref().queryAsArrow(statement, chunkSize);
     py::gil_scoped_acquire acquire;
     sigint.raiseIfInterrupted();
-    return checkAndWrapQueryResult(queryResult, state);
+    return checkAndWrapQueryResult(queryResult, keepAlive);
 }
 
 void PyConnection::setMaxNumThreadForExec(uint64_t numThreads) {
     refState().ref().setMaxNumThreadForExec(numThreads);
 }
 
+// Not guarded by SigintInterruptsQuery: a Ctrl-C while the statement compiles takes effect when
+// prepare() returns.
 PyPreparedStatement PyConnection::prepare(const std::string& query, const py::dict& parameters) {
     auto& stateRef = refState();
     auto params = transformPythonParameters(parameters, &stateRef.ref());

test/capi_xfails.py

@@ -73,6 +73,8 @@
         "test/test_keyboard_interrupt.py::test_keyboard_interrupt_stops_the_query[execute_with_parameters]",
         "test/test_keyboard_interrupt.py::test_keyboard_interrupt_stops_the_query[query]",
         "test/test_keyboard_interrupt.py::test_keyboard_interrupt_stops_the_query[query_as_arrow]",
+        "test/test_keyboard_interrupt.py::test_keyboard_interrupt_after_close_from_another_thread",
+        "test/test_keyboard_interrupt.py::test_failed_handler_check_is_reported_not_raised",
         # UDF registration is still routed through pybind.
         "test/test_blob_parameter.py::test_bytes_param_udf",
         "test/test_udf.py::test_udf",

test/test_keyboard_interrupt.py

@@ -77,3 +77,49 @@ def test_custom_sigint_handler_still_runs(tmp_path: Path) -> None:
     stdout = run_script(tmp_path, CUSTOM_HANDLER)
     # A custom handler is left alone: the query finishes, then the handler runs.
     assert "OUTCOME finished handled" in stdout, stdout
+
+
+CLOSE_DURING_QUERY = """
+threading.Timer(0.5, conn._connection.close).start()
+started = time.monotonic()
+try:
+    conn.execute(LONG_QUERY.replace('$n', '40000'))
+    print("OUTCOME finished")
+except KeyboardInterrupt:
+    print("OUTCOME KeyboardInterrupt")
+print("ELAPSED", time.monotonic() - started)
+"""
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT is POSIX only")
+def test_keyboard_interrupt_after_close_from_another_thread(tmp_path: Path) -> None:
+    # The running call keeps the connection alive, so closing it from another thread
+    # returns at once instead of waiting for the query with the GIL held. (Whether the
+    # context was freed under the running query would need ASan to observe.)
+    stdout = run_script(tmp_path, CLOSE_DURING_QUERY)
+    assert "OUTCOME KeyboardInterrupt" in stdout, stdout
+    assert float(stdout.split("ELAPSED")[1].split()[0]) < 10, stdout
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT is POSIX only")
+def test_failed_handler_check_is_reported_not_raised(tmp_path: Path) -> None:
+    script = """
+import signal, sys
+import ladybug as lb
+
+def broken(signum):
+    raise RuntimeError("getsignal is broken")
+
+signal.getsignal = broken
+conn = lb.Connection(lb.Database(sys.argv[1]))
+print("RESULT", conn.execute("RETURN 1").get_next()[0])
+"""
+    proc = subprocess.run(
+        [sys.executable, "-c", script, str(tmp_path / "db")],
+        capture_output=True,
+        text=True,
+        timeout=120,
+    )
+    assert proc.returncode == 0, proc.stderr[-2000:]
+    assert "RESULT 1" in proc.stdout, proc.stdout
+    assert "getsignal is broken" in proc.stderr, proc.stderr