diff --git a/src_cpp/include/cached_import/py_cached_import.h b/src_cpp/include/cached_import/py_cached_import.h index 734550b..b85c74a 100644 --- a/src_cpp/include/cached_import/py_cached_import.h +++ b/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: diff --git a/src_cpp/include/cached_import/py_cached_modules.h b/src_cpp/include/cached_import/py_cached_modules.h index 0d18bde..dbf840e 100644 --- a/src_cpp/include/cached_import/py_cached_modules.h +++ b/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) {} diff --git a/src_cpp/py_connection.cpp b/src_cpp/py_connection.cpp index de7c09d..477151b 100644 --- a/src_cpp/py_connection.cpp +++ b/src_cpp/py_connection.cpp @@ -1,9 +1,14 @@ #include "include/py_connection.h" #include +#include #include #include +#ifndef _WIN32 +#include +#endif + #include "cached_import/py_cached_import.h" #include "common/constants.h" #include "common/exception/not_implemented.h" @@ -183,6 +188,91 @@ void PyConnection::interrupt() { static std::unordered_map> 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. +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() != 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&) { + return false; + } + } + + 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 target{nullptr}; + static inline std::atomic received{false}; + static inline struct sigaction previous {}; + bool installed = false; +#endif +}; + +} // namespace + std::unique_ptr PyConnection::execute(PyPreparedStatement* preparedStatement, const py::dict& params) { auto& stateRef = refState(); @@ -190,27 +280,33 @@ std::unique_ptr PyConnection::execute(PyPreparedStatement* prepar throw RuntimeException("Prepared statement is closed."); } auto parameters = transformPythonParameters(params, &stateRef.ref()); + 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); } std::unique_ptr PyConnection::query(const std::string& statement) { auto& stateRef = refState(); + 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); } std::unique_ptr PyConnection::queryAsArrow(const std::string& statement, int64_t chunkSize) { auto& stateRef = refState(); + 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); } diff --git a/test/capi_xfails.py b/test/capi_xfails.py index c928e6f..c47026a 100644 --- a/test/capi_xfails.py +++ b/test/capi_xfails.py @@ -69,6 +69,10 @@ # 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]", # UDF registration is still routed through pybind. "test/test_blob_parameter.py::test_bytes_param_udf", "test/test_udf.py::test_udf", diff --git a/test/test_keyboard_interrupt.py b/test/test_keyboard_interrupt.py new file mode 100644 index 0000000..3062a06 --- /dev/null +++ b/test/test_keyboard_interrupt.py @@ -0,0 +1,79 @@ +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