console.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202
  1. from __future__ import unicode_literals
  2. import os
  3. import re
  4. import sys
  5. from prompt_toolkit import prompt, AbortAction
  6. from prompt_toolkit.history import FileHistory
  7. from prompt_toolkit.contrib.completers import WordCompleter
  8. from pygments.lexers import SqlLexer
  9. from pygments.style import Style
  10. from pygments.token import Token
  11. from pygments.styles.default import DefaultStyle
  12. from six.moves.urllib import parse
  13. from tabulate import tabulate
  14. from pydruid.db.api import connect
  15. keywords = [
  16. 'EXPLAIN PLAN FOR',
  17. 'WITH',
  18. 'SELECT',
  19. 'ALL',
  20. 'DISTINCT',
  21. 'FROM',
  22. 'WHERE',
  23. 'GROUP BY',
  24. 'HAVING',
  25. 'ORDER BY',
  26. 'ASC',
  27. 'DESC',
  28. 'LIMIT',
  29. ]
  30. aggregate_functions = [
  31. 'COUNT',
  32. 'SUM',
  33. 'MIN',
  34. 'MAX',
  35. 'AVG',
  36. 'APPROX_COUNT_DISTINCT',
  37. 'APPROX_QUANTILE',
  38. ]
  39. numeric_functions = [
  40. 'ABS',
  41. 'CEIL',
  42. 'EXP',
  43. 'FLOOR',
  44. 'LN',
  45. 'LOG10',
  46. 'POW',
  47. 'SQRT',
  48. ]
  49. string_functions = [
  50. 'CHARACTER_LENGTH',
  51. 'LOOKUP',
  52. 'LOWER',
  53. 'REGEXP_EXTRACT',
  54. 'REPLACE',
  55. 'SUBSTRING',
  56. 'TRIM',
  57. 'BTRIM',
  58. 'RTRIM',
  59. 'LTRIM',
  60. 'UPPER',
  61. ]
  62. time_functions = [
  63. 'CURRENT_TIMESTAMP',
  64. 'CURRENT_DATE',
  65. 'TIME_FLOOR',
  66. 'TIME_SHIFT',
  67. 'TIME_EXTRACT',
  68. 'TIME_PARSE',
  69. 'TIME_FORMAT',
  70. 'MILLIS_TO_TIMESTAMP',
  71. 'TIMESTAMP_TO_MILLIS',
  72. 'EXTRACT',
  73. 'FLOOR',
  74. 'CEIL',
  75. ]
  76. other_functions = [
  77. 'CAST',
  78. 'CASE',
  79. 'WHEN',
  80. 'THEN',
  81. 'END',
  82. 'NULLIF',
  83. 'COALESCE',
  84. ]
  85. replacements = {
  86. '^SHOW SCHEMAS': 'SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA',
  87. '^SHOW TABLES': 'SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES',
  88. '^DESC (?P<table>[^;\s]*)': r"""
  89. SELECT COLUMN_NAME,
  90. ORDINAL_POSITION,
  91. COLUMN_DEFAULT,
  92. IS_NULLABLE,
  93. DATA_TYPE
  94. FROM INFORMATION_SCHEMA.COLUMNS
  95. WHERE TABLE_NAME='\1'
  96. """.strip(),
  97. }
  98. class DocumentStyle(Style):
  99. styles = {
  100. Token.Menu.Completions.Completion.Current: 'bg:#00aaaa #000000',
  101. Token.Menu.Completions.Completion: 'bg:#008888 #ffffff',
  102. Token.Menu.Completions.ProgressButton: 'bg:#003333',
  103. Token.Menu.Completions.ProgressBar: 'bg:#00aaaa',
  104. }
  105. styles.update(DefaultStyle.styles)
  106. def get_connection_kwargs(url):
  107. parts = parse.urlparse(url)
  108. if ':' in parts.netloc:
  109. host, port = parts.netloc.split(':', 1)
  110. port = int(port)
  111. else:
  112. host = parts.netloc
  113. port = 8082
  114. return {
  115. 'host': host,
  116. 'port': port,
  117. 'path': parts.path,
  118. 'scheme': parts.scheme,
  119. }
  120. def get_tables(connection):
  121. cursor = connection.cursor()
  122. return [
  123. row.TABLE_NAME for row in
  124. cursor.execute('SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES')
  125. ]
  126. def get_autocomplete(connection):
  127. return (
  128. keywords +
  129. aggregate_functions +
  130. numeric_functions +
  131. string_functions +
  132. time_functions +
  133. other_functions +
  134. get_tables(connection)
  135. )
  136. def main():
  137. history = FileHistory(os.path.expanduser('~/.pydruid_history'))
  138. try:
  139. url = sys.argv[1]
  140. except IndexError:
  141. url = 'http://localhost:8082/druid/v2/sql/'
  142. kwargs = get_connection_kwargs(url)
  143. connection = connect(**kwargs)
  144. cursor = connection.cursor()
  145. words = get_autocomplete(connection)
  146. sql_completer = WordCompleter(words, ignore_case=True)
  147. while True:
  148. try:
  149. query = prompt(
  150. '> ', lexer=SqlLexer, completer=sql_completer,
  151. style=DocumentStyle, history=history,
  152. on_abort=AbortAction.RETRY)
  153. except EOFError:
  154. break # Control-D pressed.
  155. # run query
  156. query = query.strip('; ')
  157. if query:
  158. # shortcuts
  159. for pattern, repl in replacements.items():
  160. query = re.sub(pattern, repl, query)
  161. try:
  162. result = cursor.execute(query)
  163. except Exception as e:
  164. print(e)
  165. continue
  166. headers = [t[0] for t in cursor.description or []]
  167. print(tabulate(result, headers=headers))
  168. print('GoodBye!')
  169. if __name__ == '__main__':
  170. main()