Merge branch 'master' of github.com:dbcli/mycli into tsroten/config_revamptsroten/config_revamp

* 'master' of github.com:dbcli/mycli: (23 commits)
  Fix pytest in Fedora mock
  Add Python 3.7 trove classifier.
  Releasing version 1.18.2
  Prep for v1.18.2 release.
  Update changelog.md
  typo
  Move behave logging changelog entry.
  Fix warnings for running tests on Python 3.7
  Releasing version 1.18.1
  Move update to TBD.
  Add twine update to changelog.
  Update to the newest version of Twine.
  Prep for v1.18.1 release.
  reformat code
  Update changelog.md
  update changelog.md
  feature: add keywords
  Move prompt toolkit to TBD release
  Clean up and add behave logging
  merge pgcli jonathanslenders/prompt-toolkit-2.0 3c4b1ce..6fc9c56
  ...

24 files changed, 414 insertions, 264 deletions

diff --git a/changelog.md b/changelog.md
index 273aceb..19ba2b0 100644
--- a/changelog.md
+++ b/changelog.md
@@ -1,3 +1,40 @@
+TBD
+===
+
+Internal:
+---------
+
+* Add Python 3.7 trove classifier (Thanks: [Thomas Roten]).
+* Fix pytest in Fedora mock (Thanks: [Dick Marinus]).
+
+1.18.2
+======
+
+Bug Fixes:
+----------
+
+* Fixes database reconnecting feature (Thanks: [Yang Zou]).
+
+Internal:
+---------
+
+* Update Twine version to 1.12.1 (Thanks: [Thomas Roten]).
+* Fix warnings for running tests on Python 3.7 (Thanks: [Dick Marinus]).
+* Clean up and add behave logging (Thanks: [Dick Marinus]).
+
+1.18.1
+======
+
+Features:
+---------
+
+* Add Keywords: TINYINT, SMALLINT, MEDIUMINT, INT, BIGINT (Thanks: [QiaoHou Peng]).
+
+Internal:
+---------
+
+* Update prompt toolkit (Thanks: [Jonathan Slenders], [Irina Truong], [Dick Marinus]).
+
 1.18.0
 ======
 
diff --git a/mycli/AUTHORS b/mycli/AUTHORS
index 7bd0127..034c9d3 100644
--- a/mycli/AUTHORS
+++ b/mycli/AUTHORS
@@ -61,6 +61,7 @@ Contributors:
   * Zhongyang Guan
   * Huachao Mao
   * QiaoHou Peng
+  * Yang Zou
 
 Creator:
 --------
diff --git a/mycli/__init__.py b/mycli/__init__.py
index 9eb549a..5723b02 100644
--- a/mycli/__init__.py
+++ b/mycli/__init__.py
@@ -1 +1 @@
-__version__ = '1.18.0'
+__version__ = '1.18.2'
diff --git a/mycli/clibuffer.py b/mycli/clibuffer.py
index 41a63df..f6cc737 100644
--- a/mycli/clibuffer.py
+++ b/mycli/clibuffer.py
@@ -1,17 +1,21 @@
-from prompt_toolkit.buffer import Buffer
+from __future__ import unicode_literals
+
+from prompt_toolkit.enums import DEFAULT_BUFFER
 from prompt_toolkit.filters import Condition
+from prompt_toolkit.application import get_app
+from .packages.parseutils import is_open_quote
 
-class CLIBuffer(Buffer):
-    def __init__(self, always_multiline, *args, **kwargs):
-        self.always_multiline = always_multiline
 
-        @Condition
-        def is_multiline():
-            doc = self.document
-            return self.always_multiline and not _multiline_exception(doc.text)
+def cli_is_multiline(mycli):
+    @Condition
+    def cond():
+        doc = get_app().layout.get_buffer_by_name(DEFAULT_BUFFER).document
 
-        super(self.__class__, self).__init__(*args, is_multiline=is_multiline,
-                                             tempfile_suffix='.sql', **kwargs)
+        if not mycli.multi_line:
+            return False
+        else:
+            return not _multiline_exception(doc.text)
+    return cond
 
 def _multiline_exception(text):
     orig = text
diff --git a/mycli/clistyle.py b/mycli/clistyle.py
index 89caf71..bde464c 100644
--- a/mycli/clistyle.py
+++ b/mycli/clistyle.py
@@ -1,28 +1,118 @@
-import pygments.style
+from __future__ import unicode_literals
+
+import logging
+
 import pygments.styles
