Merge pull request #877 from catherinedevlin/keep_callables

Keep callables

1 files changed, 14 insertions, 13 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index 73c1a247..ca4d4efa 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -454,19 +454,20 @@ class PGCli(object):
         # It's internal api of prompt_toolkit that may change. This was added to fix #668.
         # We may find a better way to do it in the future.
         saved_callables = cli.application.pre_run_callables
-        while special.editor_command(document.text):
-            filename = special.get_filename(document.text)
-            query = (special.get_editor_query(document.text) or
-                     self.get_last_query())
-            sql, message = special.open_external_editor(filename, sql=query)
-            if message:
-                # Something went wrong. Raise an exception and bail.
-                raise RuntimeError(message)
-            cli.current_buffer.document = Document(sql, cursor_position=len(sql))
-            cli.application.pre_run_callables = []
-            document = cli.run()
-            continue
-        cli.application.pre_run_callables = saved_callables
+        try:
+            while special.editor_command(document.text):
+                filename = special.get_filename(document.text)
+                query = (special.get_editor_query(document.text) or
+                        self.get_last_query())
+                sql, message = special.open_external_editor(filename, sql=query)
+                if message:
+                    # Something went wrong. Raise an exception and bail.
+                    raise RuntimeError(message)
+                cli.current_buffer.document = Document(sql, cursor_position=len(sql))
+                cli.application.pre_run_callables = []
+                document = cli.run()
+        finally:
+            cli.application.pre_run_callables = saved_callables
         return document
 
     def execute_command(self, text, query):