The wait_select callback can cancel a query using Ctrl-C

Fixes #333.

NEWS

@@ -5,6 +5,9 @@ What's new in psycopg 2.6.2
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 - Report the server response status on errors (such as :ticket:`#281`).
+- The `~psycopg2.extras.wait_select` callback allows interrupting a
+  long-running query in an interactive shell using :kbd:`Ctrl-C`
+  (:ticket:`#333`).
 - Raise `!NotSupportedError` on unhandled server response status
   (:ticket:`#352`).
 - Fixed `!PersistentConnectionPool` on Python 3 (:ticket:`#348`).

doc/src/extras.rst

@@ -611,3 +611,6 @@ Coroutine support
 
 .. autofunction:: wait_select(conn)
 
+    .. versionchanged:: 2.6.2
+        allow to cancel a query using :kbd:`Ctrl-C`, see
+        :ref:`the FAQ <faq-interrupt-query>` for an example.

doc/src/faq.rst

@@ -223,6 +223,37 @@ What are the advantages or disadvantages of using named cursors?
     little memory on the client and to skip or discard parts of the result set.
 
 
+.. _faq-interrupt-query:
+.. cssclass:: faq
+
+How do I interrupt a long-running query in an interactive shell?
+    Normally the interactive shell becomes unresponsive to :kbd:`Ctrl-C` when
+    running a query. Using a connection in green mode allows Python to
+    receive and handle the interrupt, although it may leave the connection
+    broken, if the async callback doesn't handle the `!KeyboardInterrupt`
+    correctly.
+
+    Starting from psycopg 2.6.2, the `~psycopg2.extras.wait_select` callback
+    can handle a :kbd:`Ctrl-C` correctly. For previous versions, you can use
+    `this implementation`__.
+
+    .. __: http://initd.org/psycopg/articles/2014/07/20/cancelling-postgresql-statements-python/
+
+    .. code-block:: pycon
+
+        >>> psycopg2.extensions.set_wait_callback(psycopg2.extensions.wait_select)
+        >>> cnn = psycopg2.connect('')
+        >>> cur = cnn.cursor()
+        >>> cur.execute("select pg_sleep(10)")
+        ^C
+        Traceback (most recent call last):
+          File "<stdin>", line 1, in <module>
+          QueryCanceledError: canceling statement due to user request
+
+        >>> cnn.rollback()
+        >>> # You can use the connection and cursor again from here
+
+
 .. _faq-compile:
 
 Problems compiling and deploying psycopg2

lib/extras.py

@@ -575,15 +575,20 @@ def wait_select(conn):
     from psycopg2.extensions import POLL_OK, POLL_READ, POLL_WRITE
 
     while 1:
-        state = conn.poll()
+        try:
-        if state == POLL_OK:
+            state = conn.poll()
-            break
+            if state == POLL_OK:
-        elif state == POLL_READ:
+                break
-            select.select([conn.fileno()], [], [])
+            elif state == POLL_READ:
-        elif state == POLL_WRITE:
+                select.select([conn.fileno()], [], [])
-            select.select([], [conn.fileno()], [])
+            elif state == POLL_WRITE:
-        else:
+                select.select([], [conn.fileno()], [])
-            raise conn.OperationalError("bad state from poll: %s" % state)
+            else:
+                raise conn.OperationalError("bad state from poll: %s" % state)
+        except KeyboardInterrupt:
+            conn.cancel()
+            # the loop will be broken by a server error
+            continue
 
 
 def _solve_conn_curs(conn_or_curs):