diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c index f472c4bb..ae2f63b7 100644 --- a/psycopg/connection_type.c +++ b/psycopg/connection_type.c @@ -1395,7 +1395,15 @@ exit: PyObject *ptype = NULL, *pvalue = NULL, *ptb = NULL; PyErr_Fetch(&ptype, &pvalue, &ptb); obscure_password(self); - PyErr_Restore(ptype, pvalue, ptb); + // if we got a system exit exception (not a subclass of Exception), + // leave it as is, do not restore + if (PyErr_Occurred() && ! PyErr_ExceptionMatches(PyExc_Exception)) { + Py_XDECREF(ptype); + Py_XDECREF(pvalue); + Py_XDECREF(ptb); + } else { + PyErr_Restore(ptype, pvalue, ptb); + } } return rv; } diff --git a/tests/test_connection.py b/tests/test_connection.py index b5627075..ef3dea7f 100755 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -35,6 +35,8 @@ import subprocess as sp from collections import deque from operator import attrgetter from weakref import ref +import multiprocessing +import signal import psycopg2 import psycopg2.extras @@ -408,6 +410,26 @@ t.join() finally: shutil.rmtree(dir, ignore_errors=True) + @slow + def test_handles_keyboardinterrupt(self): + def conn(queue): + host = "10.255.255.1" # will timeout + queue.put(1) + try: + self.connect(host=host, password="x", connect_timeout=1) + except KeyboardInterrupt: + queue.put("KeyboardInterrupt") + except psycopg2.OperationalError: + queue.put("OperationalError") + + queue = multiprocessing.Queue() + process = multiprocessing.Process(target=conn, args=(queue,)) + process.start() + queue.get() + os.kill(process.pid, signal.SIGINT) + process.join() + self.assertEqual(queue.get(), "KeyboardInterrupt") + class ParseDsnTestCase(ConnectingTestCase): def test_parse_dsn(self):