Support for cancelling queries

2 files changed, 11 insertions, 0 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index 891ea5a4..87b8bda8 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -156,6 +156,10 @@ def cli(database, user, host, port, prompt_passwd, never_prompt):
                     _logger.debug("status: %r", status)
                     output.extend(format_output(rows, headers, status))
                 click.echo_via_pager('\n'.join(output))
+            except KeyboardInterrupt:
+                pgexecute.reconnect()
+                _logger.debug("cancelled query, sql: %r", document.text)
+                click.secho("cancelled query", err=True, fg='red')
             except Exception as e:
                 _logger.error("sql: %r, error: %r", document.text, e)
                 _logger.error("traceback: %r", traceback.format_exc())
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 0c9d45b5..e4b31d51 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -1,5 +1,6 @@
 import logging
 import psycopg2
+import psycopg2.extras
 import psycopg2.extensions
 from collections import defaultdict
 from .packages import pgspecial
@@ -11,6 +12,9 @@ _logger = logging.getLogger(__name__)
 psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
 psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
 
+# When running a query, make pressing CTRL+C raise a KeyboardInterrupt
+psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
+
 def _parse_dsn(dsn, default_user, default_password, default_host,
         default_port):
     """
@@ -70,6 +74,9 @@ class PGExecute(object):
                 _parse_dsn(database, default_user=user,
                         default_password=password, default_host=host,
                         default_port=port)
+        self.reconnect()
+
+    def reconnect(self):
         self.conn = psycopg2.connect(database=self.dbname, user=self.user,
                 password=self.password, host=self.host, port=self.port)
         self.conn.autocommit = True