-from pygments.token import string_to_tokentype
+from pygments.token import string_to_tokentype, Token
+from pygments.style import Style as PygmentsStyle
 from pygments.util import ClassNotFound
+from prompt_toolkit.styles.pygments import style_from_pygments_cls
+from prompt_toolkit.styles import merge_styles, Style
 
+logger = logging.getLogger(__name__)
 
-def style_factory(name, cli_style):
-    """Create a Pygments Style class based on the user's preferences.
+# map Pygments tokens (ptk 1.0) to class names (ptk 2.0).
+TOKEN_TO_PROMPT_STYLE = {
+    Token.Menu.Completions.Completion.Current: 'completion-menu.completion.current',
+    Token.Menu.Completions.Completion: 'completion-menu.completion',
+    Token.Menu.Completions.Meta.Current: 'completion-menu.meta.completion.current',
+    Token.Menu.Completions.Meta: 'completion-menu.meta.completion',
+    Token.Menu.Completions.MultiColumnMeta: 'completion-menu.multi-column-meta',
+    Token.Menu.Completions.ProgressButton: 'scrollbar.arrow',  # best guess
+    Token.Menu.Completions.ProgressBar: 'scrollbar',  # best guess
+    Token.SelectedText: 'selected',
+    Token.SearchMatch: 'search',
+    Token.SearchMatch.Current: 'search.current',
+    Token.Toolbar: 'bottom-toolbar',
+    Token.Toolbar.Off: 'bottom-toolbar.off',
+    Token.Toolbar.On: 'bottom-toolbar.on',
+    Token.Toolbar.Search: 'search-toolbar',
+    Token.Toolbar.Search.Text: 'search-toolbar.text',
+    Token.Toolbar.System: 'system-toolbar',
+    Token.Toolbar.Arg: 'arg-toolbar',
+    Token.Toolbar.Arg.Text: 'arg-toolbar.text',
+    Token.Toolbar.Transaction.Valid: 'bottom-toolbar.transaction.valid',
+    Token.Toolbar.Transaction.Failed: 'bottom-toolbar.transaction.failed',
+    Token.Output.Header: 'output.header',
+    Token.Output.OddRow: 'output.odd-row',
+    Token.Output.EvenRow: 'output.even-row',
+}
+
+# reverse dict for cli_helpers, because they still expect Pygments tokens.
+PROMPT_STYLE_TO_TOKEN = {
+    v: k for k, v in TOKEN_TO_PROMPT_STYLE.items()
+}
 
-    :param str name: The name of a built-in Pygments style.
-    :param dict cli_style: The user's token-type style preferences.
+
+def parse_pygments_style(token_name, style_object, style_dict):
+    """Parse token type and style string.
+
+    :param token_name: str name of Pygments token. Example: "Token.String"
+    :param style_object: pygments.style.Style instance to use as base
+    :param style_dict: dict of token names and their styles, customized to this cli
 
     """
+    token_type = string_to_tokentype(token_name)
+    try:
+        other_token_type = string_to_tokentype(style_dict[token_name])
+        return token_type, style_object.styles[other_token_type]
+    except AttributeError as err:
+        return token_type, style_dict[token_name]
+
+
+def style_factory(name, cli_style):
     try:
         style = pygments.styles.get_style_by_name(name)
     except ClassNotFound:
         style = pygments.styles.get_style_by_name('native')
 
-    style_tokens = {}
-    style_tokens.update(style.styles)
-    custom_styles = {string_to_tokentype(x): y for x, y in cli_style.items()}
-    style_tokens.update(custom_styles)
+    prompt_styles = []
+    # prompt-toolkit used pygments tokens for styling before, switched to style
+    # names in 2.0. Convert old token types to new style names, for backwards compatibility.
+    for token in cli_style:
+        if token.startswith('Token.'):
+            # treat as pygments token (1.0)
+            token_type, style_value = parse_pygments_style(
+                token, style, cli_style)
+            if token_type in TOKEN_TO_PROMPT_STYLE:
+                prompt_style = TOKEN_TO_PROMPT_STYLE[token_type]
+                prompt_styles.append((prompt_style, style_value))
+            else:
+                # we don't want to support tokens anymore
+                logger.error('Unhandled style / class name: %s', token)
+        else:
+            # treat as prompt style name (2.0). See default style names here:
+            # https://github.com/jonathanslenders/python-prompt-toolkit/blob/master/prompt_toolkit/styles/defaults.py
+            prompt_styles.append((token, cli_style[token]))
+
+    override_style = Style([('bottom-toolbar', 'noreverse')])
+    return merge_styles([
+        style_from_pygments_cls(style),
+        override_style,
+        Style(prompt_styles)
+    ])
+
+
+def style_factory_output(name, cli_style):
+    try:
+        style = pygments.styles.get_style_by_name(name).styles
+    except ClassNotFound:
+        style = pygments.styles.get_style_by_name('native').styles
+
+    for token in cli_style:
+        if token.startswith('Token.'):
+            token_type, style_value = parse_pygments_style(
+                token, style, cli_style)
+            style.update({token_type: style_value})
+        elif token in PROMPT_STYLE_TO_TOKEN:
+            token_type = PROMPT_STYLE_TO_TOKEN[token]
+            style.update({token_type: cli_style[token]})
+        else:
+            # TODO: cli helpers will have to switch to ptk.Style
+            logger.error('Unhandled style / class name: %s', token)
 
-    class MycliStyle(pygments.style.Style):
-        default_styles = ''
-        styles = style_tokens
+    class OutputStyle(PygmentsStyle):
+        default_style = ""
+        styles = style
 
-    return MycliStyle
+    return OutputStyle
diff --git a/mycli/clitoolbar.py b/mycli/clitoolbar.py
index bb638d5..89e6afa 100644
--- a/mycli/clitoolbar.py
+++ b/mycli/clitoolbar.py
@@ -1,48 +1,48 @@
-from pygments.token import Token
-from prompt_toolkit.enums import DEFAULT_BUFFER, EditingMode
-from prompt_toolkit.key_binding.vi_state import InputMode
+from __future__ import unicode_literals
 
+from prompt_toolkit.key_binding.vi_state import InputMode
+from prompt_toolkit.application import get_app
+from prompt_toolkit.enums import EditingMode
 
-def create_toolbar_tokens_func(get_is_refreshing, show_fish_help):
-    """
-    Return a function that generates the toolbar tokens.
-    """
-    token = Token.Toolbar
 
-    def get_toolbar_tokens(cli):
+def create_toolbar_tokens_func(mycli, show_fish_help):
+    """Return a function that generates the toolbar tokens."""
+    def get_toolbar_tokens():
         result = []
-        result.append((token, ' '))
+        result.append(('class:bottom-toolbar', ' '))
 
-        if cli.buffers[DEFAULT_BUFFER].always_multiline:
-            result.append((token.On, '[F3] Multiline: ON  '))
-        else:
-            result.append((token.Off, '[F3] Multiline: OFF  '))
+        if mycli.multi_line:
+            result.append(
+                ('class:bottom-toolbar', ' (Semi-colon [;] will end the line) '))
 
-        if cli.buffers[DEFAULT_BUFFER].always_multiline:
-            result.append((token,
-                ' (Semi-colon [;] will end the line)'))
-
-        if cli.editing_mode == EditingMode.VI:
+        if mycli.multi_line:
+            result.append(('class:bottom-toolbar.on', '[F3] Multiline: ON  '))
+        else:
+            result.append(('class:bottom-toolbar.off',
+                           '[F3] Multiline: OFF  '))
+        if mycli.prompt_app.editing_mode == EditingMode.VI:
             result.append((
-                token.On,
-                'Vi-mode ({})'.format(_get_vi_mode(cli))
+                'class:botton-toolbar.on',
+                'Vi-mode ({})'.format(_get_vi_mode())
             ))
 
         if show_fish_help():
-            result.append((token, '  Right-arrow to complete suggestion'))
+            result.append(
+                ('class:bottom-toolbar', '  Right-arrow to complete suggestion'))
 
-        if get_is_refreshing():
-            result.append((token, '     Refreshing completions...'))
+        if mycli.completion_refresher.is_refreshing():
+            result.append(
+                ('class:bottom-toolbar', '     Refreshing completions...'))
 
         return result
     return get_toolbar_tokens
 
 
-def _get_vi_mode(cli):
+def _get_vi_mode():
     """Get the current vi mode for display."""
     return {
         InputMode.INSERT: 'I',
         InputMode.NAVIGATION: 'N',
         InputMode.REPLACE: 'R',
-        InputMode.INSERT_MULTIPLE: 'M'
-    }[cli.vi_state.input_mode]
+        InputMode.INSERT_MULTIPLE: 'M',
+    }[get_app().vi_state.input_mode]
diff --git a/mycli/config b/mycli/config
index 3c91dc0..e04c3d3 100644
--- a/mycli/config
+++ b/mycli/config
@@ -85,41 +85,33 @@ pager = option('on', 'off', 'auto', default='auto')
 
 pager_command = string(default='')
 
-# disabled pager on startup
-enable_pager = True
-
 # Custom colors for the completion menu, toolbar, etc.
 [colors]
-# Completion menus.
-Token.Menu.Completions.Completion.Current = string(default='bg:#00aaaa #000000')
-Token.Menu.Completions.Completion = string(default='bg:#008888 #ffffff')
-Token.Menu.Completions.MultiColumnMeta = string(default='bg:#aaffff #000000')
-Token.Menu.Completions.ProgressButton = string(default='bg:#003333')
-Token.Menu.Completions.ProgressBar = string(default='bg:#00aaaa')
-
-# Query results
-Token.Output.Header = string(default='bold')
-Token.Output.OddRow = string(default='')
-Token.Output.EvenRow = string(default='')
-
-# Selected text.
-Token.SelectedText = string(default='#ffffff bg:#6666aa')
-
-# Search matches. (reverse-i-search)
-Token.SearchMatch = string(default='#ffffff bg:#4444aa')
-Token.SearchMatch.Current = string(default='#ffffff bg:#44aa44')
-
-# The bottom toolbar.
-Token.Toolbar = string(default='bg:#222222 #aaaaaa')
-Token.Toolbar.Off = string(default='bg:#222222 #888888')
-Token.Toolbar.On = string(default='bg:#222222 #ffffff')
-
-# Search/arg/system toolbars.
-Token.Toolbar.Search = string(default='noinherit bold')
-Token.Toolbar.Search.Text = string(default='nobold')
-Token.Toolbar.System = string(default='noinherit bold')
-Token.Toolbar.Arg = string(default='noinherit bold')
-Token.Toolbar.Arg.Text = string(default='nobold')
+completion-menu.completion.current = string(default=''bg:#ffffff #000000')
+completion-menu.completion = string(default='bg:#008888 #ffffff')
+completion-menu.meta.completion.current = string(default='bg:#44aaaa #000000')
+completion-menu.meta.completion = string(default='bg:#448888 #ffffff')
+completion-menu.multi-column-meta = string(default='bg:#aaffff #000000')
+scrollbar.arrow = string(default='bg:#003333')
+scrollbar = string(default='bg:#00aaaa')
+selected = string(default='#ffffff bg:#6666aa')
+search = string(default='#ffffff bg:#4444aa')
+search.current = string(default='#ffffff bg:#44aa44')
+bottom-toolbar = string(default='bg:#222222 #aaaaaa')
+bottom-toolbar.off = string(default='bg:#222222 #888888')
+bottom-toolbar.on = string(default='bg:#222222 #ffffff')
+search-toolbar = string(default='noinherit bold')
+search-toolbar.text = string(default='nobold')
+system-toolbar = string(default='noinherit bold')
+arg-toolbar = string(default='noinherit bold')
+arg-toolbar.text = string(default='nobold')
+bottom-toolbar.transaction.valid = string(default='bg:#222222 #00ff5f bold')
+bottom-toolbar.transaction.failed = string(default='bg:#222222 #ff005f bold')
+
+# style classes for colored table output
+output.header = string(default='#00ff5f bold')
+output.odd-row = string(default='')
+output.even-row = string(default='')
 
 # Favorite queries.
 [favorite_queries]
diff --git a/mycli/filters.py b/mycli/filters.py
deleted file mode 100644
index 6a8075f..0000000
--- a/mycli/filters.py
+++ /dev/null
@@ -1,12 +0,0 @@
-from prompt_toolkit.filters import Filter
-
-class HasSelectedCompletion(Filter):
-    """Enable when the current buffer has a selected completion."""
-
-    def __call__(self, cli):
-        complete_state = cli.current_buffer.complete_state
-        return (complete_state is not None and
-                 complete_state.current_completion is not None)
-
-    def __repr__(self):
-        return "HasSelectedCompletion()"
diff --git a/mycli/key_bindings.py b/mycli/key_bindings.py
index 33bb4f2..abc559f 100644
--- a/mycli/key_bindings.py
+++ b/mycli/key_bindings.py
@@ -1,65 +1,50 @@
+from __future__ import unicode_literals
 import logging
 from prompt_toolkit.enums import EditingMode
-from prompt_toolkit.keys import Keys
-from prompt_toolkit.key_binding.manager import KeyBindingManager
-from .filters import HasSelectedCompletion
+from prompt_toolkit.filters import completion_is_selected
+from prompt_toolkit.key_binding import KeyBindings
 
 _logger = logging.getLogger(__name__)
 
 
-def mycli_bindings():
-    """
-    Custom key bindings for mycli.
-    """
-    key_binding_manager = KeyBindingManager(
-        enable_open_in_editor=True,
-        enable_system_bindings=True,
-        enable_auto_suggest_bindings=True,
-        enable_search=True,
-        enable_abort_and_exit_bindings=True)
+def mycli_bindings(mycli):
+    """Custom key bindings for mycli."""
+    kb = KeyBindings()
 
-    @key_binding_manager.registry.add_binding(Keys.F2)
+    @kb.add('f2')
     def _(event):
-        """
-        Enable/Disable SmartCompletion Mode.
-        """
+        """Enable/Disable SmartCompletion Mode."""
         _logger.debug('Detected F2 key.')
-        buf = event.cli.current_buffer
-        buf.completer.smart_completion = not buf.completer.smart_completion
+        mycli.completer.smart_completion = not mycli.completer.smart_completion
 
-    @key_binding_manager.registry.add_binding(Keys.F3)
+    @kb.add('f3')
     def _(event):
-        """
-        Enable/Disable Multiline Mode.
-        """
+        """Enable/Disable Multiline Mode."""
         _logger.debug('Detected F3 key.')
-        buf = event.cli.current_buffer
-        buf.always_multiline = not buf.always_multiline
+        mycli.multi_line = not mycli.multi_line
 
-    @key_binding_manager.registry.add_binding(Keys.F4)
+    @kb.add('f4')
     def _(event):
-        """
-        Toggle between Vi and Emacs mode.
-        """
+        """Toggle between Vi and Emacs mode."""
         _logger.debug('Detected F4 key.')
-        if event.cli.editing_mode == EditingMode.VI:
-            event.cli.editing_mode = EditingMode.EMACS
+        if mycli.key_bindings == "vi":
+            event.app.editing_mode = EditingMode.EMACS
+            mycli.key_bindings = "emacs"
         else:
-            event.cli.editing_mode = EditingMode.VI
+            event.app.editing_mode = EditingMode.VI
+            mycli.key_bindings = "vi"
 
-    @key_binding_manager.registry.add_binding(Keys.Tab)
+    @kb.add('tab')
     def _(event):
-        """
-        Force autocompletion at cursor.
-        """
+        """Force autocompletion at cursor."""
         _logger.debug('Detected <Tab> key.')
-        b = event.cli.current_buffer
+        b = event.app.current_buffer
         if b.complete_state:
             b.complete_next()
         else:
-            event.cli.start_completion(select_first=True)
+            b.start_completion(select_first=True)
 
-    @key_binding_manager.registry.add_binding(Keys.ControlSpace)
+    @kb.add('c-space')
     def _(event):
         """
         Initialize autocompletion at cursor.
