CTRL+C while querying update: explicitly cancel old connection

2 files changed, 7 insertions, 3 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index 87b8bda8..f88c0c48 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -157,7 +157,8 @@ def cli(database, user, host, port, prompt_passwd, never_prompt):
                     output.extend(format_output(rows, headers, status))
                 click.echo_via_pager('\n'.join(output))
             except KeyboardInterrupt:
-                pgexecute.reconnect()
+                # Restart connection to the database
+                pgexecute.connect()
                 _logger.debug("cancelled query, sql: %r", document.text)
                 click.secho("cancelled query", err=True, fg='red')
             except Exception as e:
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index e4b31d51..27af3085 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -13,6 +13,7 @@ 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
+# See http://initd.org/psycopg/articles/2014/07/20/cancelling-postgresql-statements-python/
 psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
 
 def _parse_dsn(dsn, default_user, default_password, default_host,
@@ -74,9 +75,11 @@ class PGExecute(object):
                 _parse_dsn(database, default_user=user,
                         default_password=password, default_host=host,
                         default_port=port)
-        self.reconnect()
+        self.connect()
 
-    def reconnect(self):
+    def connect(self):
+        if hasattr(self, 'conn'):
+            self.conn.close()
         self.conn = psycopg2.connect(database=self.dbname, user=self.user,
                 password=self.password, host=self.host, port=self.port)
         self.conn.autocommit = True