Merge pull request #1017 from dbcli/j-bennet/bugfix-should-exit-even-with-connection-closed

Fix for https://github.com/dbcli/pgcli/issues/1014.

3 files changed, 39 insertions, 5 deletions

diff --git a/changelog.rst b/changelog.rst
index 069d782e..7db254c5 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -13,6 +13,7 @@ Bug fixes:
 * All pexpect submodules have been moved into the pexpect package as of version 3.0. Use pexpect.TIMEOUT (Thanks: `Marcin Cieślak`_)
 * Resizing pgcli terminal kills the connection to postgres in python 2.7 (Thanks: `Amjith Ramanujam`_)
 * Fix crash retrieving server version with ``--single-connection``. (Thanks: `Irina Truong`_)
+* Cannot quit application without reconnecting to database (#1014). (Thanks: `Irina Truong`_)
 
 Internal:
 ---------
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 4deaf279..8c3bfa82 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -354,7 +354,13 @@ class PGExecute(object):
                 if pgspecial:
                     # First try to run each query as special
                     _logger.debug('Trying a pgspecial command. sql: %r', sql)
-                    cur = self.conn.cursor()
+                    try:
+                        cur = self.conn.cursor()
+                    except psycopg2.InterfaceError:
+                        # edge case when connection is already closed, but we
+                        # don't need cursor for special_cmd.arg_type == NO_QUERY.
+                        # See https://github.com/dbcli/pgcli/issues/1014.
+                        cur = None
                     try:
                         for result in pgspecial.execute(cur, sql):
                             # e.g. execute_from_file already appends these
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index 39788950..f49cb6bc 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -1,13 +1,16 @@
 # coding=UTF-8
 from __future__ import print_function
 
-import pytest
-import psycopg2
-from mock import patch
-from pgcli.packages.parseutils.meta import FunctionMetadata
 from textwrap import dedent
+
+import psycopg2
+import pytest
+from mock import patch, MagicMock
+from pgspecial.main import PGSpecial, NO_QUERY
 from utils import run, dbtest, requires_json, requires_jsonb
+
 from pgcli.main import PGCli
+from pgcli.packages.parseutils.meta import FunctionMetadata
 
 
 def function_meta_data(
@@ -425,3 +428,27 @@ def test_short_host(executor):
         assert executor.short_host == 'localhost'
     with patch.object(executor, 'host', 'localhost1.example.org,localhost2.example.org'):
         assert executor.short_host == 'localhost1'
+
+
+class BrokenConnection(object):
+    """Mock a connection that failed."""
+
+    def cursor(self):
+        raise psycopg2.InterfaceError("I'm broken!")
+
+
+@dbtest
+def test_exit_without_active_connection(executor):
+    quit_handler = MagicMock()
+    pgspecial = PGSpecial()
+    pgspecial.register(quit_handler, '\\q', '\\q', 'Quit pgcli.',
+                       arg_type=NO_QUERY, case_sensitive=True, aliases=(':q',))
+
+    with patch.object(executor, "conn", BrokenConnection()):
+        # we should be able to quit the app, even without active connection
+        run(executor, "\\q", pgspecial=pgspecial)
+        quit_handler.assert_called_once()
+
+        # an exception should be raised when running a query without active connection
+        with pytest.raises(psycopg2.InterfaceError):
+            run(executor, "select 1", pgspecial=pgspecial)