Handle the case where \e is issued multiple times.

3 files changed, 28 insertions, 20 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index 59504caf..35664cdd 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -197,17 +197,29 @@ class PGCli(object):
                 if quit_command(document.text):
                     raise Exit
 
-                # Editor command is a special case. We don't have to execute it
-                # and print out the result. We have to extend our output and
-                # continue typing in SQL / commands.
-                if iospecial.editor_command(document.text):
-                    sql, _, _, message = list(pgexecute.run(document.text))[0]
-                    if sql or not message:
-                        # We either have some SQL, or no SQL, but no error
-                        # message either.
+                try:
+                    # Editor command is any query that is prefixed or suffixed
+                    # by a '\e'. The reason for a while loop is because a user
+                    # might edit a query multiple times.
+                    # For eg:
+                    # "select * from \e"<enter> to edit it in vim, then come
+                    # back to the prompt with the edited query "select * from
+                    # blah where q = 'abc'\e" to edit it again.
+                    while iospecial.editor_command(document.text):
+                        filename = iospecial.get_filename(document.text)
+                        sql, message = iospecial.open_external_editor(filename, sql=document.text)
+                        if message:
+                            # Something went wrong. Raise an exception and bail.
+                            raise RuntimeError(message)
                         document = cli.read_input(
-                            initial_document=Document(sql, cursor_position=len(sql)),
-                            on_exit=AbortAction.RAISE_EXCEPTION)
+                                initial_document=Document(sql, cursor_position=len(sql)),
+                                on_exit=AbortAction.RAISE_EXCEPTION)
+                        continue
+                except RuntimeError as e:
+                    logger.error("sql: %r, error: %r", document.text, e)
+                    logger.error("traceback: %r", traceback.format_exc())
+                    click.secho(str(e), err=True, fg='red')
+                    continue
 
                 # Keep track of whether or not the query is mutating. In case
                 # of a multi-statement query, the overall query is considered
diff --git a/pgcli/packages/iospecial.py b/pgcli/packages/iospecial.py
index e0b38c93..42984213 100644
--- a/pgcli/packages/iospecial.py
+++ b/pgcli/packages/iospecial.py
@@ -14,8 +14,12 @@ def editor_command(command):
     """
     # It is possible to have `\e filename` or `SELECT * FROM \e`. So we check
     # for both conditions.
-    return command.strip().endswith('\e') or command.strip().startswith('\e')
+    return command.strip().endswith('\\e') or command.strip().startswith('\\e')
 
+def get_filename(sql):
+    if sql.strip().startswith('\\e'):
+        command, _, filename = sql.partition(' ')
+        return filename.strip() or None
 
 def open_external_editor(filename=None, sql=''):
     """
@@ -49,4 +53,4 @@ def open_external_editor(filename=None, sql=''):
         # Empty string is ok.
         query = sql
 
-    yield (query, None, None, message)
+    return (query, message)
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index eee9e9c8..c8bbe0a9 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -190,14 +190,6 @@ class PGExecute(object):
                 _logger.debug('Successfully switched to DB: %r', dbname)
                 yield (None, None, None, 'You are now connected to database "%s" as '
                         'user "%s"' % (self.dbname, self.user))
-            elif iospecial.editor_command(sql):
-                filename = None
-                # If the query starts with \e then it can be followed by a filename.
-                if sql.startswith('\e'):
-                    tokens = sql.split(' ', 2)
-                    filename = tokens[1] if len(tokens) > 1 and tokens[1] else None
-                for result in iospecial.open_external_editor(filename=filename, sql=sql):
-                    yield result
             else:
                 try:   # Special command
                     _logger.debug('Trying a pgspecial command. sql: %r', sql)