@@ -71,21 +56,25 @@ def mycli_bindings():
         """
         _logger.debug('Detected <C-Space> key.')
 
-        b = event.cli.current_buffer
+        b = event.app.current_buffer
         if b.complete_state:
             b.complete_next()
         else:
-            event.cli.start_completion(select_first=False)
+            b.start_completion(select_first=False)
 
-    @key_binding_manager.registry.add_binding(Keys.ControlJ, filter=HasSelectedCompletion())
+    @kb.add('enter', filter=completion_is_selected)
     def _(event):
+        """Makes the enter key work as the tab key only when showing the menu.
+
+        In other words, don't execute query when enter is pressed in
+        the completion dropdown menu, instead close the dropdown menu
+        (accept current selection).
+
         """
-        Makes the enter key work as the tab key only when showing the menu.
-        """
-        _logger.debug('Detected <C-J> key.')
+        _logger.debug('Detected enter key.')
 
         event.current_buffer.complete_state = None
-        b = event.cli.current_buffer
+        b = event.app.current_buffer
         b.complete_state = None
 
-    return key_binding_manager
+    return kb
diff --git a/mycli/main.py b/mycli/main.py
index 4e87517..28c3e86 100755
--- a/mycli/main.py
+++ b/mycli/main.py
@@ -15,18 +15,17 @@ from cli_helpers.tabular_output import TabularOutputFormatter
 from cli_helpers.tabular_output import preprocessors
 import click
 import sqlparse
