from __future__ import unicode_literals import os import re import sys from prompt_toolkit import prompt, AbortAction from prompt_toolkit.history import FileHistory from prompt_toolkit.contrib.completers import WordCompleter from pygments.lexers import SqlLexer from pygments.style import Style from pygments.token import Token from pygments.styles.default import DefaultStyle from six.moves.urllib import parse from tabulate import tabulate from pydruid.db.api import connect keywords = [ 'EXPLAIN PLAN FOR', 'WITH', 'SELECT', 'ALL', 'DISTINCT', 'FROM', 'WHERE', 'GROUP BY', 'HAVING', 'ORDER BY', 'ASC', 'DESC', 'LIMIT', ] aggregate_functions = [ 'COUNT', 'SUM', 'MIN', 'MAX', 'AVG', 'APPROX_COUNT_DISTINCT', 'APPROX_QUANTILE', ] numeric_functions = [ 'ABS', 'CEIL', 'EXP', 'FLOOR', 'LN', 'LOG10', 'POW', 'SQRT', ] string_functions = [ 'CHARACTER_LENGTH', 'LOOKUP', 'LOWER', 'REGEXP_EXTRACT', 'REPLACE', 'SUBSTRING', 'TRIM', 'BTRIM', 'RTRIM', 'LTRIM', 'UPPER', ] time_functions = [ 'CURRENT_TIMESTAMP', 'CURRENT_DATE', 'TIME_FLOOR', 'TIME_SHIFT', 'TIME_EXTRACT', 'TIME_PARSE', 'TIME_FORMAT', 'MILLIS_TO_TIMESTAMP', 'TIMESTAMP_TO_MILLIS', 'EXTRACT', 'FLOOR', 'CEIL', ] other_functions = [ 'CAST', 'CASE', 'WHEN', 'THEN', 'END', 'NULLIF', 'COALESCE', ] replacements = { '^SHOW SCHEMAS': 'SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA', '^SHOW TABLES': 'SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES', '^DESC (?P[^;\s]*)': r""" SELECT COLUMN_NAME, ORDINAL_POSITION, COLUMN_DEFAULT, IS_NULLABLE, DATA_TYPE FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME='\1' """.strip(), } class DocumentStyle(Style): styles = { Token.Menu.Completions.Completion.Current: 'bg:#00aaaa #000000', Token.Menu.Completions.Completion: 'bg:#008888 #ffffff', Token.Menu.Completions.ProgressButton: 'bg:#003333', Token.Menu.Completions.ProgressBar: 'bg:#00aaaa', } styles.update(DefaultStyle.styles) def get_connection_kwargs(url): parts = parse.urlparse(url) if ':' in parts.netloc: host, port = parts.netloc.split(':', 1) port = int(port) else: host = parts.netloc port = 8082 return { 'host': host, 'port': port, 'path': parts.path, 'scheme': parts.scheme, } def get_tables(connection): cursor = connection.cursor() return [ row.TABLE_NAME for row in cursor.execute('SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES') ] def get_autocomplete(connection): return ( keywords + aggregate_functions + numeric_functions + string_functions + time_functions + other_functions + get_tables(connection) ) def main(): history = FileHistory(os.path.expanduser('~/.pydruid_history')) try: url = sys.argv[1] except IndexError: url = 'http://localhost:8082/druid/v2/sql/' kwargs = get_connection_kwargs(url) connection = connect(**kwargs) cursor = connection.cursor() words = get_autocomplete(connection) sql_completer = WordCompleter(words, ignore_case=True) while True: try: query = prompt( '> ', lexer=SqlLexer, completer=sql_completer, style=DocumentStyle, history=history, on_abort=AbortAction.RETRY) except EOFError: break # Control-D pressed. # run query query = query.strip('; ') if query: # shortcuts for pattern, repl in replacements.items(): query = re.sub(pattern, repl, query) try: result = cursor.execute(query) except Exception as e: print(e) continue headers = [t[0] for t in cursor.description or []] print(tabulate(result, headers=headers)) print('GoodBye!') if __name__ == '__main__': main()