Handle exit exceptions in connect

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;
 }

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):