-from prompt_toolkit import CommandLineInterface, Application, AbortAction
-from prompt_toolkit.interface import AcceptAction
+from prompt_toolkit.completion import DynamicCompleter
 from prompt_toolkit.enums import DEFAULT_BUFFER, EditingMode
-from prompt_toolkit.shortcuts import create_prompt_layout, create_eventloop
-from prompt_toolkit.styles.from_pygments import style_from_pygments
+from prompt_toolkit.shortcuts import PromptSession, CompleteStyle
+from prompt_toolkit.styles.pygments import style_from_pygments_cls
 from prompt_toolkit.document import Document
-from prompt_toolkit.filters import Always, HasFocus, IsDone
+from prompt_toolkit.filters import HasFocus, IsDone
 from prompt_toolkit.layout.processors import (HighlightMatchingBracketProcessor,
                                               ConditionalProcessor)
+from prompt_toolkit.lexers import PygmentsLexer
 from prompt_toolkit.history import FileHistory
 from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
-from pygments.token import Token
 
 from .packages.special.main import NO_QUERY
 from .packages.prompt_utils import confirm, confirm_destructive_query, prompt
@@ -34,9 +33,9 @@ from .packages.tabular_output import sql_format
 import mycli.packages.special as special
 from .sqlcompleter import SQLCompleter
 from .clitoolbar import create_toolbar_tokens_func
