add un/prettify of current statement via sqlglot

10 files changed, 163 insertions, 5 deletions

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index b678f57..b7d1d38 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -9,10 +9,8 @@ jobs:
   linux:
     strategy:
       matrix:
-        python-version: [3.6, 3.7, 3.8, 3.9]
+        python-version: [3.7, 3.8, 3.9]
         include:
-          - python-version: 3.6
-            os: ubuntu-18.04 # MySQL 5.7.32
           - python-version: 3.7
             os: ubuntu-18.04 # MySQL 5.7.32
           - python-version: 3.8
diff --git a/README.md b/README.md
index e77ba7b..9e177b7 100644
--- a/README.md
+++ b/README.md
@@ -135,6 +135,7 @@ Features
 * Log every query and its results to a file (disabled by default).
 * Pretty prints tabular data (with colors!)
 * Support for SSL connections
+* Some features are only exposed as [key bindings](doc/key_bindings.rst)
 
 Contributions:
 --------------
@@ -219,7 +220,7 @@ Thanks to [PyMysql](https://github.com/PyMySQL/PyMySQL) for a pure python adapte
 
 ### Compatibility
 
-Mycli is tested on macOS and Linux.
+Mycli is tested on macOS and Linux, and requires Python 3.7 or better.
 
 **Mycli is not tested on Windows**, but the libraries used in this app are Windows-compatible.
 This means it should work without any modifications. If you're unable to run it
diff --git a/changelog.md b/changelog.md
index d79389b..6685329 100644
--- a/changelog.md
+++ b/changelog.md
@@ -7,6 +7,7 @@ Features:
 
 * Add `--ssl` flag to enable ssl/tls.
 * Add `pager` option to `~/.myclirc`, for instance `pager = 'pspg --csv'` (Thanks: [BuonOmo])
+* Add prettify/unprettify keybindings to format the current statement using `sqlglot`.
 
 
 Internal:
diff --git a/doc/key_bindings.rst b/doc/key_bindings.rst
new file mode 100644
index 0000000..0534870
--- /dev/null
+++ b/doc/key_bindings.rst
@@ -0,0 +1,65 @@
+*************
+Key Bindings:
+*************
+
+Most key bindings are simply inherited from `prompt-toolkit <https://python-prompt-toolkit.readthedocs.io/en/master/index.html>`_ .
+
+The following key bindings are special to mycli:
+
+###
+F2
+###
+
+Enable/Disable SmartCompletion Mode.
+
+###
+F3
+###
+
+Enable/Disable Multiline Mode.
+
+###
+F4
+###
+
+Toggle between Vi and Emacs mode.
+
+###
+Tab
+###
+
+Force autocompletion at cursor.
+
+#######
+C-space
+#######
+
+Initialize autocompletion at cursor.
+
+If the autocompletion menu is not showing, display it with the appropriate completions for the context.
+
+If the menu is showing, select the next completion.
+
+#########
+ESC Enter
+#########
+
+Introduce a line break in multi-line mode, or dispatch the command in single-line mode.
+
+The sequence ESC-Enter is often sent by Alt-Enter.
+
+#################################
+C-x p (Emacs-mode) or > (Vi-mode)
+#################################
+
+Prettify and indent current statement, usually into multiple lines.
+
+Only accepts buffers containing single SQL statements.
+
+#################################
+C-x u (Emacs-mode) or < (Vi-mode)
+#################################
+
+Unprettify and dedent current statement, usually into one line.
+
+Only accepts buffers containing single SQL statements.
diff --git a/mycli/clitoolbar.py b/mycli/clitoolbar.py
index eec2978..24d1108 100644
--- a/mycli/clitoolbar.py
+++ b/mycli/clitoolbar.py
@@ -30,6 +30,11 @@ def create_toolbar_tokens_func(mycli, show_fish_help):
                 'Vi-mode ({})'.format(_get_vi_mode())
             ))
 
+        if mycli.toolbar_error_message:
+            result.append(
+                ('class:bottom-toolbar', '  ' + mycli.toolbar_error_message))
+            mycli.toolbar_error_message = None
+
         if show_fish_help():
             result.append(
                 ('class:bottom-toolbar', '  Right-arrow to complete suggestion'))
diff --git a/mycli/key_bindings.py b/mycli/key_bindings.py
index 4a24c82..03e4ace 100644
--- a/mycli/key_bindings.py
+++ b/mycli/key_bindings.py
@@ -1,6 +1,6 @@
 import logging
 from prompt_toolkit.enums import EditingMode
-from prompt_toolkit.filters import completion_is_selected
+from prompt_toolkit.filters import completion_is_selected, emacs_mode, vi_mode
 from prompt_toolkit.key_binding import KeyBindings
 
 _logger = logging.getLogger(__name__)
@@ -61,6 +61,48 @@ def mycli_bindings(mycli):
         else:
             b.start_completion(select_first=False)
 
