fix(pybind): let Ctrl-C interrupt a running query

src_cpp/include/cached_import/py_cached_import.h

@@ -25,6 +25,8 @@ class PythonCachedImport {
     PandasCachedItem pandas;
     PolarsCachedItem polars;
     PyarrowCachedItem pyarrow;
+    SignalCachedItem signal;
+    ThreadingCachedItem threading;
     UUIDCachedItem uuid;
 
 private:

src_cpp/include/cached_import/py_cached_modules.h

@@ -154,6 +154,24 @@ class PyarrowCachedItem : public PythonCachedItem {
     LibCachedItem lib;
 };
 
+class SignalCachedItem : public PythonCachedItem {
+public:
+    SignalCachedItem()
+        : PythonCachedItem("signal"), getsignal("getsignal", this), sigint("SIGINT", this),
+          default_int_handler("default_int_handler", this) {}
+
+    PythonCachedItem getsignal;
+    PythonCachedItem sigint;
+    PythonCachedItem default_int_handler;
+};
+
+class ThreadingCachedItem : public PythonCachedItem {
+public:
+    ThreadingCachedItem() : PythonCachedItem("threading"), main_thread("main_thread", this) {}
+
+    PythonCachedItem main_thread;
+};
+
 class UUIDCachedItem : public PythonCachedItem {
 public:
     UUIDCachedItem() : PythonCachedItem("uuid"), UUID("UUID", this) {}

src_cpp/py_connection.cpp

@@ -1,9 +1,18 @@
 #include "include/py_connection.h"
 
 #include <algorithm>
+#include <atomic>
 #include <cctype>
 #include <utility>
 
+#ifndef _WIN32
+#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"
@@ -183,41 +192,154 @@ void PyConnection::interrupt() {
 static std::unordered_map<std::string, std::unique_ptr<Value>> transformPythonParameters(
     const py::dict& params, Connection* conn);
 
+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. 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) {
+#ifndef _WIN32
+        if (target.load() != nullptr || !onMainThreadWithDefaultHandler()) {
+            return;
+        }
+        if (sigaction(SIGINT, nullptr, &previous) != 0) {
+            return;
+        }
+        struct sigaction action {};
+        action.sa_sigaction = handle;
+        action.sa_flags = SA_SIGINFO | (previous.sa_flags & SA_ONSTACK);
+        sigemptyset(&action.sa_mask);
+        received = false;
+        target = context;
+        installed = sigaction(SIGINT, &action, nullptr) == 0;
+        if (!installed) {
+            target = nullptr;
+        }
+#endif
+    }
+
+    ~SigintInterruptsQuery() {
+#ifndef _WIN32
+        if (installed) {
+            sigaction(SIGINT, &previous, nullptr);
+            target = nullptr;
+        }
+#endif
+    }
+
+    // Call with the GIL held once the query has returned.
+    void raiseIfInterrupted() const {
+#ifndef _WIN32
+        if (installed && received && PyErr_CheckSignals() != 0) {
+            throw py::error_already_set();
+        }
+#endif
+    }
+
+private:
+#ifndef _WIN32
+    static bool onMainThreadWithDefaultHandler() {
+        // The check must never fail a query, e.g. once imports stop working at shutdown.
+        try {
+            auto mainThread = importCache->threading.main_thread()();
+            if (mainThread.attr("ident").cast<unsigned long>() != PyThread_get_thread_ident()) {
+                return false;
+            }
+            auto& signal = importCache->signal;
+            return signal.getsignal()(signal.sigint()).is(signal.default_int_handler());
+        } 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();
+        }
+        received = true;
+        // Python's handler only records the signal; KeyboardInterrupt is raised later.
+        if (previous.sa_flags & SA_SIGINFO) {
+            previous.sa_sigaction(signum, info, ucontext);
+        } else if (previous.sa_handler != SIG_DFL && previous.sa_handler != SIG_IGN) {
+            previous.sa_handler(signum);
+        }
+    }
+
+    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
+};
+
+} // namespace
+
 std::unique_ptr<PyQueryResult> PyConnection::execute(PyPreparedStatement* preparedStatement,
     const py::dict& params) {
     auto& stateRef = refState();
     if (preparedStatement == nullptr || preparedStatement->state == nullptr) {
         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;
-    return checkAndWrapQueryResult(queryResult, state);
+    sigint.raiseIfInterrupted();
+    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;
-    return checkAndWrapQueryResult(queryResult, state);
+    sigint.raiseIfInterrupted();
+    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;
-    return checkAndWrapQueryResult(queryResult, state);
+    sigint.raiseIfInterrupted();
+    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

@@ -69,6 +69,12 @@
         # C API scan rewriting uses temporary Arrow-backed tables, which cannot be
         # created through a read-only connection.
         "test/test_scan_pyarrow.py::test_pyarrow_basic",
+        # Ctrl-C interrupts a running query only on the pybind backend.
+        "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

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+import textwrap
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+# Each scenario sends SIGINT to itself one second into a query that runs for many
+# seconds, and reports how the query ended.
+PRELUDE = """
+import os, signal, sys, threading, time
+import ladybug as lb
+
+db = lb.Database(sys.argv[1])
+conn = lb.Connection(db)
+LONG_QUERY = "WITH range(1, $n) AS r UNWIND r AS x UNWIND r AS y RETURN sum(x * y)"
+threading.Timer(1.0, os.kill, (os.getpid(), signal.SIGINT)).start()
+"""
+
+SCENARIOS = {
+    "query": "conn.execute(LONG_QUERY.replace('$n', '40000'))",
+    "execute_with_parameters": "conn.execute(LONG_QUERY, {'n': 40000})",
+    "query_as_arrow": "conn.query_as_arrow(LONG_QUERY.replace('$n', '40000'), 1000)",
+}
+
+RUN = """
+started = time.monotonic()
+try:
+    {call}
+    print("OUTCOME finished")
+except KeyboardInterrupt:
+    print("OUTCOME KeyboardInterrupt")
+print("ELAPSED", time.monotonic() - started)
+print("AFTER", conn.execute("RETURN 1").get_next()[0])
+"""
+
+CUSTOM_HANDLER = """
+received = []
+signal.signal(signal.SIGINT, lambda signum, frame: received.append(signum))
+result = conn.execute(LONG_QUERY.replace("$n", "20000")).get_next()[0]
+time.sleep(0.2)
+print("OUTCOME", "finished" if result else "empty", "handled" if received else "missed")
+"""
+
+
+def run_script(tmp_path: Path, body: str) -> str:
+    script = textwrap.dedent(PRELUDE) + textwrap.dedent(body)
+    try:
+        proc = subprocess.run(
+            [sys.executable, "-c", script, str(tmp_path / "db")],
+            capture_output=True,
+            text=True,
+            timeout=120,
+        )
+    except subprocess.TimeoutExpired:
+        pytest.fail("the query was not interrupted")
+    assert proc.returncode == 0, proc.stderr[-2000:]
+    return proc.stdout
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT is POSIX only")
+@pytest.mark.parametrize("scenario", sorted(SCENARIOS))
+def test_keyboard_interrupt_stops_the_query(tmp_path: Path, scenario: str) -> None:
+    stdout = run_script(tmp_path, RUN.format(call=SCENARIOS[scenario]))
+    assert "OUTCOME KeyboardInterrupt" in stdout, stdout
+    assert float(stdout.split("ELAPSED")[1].split()[0]) < 10, stdout
+    assert "AFTER 1" in stdout, stdout
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGINT is POSIX only")
+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