-from .clistyle import style_factory
+from .clistyle import style_factory, style_factory_output
 from .sqlexecute import FIELD_TYPES, SQLExecute
-from .clibuffer import CLIBuffer
+from .clibuffer import cli_is_multiline
 from .completion_refresher import CompletionRefresher
 from .config import MyCliConfig
 from .key_bindings import mycli_bindings
@@ -53,7 +52,6 @@ click.disable_unicode_literals_warning = True
 try:
     from urlparse import urlparse
     from urlparse import unquote
-    FileNotFoundError = OSError
 except ImportError:
     from urllib.parse import urlparse
     from urllib.parse import unquote
@@ -102,8 +100,13 @@ class MyCli(object):
             format_name=c['main']['table_format'])
         sql_format.register_new_formatter(self.formatter)
         self.formatter.mycli = self
+        self.syntax_style = c['main']['syntax_style']
         self.less_chatty = c['main']['less_chatty']
-        self.output_style = style_factory(c['main']['syntax_style'], c['colors'])
+        self.cli_style = c['colors']
+        self.output_style = style_factory_output(
+            self.syntax_style,
+            self.cli_style
+        )
         self.wider_completion_menu = c['main']['wider_completion_menu']
         c_dest_warning = c['main']['destructive_warning']
         self.destructive_warning = c_dest_warning if warn is None else warn