+    @kb.add('>', filter=vi_mode)
+    @kb.add('c-x', 'p', filter=emacs_mode)
+    def _(event):
+        """
+        Prettify and indent current statement, usually into multiple lines.
+
+        Only accepts buffers containing single SQL statements.
+        """
+        _logger.debug('Detected <C-x p>/> key.')
+
+        b = event.app.current_buffer
+        cursorpos_relative = b.cursor_position / len(b.text)
+        pretty_text = mycli.handle_prettify_binding(b.text)
+        if len(pretty_text) > 0:
+            b.text = pretty_text
+            cursorpos_abs = int(round(cursorpos_relative * len(b.text)))
+            while 0 < cursorpos_abs < len(b.text) \
+                  and b.text[cursorpos_abs] in (' ', '\n'):
+                cursorpos_abs -= 1
+            b.cursor_position = min(cursorpos_abs, len(b.text))
+
+    @kb.add('<', filter=vi_mode)
+    @kb.add('c-x', 'u', filter=emacs_mode)
+    def _(event):
+        """
+        Unprettify and dedent current statement, usually into one line.
+
+        Only accepts buffers containing single SQL statements.
+        """
+        _logger.debug('Detected <C-x u>/< key.')
+
+        b = event.app.current_buffer
+        cursorpos_relative = b.cursor_position / len(b.text)
+        unpretty_text = mycli.handle_unprettify_binding(b.text)
+        if len(unpretty_text) > 0:
+            b.text = unpretty_text
+            cursorpos_abs = int(round(cursorpos_relative * len(b.text)))
+            while 0 < cursorpos_abs < len(b.text) \
+                  and b.text[cursorpos_abs] in (' ', '\n'):
+                cursorpos_abs -= 1
+            b.cursor_position = min(cursorpos_abs, len(b.text))
+
     @kb.add('enter', filter=completion_is_selected)
     def _(event):
         """Makes the enter key work as the tab key only when showing the menu.
diff --git a/mycli/main.py b/mycli/main.py
index 64d00f0..208572d 100755
--- a/mycli/main.py
+++ b/mycli/main.py
@@ -24,6 +24,7 @@ from cli_helpers.tabular_output import preprocessors
 from cli_helpers.utils import strip_ansi
 import click
 import sqlparse
+import sqlglot
 from mycli.packages.parseutils import is_dropping_database, is_destructive
 from prompt_toolkit.completion import DynamicCompleter
 from prompt_toolkit.enums import DEFAULT_BUFFER, EditingMode
@@ -123,6 +124,7 @@ class MyCli(object):
         self.logfile = logfile
         self.defaults_suffix = defaults_suffix
         self.login_path = login_path
+        self.toolbar_error_message = None
 
         # self.cnf_files is a class variable that stores the list of mysql
         # config files to read in at launch.
@@ -582,6 +584,34 @@ class MyCli(object):
             return True
         return False
 
+    def handle_prettify_binding(self, text):
+        try:
+            statements = sqlglot.parse(text, read='mysql')
+        except Exception as e:
+            statements = []
+        if len(statements) == 1:
+            pretty_text = statements[0].sql(pretty=True, pad=4, dialect='mysql')
+        else:
+            pretty_text = ''
+            self.toolbar_error_message = 'Prettify failed to parse statement'
+        if len(pretty_text) > 0:
+            pretty_text = pretty_text + ';'
+        return pretty_text
+
+    def handle_unprettify_binding(self, text):
+        try:
+            statements = sqlglot.parse(text, read='mysql')
+        except Exception as e:
+            statements = []
+        if len(statements) == 1:
+            unpretty_text = statements[0].sql(pretty=False, dialect='mysql')
+        else:
+            unpretty_text = ''
+            self.toolbar_error_message = 'Unprettify failed to parse statement'
+        if len(unpretty_text) > 0:
+            unpretty_text = unpretty_text + ';'
+        return unpretty_text
+
     def run_cli(self):
         iterations = 0
         sqlexecute = self.sqlexecute
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 3d10e9a..955a9f5 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -14,3 +14,4 @@ paramiko==2.11.0
 pyperclip>=1.8.1
 importlib_resources>=5.0.0
 pyaes>=1.6.1
+sqlglot>=5.1.3
diff --git a/setup.py b/setup.py
index c9dcc44..20afeef 100755
--- a/setup.py
+++ b/setup.py
@@ -26,6 +26,7 @@ install_requirements = [
     'prompt_toolkit>=3.0.6,<4.0.0',
     'PyMySQL >= 0.9.2',
     'sqlparse>=0.3.0,<0.5.0',
+    'sqlglot>=5.1.3',
     'configobj >= 5.0.5',
     'cli_helpers[styles] >= 2.2.1',
     'pyperclip >= 1.8.1',
diff --git a/test/test_main.py b/test/test_main.py
index c3351ec..64cba0a 100644
--- a/test/test_main.py
+++ b/test/test_main.py
@@ -283,6 +283,20 @@ def test_list_dsn():
         assert result.output == "test : mysql://test/test\n"
 
 
+def test_prettify_statement():
+    statement = 'SELECT 1'
+    m = MyCli()
+    pretty_statement = m.handle_prettify_binding(statement)
+    assert pretty_statement == 'SELECT\n    1;'
+
+
+def test_unprettify_statement():
+    statement = 'SELECT\n    1'
+    m = MyCli()
+    unpretty_statement = m.handle_unprettify_binding(statement)
+    assert unpretty_statement == 'SELECT 1;'
+
+
 def test_list_ssh_config():
     runner = CliRunner()
     with NamedTemporaryFile(mode="w") as ssh_config: