api.py 9.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352
  1. from __future__ import absolute_import
  2. from __future__ import division
  3. from __future__ import print_function
  4. from __future__ import unicode_literals
  5. from collections import namedtuple
  6. import itertools
  7. import json
  8. from six import string_types
  9. from six.moves.urllib import parse
  10. import requests
  11. from pydruid.db import exceptions
  12. class Type(object):
  13. STRING = 1
  14. NUMBER = 2
  15. BOOLEAN = 3
  16. def connect(host='localhost', port=8082, path='/druid/v2/sql/', scheme='http'):
  17. """
  18. Constructor for creating a connection to the database.
  19. >>> conn = connect('localhost', 8082)
  20. >>> curs = conn.cursor()
  21. """
  22. return Connection(host, port, path, scheme)
  23. def check_closed(f):
  24. """Decorator that checks if connection/cursor is closed."""
  25. def g(self, *args, **kwargs):
  26. if self.closed:
  27. raise exceptions.Error(
  28. '{klass} already closed'.format(klass=self.__class__.__name__))
  29. return f(self, *args, **kwargs)
  30. return g
  31. def check_result(f):
  32. """Decorator that checks if the cursor has results from `execute`."""
  33. def g(self, *args, **kwargs):
  34. if self._results is None:
  35. raise exceptions.Error('Called before `execute`')
  36. return f(self, *args, **kwargs)
  37. return g
  38. def get_description_from_row(row):
  39. """
  40. Return description from a single row.
  41. We only return the name, type (inferred from the data) and if the values
  42. can be NULL. String columns in Druid are NULLable. Numeric columns are NOT
  43. NULL.
  44. """
  45. return [
  46. (
  47. name, # name
  48. get_type(value), # type_code
  49. None, # [display_size]
  50. None, # [internal_size]
  51. None, # [precision]
  52. None, # [scale]
  53. get_type(value) == Type.STRING, # [null_ok]
  54. )
  55. for name, value in row.items()
  56. ]
  57. def get_type(value):
  58. """Infer type from value."""
  59. if isinstance(value, string_types) or value is None:
  60. return Type.STRING
  61. elif isinstance(value, (int, float)):
  62. return Type.NUMBER
  63. elif isinstance(value, bool):
  64. return Type.BOOLEAN
  65. raise exceptions.Error(
  66. 'Value of unknown type: {value}'.format(value=value))
  67. class Connection(object):
  68. """Connection to a Druid database."""
  69. def __init__(
  70. self,
  71. host='localhost',
  72. port=8082,
  73. path='/druid/v2/sql/',
  74. scheme='http',
  75. ):
  76. netloc = '{host}:{port}'.format(host=host, port=port)
  77. self.url = parse.urlunparse(
  78. (scheme, netloc, path, None, None, None))
  79. self.closed = False
  80. self.cursors = []
  81. @check_closed
  82. def close(self):
  83. """Close the connection now."""
  84. self.closed = True
  85. for cursor in self.cursors:
  86. try:
  87. cursor.close()
  88. except exceptions.Error:
  89. pass # already closed
  90. @check_closed
  91. def commit(self):
  92. """
  93. Commit any pending transaction to the database.
  94. Not supported.
  95. """
  96. pass
  97. @check_closed
  98. def cursor(self):
  99. """Return a new Cursor Object using the connection."""
  100. cursor = Cursor(self.url)
  101. self.cursors.append(cursor)
  102. return cursor
  103. @check_closed
  104. def execute(self, operation, parameters=None):
  105. cursor = self.cursor()
  106. return cursor.execute(operation, parameters)
  107. def __enter__(self):
  108. return self.cursor()
  109. def __exit__(self, *exc):
  110. self.close()
  111. class Cursor(object):
  112. """Connection cursor."""
  113. def __init__(self, url):
  114. self.url = url
  115. # This read/write attribute specifies the number of rows to fetch at a
  116. # time with .fetchmany(). It defaults to 1 meaning to fetch a single
  117. # row at a time.
  118. self.arraysize = 1
  119. self.closed = False
  120. # this is updated only after a query
  121. self.description = None
  122. # this is set to an iterator after a successfull query
  123. self._results = None
  124. @property
  125. @check_result
  126. @check_closed
  127. def rowcount(self):
  128. # consume the iterator
  129. results = list(self._results)
  130. n = len(results)
  131. self._results = iter(results)
  132. return n
  133. @check_closed
  134. def close(self):
  135. """Close the cursor."""
  136. self.closed = True
  137. @check_closed
  138. def execute(self, operation, parameters=None):
  139. query = apply_parameters(operation, parameters or {})
  140. # `_stream_query` returns a generator that produces the rows; we need
  141. # to consume the first row so that `description` is properly set, so
  142. # let's consume it and insert it back.
  143. results = self._stream_query(query)
  144. try:
  145. first_row = next(results)
  146. self._results = itertools.chain([first_row], results)
  147. except StopIteration:
  148. self._results = iter([])
  149. return self
  150. @check_closed
  151. def executemany(self, operation, seq_of_parameters=None):
  152. raise exceptions.NotSupportedError(
  153. '`executemany` is not supported, use `execute` instead')
  154. @check_result
  155. @check_closed
  156. def fetchone(self):
  157. """
  158. Fetch the next row of a query result set, returning a single sequence,
  159. or `None` when no more data is available.
  160. """
  161. try:
  162. return self.next()
  163. except StopIteration:
  164. return None
  165. @check_result
  166. @check_closed
  167. def fetchmany(self, size=None):
  168. """
  169. Fetch the next set of rows of a query result, returning a sequence of
  170. sequences (e.g. a list of tuples). An empty sequence is returned when
  171. no more rows are available.
  172. """
  173. size = size or self.arraysize
  174. return list(itertools.islice(self, size))
  175. @check_result
  176. @check_closed
  177. def fetchall(self):
  178. """
  179. Fetch all (remaining) rows of a query result, returning them as a
  180. sequence of sequences (e.g. a list of tuples). Note that the cursor's
  181. arraysize attribute can affect the performance of this operation.
  182. """
  183. return list(self)
  184. @check_closed
  185. def setinputsizes(self, sizes):
  186. # not supported
  187. pass
  188. @check_closed
  189. def setoutputsizes(self, sizes):
  190. # not supported
  191. pass
  192. @check_closed
  193. def __iter__(self):
  194. return self
  195. @check_closed
  196. def __next__(self):
  197. return next(self._results)
  198. next = __next__
  199. def _stream_query(self, query):
  200. """
  201. Stream rows from a query.
  202. This method will yield rows as the data is returned in chunks from the
  203. server.
  204. """
  205. self.description = None
  206. headers = {'Content-Type': 'application/json'}
  207. payload = {'query': query}
  208. r = requests.post(self.url, stream=True, headers=headers, json=payload)
  209. if r.encoding is None:
  210. r.encoding = 'utf-8'
  211. # raise any error messages
  212. if r.status_code != 200:
  213. payload = r.json()
  214. msg = (
  215. '{error} ({errorClass}): {errorMessage}'.format(**payload)
  216. )
  217. raise exceptions.ProgrammingError(msg)
  218. # Druid will stream the data in chunks of 8k bytes, splitting the JSON
  219. # between them; setting `chunk_size` to `None` makes it use the server
  220. # size
  221. chunks = r.iter_content(chunk_size=None, decode_unicode=True)
  222. Row = None
  223. for row in rows_from_chunks(chunks):
  224. # update description
  225. if self.description is None:
  226. self.description = get_description_from_row(row)
  227. # return row in namedtuple
  228. if Row is None:
  229. Row = namedtuple('Row', row.keys(), rename=True)
  230. yield Row(*row.values())
  231. def rows_from_chunks(chunks):
  232. """
  233. A generator that yields rows from JSON chunks.
  234. Druid will return the data in chunks, but they are not aligned with the
  235. JSON objects. This function will parse all complete rows inside each chunk,
  236. yielding them as soon as possible.
  237. """
  238. body = ''
  239. for chunk in chunks:
  240. if chunk:
  241. body = ''.join((body, chunk))
  242. # find last complete row
  243. boundary = 0
  244. brackets = 0
  245. in_string = False
  246. for i, char in enumerate(body):
  247. if char == '"':
  248. if not in_string:
  249. in_string = True
  250. elif body[i - 1] != '\\':
  251. in_string = False
  252. if in_string:
  253. continue
  254. if char == '{':
  255. brackets += 1
  256. elif char == '}':
  257. brackets -= 1
  258. if brackets == 0 and i > boundary:
  259. boundary = i + 1
  260. rows = body[:boundary].lstrip('[,')
  261. body = body[boundary:]
  262. for row in json.loads('[{rows}]'.format(rows=rows)):
  263. yield row
  264. def apply_parameters(operation, parameters):
  265. escaped_parameters = {
  266. key: escape(value) for key, value in parameters.items()
  267. }
  268. return operation % escaped_parameters
  269. def escape(value):
  270. if value == '*':
  271. return value
  272. elif isinstance(value, string_types):
  273. return "'{}'".format(value.replace("'", "''"))
  274. elif isinstance(value, (int, float)):
  275. return value
  276. elif isinstance(value, bool):
  277. return 'TRUE' if value else 'FALSE'
  278. elif isinstance(value, (list, tuple)):
  279. return ', '.join(escape(element) for element in value)