@@ -144,7 +147,7 @@ class MyCli(object):
         # Register custom special commands.
         self.register_special_commands()
 
-        self.cli = None
+        self.prompt_app = None
 
     def register_special_commands(self):
         special.register_special_command(self.change_db, 'use',
@@ -357,37 +360,39 @@ class MyCli(object):
             self.echo(str(e), err=True, fg='red')
             exit(1)
 
-    def handle_editor_command(self, cli, document):
-        """
-        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:
+    def handle_editor_command(self, text):
+        """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.
-        :param cli: CommandLineInterface
-        :param document: Document
+        :param text: Document
         :return: Document
+
         """
-        # FIXME: using application.pre_run_callables like this here is not the best solution.
-        # It's internal api of prompt_toolkit that may change. This was added to fix
-        # https://github.com/dbcli/pgcli/issues/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
+
+        while special.editor_command(text):
+            filename = special.get_filename(text)
+            query = (special.get_editor_query(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()
+            while True:
+                try:
+                    text = self.prompt_app.prompt(
+                        default=sql,
+                        vi_mode=self.key_bindings == 'vi'
+                    )
+                    break
+                except KeyboardInterrupt:
+                    sql = None
+
             continue
-        cli.application.pre_run_callables = saved_callables
-        return document
+        return text
 
     def run_cli(self):
         iterations = 0
@@ -412,7 +417,7 @@ class MyCli(object):
                 'Your query history will not be saved.'.format(history_file),
                 err=True, fg='red')
 
-        key_binding_manager = mycli_bindings()
+        key_bindings = mycli_bindings(self)
 
         if not self.less_chatty:
             print(' '.join(sqlexecute.server_type()))
@@ -422,39 +427,44 @@ class MyCli(object):
             print('Home: http://mycli.net')
             print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file]))
 
-        def prompt_tokens(cli):
+        def get_message():
             prompt = self.get_prompt(self.prompt_format)
             if (self.prompt_format == self.config.default_config['main']['prompt']
                     and len(prompt) > self.max_len_prompt):
                 prompt = self.get_prompt('\\d> ')
-            return [(Token.Prompt, prompt)]
+            return [('class:prompt', prompt)]
 
-        def get_continuation_tokens(cli, width):
-            continuation_prompt = self.get_prompt(self.prompt_continuation_format)
-            return [(Token.Continuation, ' ' * (width - len(continuation_prompt)) + continuation_prompt)]
+        def get_continuation(width, line_number, is_soft_wrap):
+            continuation = ' ' * (width - 1) + ' '
+            return [('class:continuation', continuation)]
 
         def show_suggestion_tip():
             return iterations < 2
 
-        def one_iteration(document=None):
-            if document is None:
-                document = self.cli.run()
+        def one_iteration(text=None):
+            if text is None:
+                try:
+                    text = self.prompt_app.prompt(
+                        vi_mode=self.key_bindings == 'vi'
+                    )
+                except KeyboardInterrupt:
+                    return
 
                 special.set_expanded_output(False)
 
                 try:
-                    document = self.handle_editor_command(self.cli, document)
+                    text = self.handle_editor_command(text)
                 except RuntimeError as e:
-                    logger.error("sql: %r, error: %r", document.text, e)
+                    logger.error("sql: %r, error: %r", text, e)
                     logger.error("traceback: %r", traceback.format_exc())
                     self.echo(str(e), err=True, fg='red')
                     return
 
-            if not document.text.strip():
+            if not text.strip():
                 return
 
             if self.destructive_warning:
-                destroy = confirm_destructive_query(document.text)
+                destroy = confirm_destructive_query(text)
                 if destroy is None:
                     pass  # Query was not destructive. Nothing to do here.
                 elif destroy is True:
@@ -469,18 +479,18 @@ class MyCli(object):
             mutating = False
 
             try:
-                logger.debug('sql: %r', document.text)
+                logger.debug('sql: %r', text)
 
-                special.write_tee(self.get_prompt(self.prompt_format) + document.text)
+                special.write_tee(self.get_prompt(self.prompt_format) + text)
                 if self.logfile:
                     self.logfile.write('\n# %s\n' % datetime.now())
-                    self.logfile.write(document.text)
+                    self.logfile.write(text)
                     self.logfile.write('\n')
 
                 successful = False
                 start = time()
-                res = sqlexecute.run(document.text)
-                self.formatter.query = document.text
+                res = sqlexecute.run(text)
+                self.formatter.query = text
                 successful = True
                 result_count = 0
                 for title, cur, headers, status in res:
@@ -497,7 +507,7 @@ class MyCli(object):
                             break
 
                     if self.auto_vertical_output:
-                        max_width = self.cli.output.get_size().columns
+                        max_width = self.prompt_app.output.get_size().columns
                     else:
                         max_width = None
 
@@ -535,7 +545,7 @@ class MyCli(object):
                         status_str = str(status).lower()
                         if status_str.find('ok') > -1:
                             logger.debug("cancelled query, connection id: %r, sql: %r",
-                                         connection_id_to_kill, document.text)
+                                         connection_id_to_kill, text)
                             self.echo("cancelled query", err=True, fg='red')
                 except Exception as e:
                     self.echo('Encountered error while