From c076fc3a264888c9e38d3f32c7a41bbc918e5d87 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 1 Oct 2015 15:26:13 +0100
Subject: [PATCH] The wait_select callback can cancel a query using Ctrl-C

Fixes #333.
---
 NEWS               |  3 +++
 doc/src/extras.rst |  3 +++
 doc/src/faq.rst    | 31 +++++++++++++++++++++++++++++++
 lib/extras.py      | 23 ++++++++++++++---------
 4 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/NEWS b/NEWS
index 3c34b6b8..b2fc084a 100644
--- a/NEWS
+++ b/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`).
diff --git a/doc/src/extras.rst b/doc/src/extras.rst
index 36ef0132..0e21ae58 100644
--- a/doc/src/extras.rst
+++ b/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.
diff --git a/doc/src/faq.rst b/doc/src/faq.rst
index 8f2f1ecc..69273ba5 100644
--- a/doc/src/faq.rst
+++ b/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
diff --git a/lib/extras.py b/lib/extras.py
index c9f1cbcd..2713d6fc 100644
--- a/lib/extras.py
+++ b/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()
-        if state == POLL_OK:
-            break
-        elif state == POLL_READ:
-            select.select([conn.fileno()], [], [])
-        elif state == POLL_WRITE:
-            select.select([], [conn.fileno()], [])
-        else:
-            raise conn.OperationalError("bad state from poll: %s" % state)
+        try:
+            state = conn.poll()
+            if state == POLL_OK:
+                break
+            elif state == POLL_READ:
+                select.select([conn.fileno()], [], [])
+            elif state == POLL_WRITE:
+                select.select([], [conn.fileno()], [])
+            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):