Handle a multi-line query on Enter key-press (fixes #1031)

Use a (conditional) Enter key-binding to force-handle a multi-line
buffer, rather than doing so by (conditionally) disabling the multiline
mode of prompt_toolkit.

This has the benefit of being more efficient (the multiline Condition
filter is called very often, which (due to the repeated query parsing)
causes editing to become slow with a large buffer that ends in a
semicolon), clearer in intent (we want to force-handle the query, rather
than (temporarily) disable multiline mode which indirectly forces the
buffer to be handled) and avoids a bug in multi-line history search
(issue #1031)

4 files changed, 43 insertions, 32 deletions

diff --git a/changelog.rst b/changelog.rst
index 1ec7b729..e33fb00b 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -14,6 +14,7 @@ Bug fixes:
 * Empty query caused error message (#1019) (Thanks: `Sebastian Janko`_)
 * History navigation bindings in multiline queries (#1004) (Thanks: `Pedro Ferrari`_)
 * Can't connect to pgbouncer database (#1093). (Thanks: `Irina Truong`_)
+* Fix broken multi-line history search (#1031). (Thanks: `Owen Stephens`_)
 
 Internal:
 ---------
diff --git a/pgcli/key_bindings.py b/pgcli/key_bindings.py
index c839256e..a8f392f6 100644
--- a/pgcli/key_bindings.py
+++ b/pgcli/key_bindings.py
@@ -9,6 +9,8 @@ from prompt_toolkit.filters import (
     has_selection,
 )
 
+from .pgbuffer import multi_line_buffer_should_be_handled
+
 _logger = logging.getLogger(__name__)
 
 
@@ -94,6 +96,15 @@ def pgcli_bindings(pgcli):
         event.current_buffer.complete_state = None
         event.app.current_buffer.complete_state = None
 
+    # When using multi_line input mode the buffer is not handled on Enter (a new line is
+    # inserted instead), so we force the handling if one of several conditions are True
+    @kb.add(
+        "enter",
+        filter=~completion_is_selected & multi_line_buffer_should_be_handled(pgcli),
+    )
+    def _(event):
+        event.current_buffer.validate_and_handle()
+
     @kb.add("escape", "enter")
     def _(event):
         """Introduces a line break regardless of multi-line mode or not."""
diff --git a/pgcli/main.py b/pgcli/main.py
index c756849b..7d5b35bd 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -35,7 +35,7 @@ from prompt_toolkit.completion import DynamicCompleter, ThreadedCompleter
 from prompt_toolkit.enums import DEFAULT_BUFFER, EditingMode
 from prompt_toolkit.shortcuts import PromptSession, CompleteStyle
 from prompt_toolkit.document import Document
-from prompt_toolkit.filters import HasFocus, IsDone
+from prompt_toolkit.filters import HasFocus, IsDone, Condition
 from prompt_toolkit.lexers import PygmentsLexer
 from prompt_toolkit.layout.processors import (
     ConditionalProcessor,
@@ -53,7 +53,6 @@ from .pgcompleter import PGCompleter
 from .pgtoolbar import create_toolbar_tokens_func
 from .pgstyle import style_factory, style_factory_output
 from .pgexecute import PGExecute
-from .pgbuffer import pg_is_multiline
 from .completion_refresher import CompletionRefresher
 from .config import (
     get_casing_file,
@@ -810,7 +809,7 @@ class PGCli(object):
                 ],
                 auto_suggest=AutoSuggestFromHistory(),
                 tempfile_suffix=".sql",
-                multiline=pg_is_multiline(self),
+                multiline=Condition(lambda: self.multi_line),
                 history=history,
                 completer=ThreadedCompleter(DynamicCompleter(lambda: self.completer)),
                 complete_while_typing=True,
diff --git a/pgcli/pgbuffer.py b/pgcli/pgbuffer.py
index 7e89542d..995ee208 100644
--- a/pgcli/pgbuffer.py
+++ b/pgcli/pgbuffer.py
@@ -6,21 +6,6 @@ from prompt_toolkit.application import get_app
 from .packages.parseutils.utils import is_open_quote
 
 
-def pg_is_multiline(pgcli):
-    @Condition
-    def cond():
-        doc = get_app().layout.get_buffer_by_name(DEFAULT_BUFFER).document
-
-        if not pgcli.multi_line:
-            return False
-        if pgcli.multiline_mode == "safe":
-            return True
-        else:
-            return not _multiline_exception(doc.text)
-
-    return cond
-
-
 def _is_complete(sql):
     # A complete command is an sql statement that ends with a semicolon, unless
     # there's an open quote surrounding it, as is common when writing a
@@ -28,17 +13,32 @@ def _is_complete(sql):
     return sql.endswith(";") and not is_open_quote(sql)
 
 
-def _multiline_exception(text):
-    text = text.strip()
-    return (
-        text.startswith("\\")
-        or text.endswith(r"\e")  # Special Command
-        or text.endswith(r"\G")  # Special Command
-        or _is_complete(text)  # Ended with \e which should launch the editor
-        or (text == "exit")  # A complete SQL command
-        or (text == "quit")  # Exit doesn't need semi-colon
-        or (text == ":q")  # Quit doesn't need semi-colon
-        or (  # To all the vim fans out there
-            text == ""
-        )  # Just a plain enter without any text
-    )
+"""
+Returns True if the input mode is multi-line using psql mode, and the main
+buffer's contents indicate that it should be handled, False otherwise. This
+method is required since by default prompt_toolkit would not handle a buffer on
+Enter keypress when in multi-line mode.
+"""
+
+
+def multi_line_buffer_should_be_handled(pgcli):
+    @Condition
+    def cond():
+        if not pgcli.multi_line or pgcli.multiline_mode == "safe":
+            return False
+
+        doc = get_app().layout.get_buffer_by_name(DEFAULT_BUFFER).document
+        text = doc.text.strip()
+
+        return (
+            text.startswith("\\")  # Special Command
+            or text.endswith(r"\e")  # Special Command
+            or text.endswith(r"\G")  # Ended with \e which should launch the editor
+            or _is_complete(text)  # A complete SQL command
+            or (text == "exit")  # Exit doesn't need semi-colon
+            or (text == "quit")  # Quit doesn't need semi-colon
+            or (text == ":q")  # To all the vim fans out there
+            or (text == "")  # Just a plain enter without any text
+        )
+
+    return cond