From 358520621328db6dabe48a59fb9842a1b98b987a Mon Sep 17 00:00:00 2001
From: Eric Lee <erhlee.bird@gmail.com>
Date: Tue, 2 Feb 2021 13:04:03 -0500
Subject: [PATCH] fix pydruid.console prompt_toolkit + pygments integration

---
 pydruid/console.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/pydruid/console.py b/pydruid/console.py
index 3218f6fe..d6cd31af 100644
--- a/pydruid/console.py
+++ b/pydruid/console.py
@@ -6,6 +6,8 @@
 from prompt_toolkit import prompt
 from prompt_toolkit.completion.word_completer import WordCompleter
 from prompt_toolkit.history import FileHistory
+from prompt_toolkit.lexers import PygmentsLexer
+from prompt_toolkit.styles.pygments import style_from_pygments_cls
 from pygments.lexers import SqlLexer
 from pygments.style import Style
 from pygments.styles.default import DefaultStyle
@@ -170,9 +172,9 @@ def main():
         try:
             query = prompt(
                 "> ",
-                lexer=SqlLexer,
+                lexer=PygmentsLexer(SqlLexer),
                 completer=sql_completer,
-                style=DocumentStyle,
+                style=style_from_pygments_cls(DocumentStyle),
                 history=history,
             )
         except (EOFError, KeyboardInterrupt):