| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352 |
- from __future__ import absolute_import
- from __future__ import division
- from __future__ import print_function
- from __future__ import unicode_literals
- from collections import namedtuple
- import itertools
- import json
- from six import string_types
- from six.moves.urllib import parse
- import requests
- from pydruid.db import exceptions
- class Type(object):
- STRING = 1
- NUMBER = 2
- BOOLEAN = 3
- def connect(host='localhost', port=8082, path='/druid/v2/sql/', scheme='http'):
- """
- Constructor for creating a connection to the database.
- >>> conn = connect('localhost', 8082)
- >>> curs = conn.cursor()
- """
- return Connection(host, port, path, scheme)
- def check_closed(f):
- """Decorator that checks if connection/cursor is closed."""
- def g(self, *args, **kwargs):
- if self.closed:
- raise exceptions.Error(
- '{klass} already closed'.format(klass=self.__class__.__name__))
- return f(self, *args, **kwargs)
- return g
- def check_result(f):
- """Decorator that checks if the cursor has results from `execute`."""
- def g(self, *args, **kwargs):
- if self._results is None:
- raise exceptions.Error('Called before `execute`')
- return f(self, *args, **kwargs)
- return g
- def get_description_from_row(row):
- """
- Return description from a single row.
- We only return the name, type (inferred from the data) and if the values
- can be NULL. String columns in Druid are NULLable. Numeric columns are NOT
- NULL.
- """
- return [
- (
- name, # name
- get_type(value), # type_code
- None, # [display_size]
- None, # [internal_size]
- None, # [precision]
- None, # [scale]
- get_type(value) == Type.STRING, # [null_ok]
- )
- for name, value in row.items()
- ]
- def get_type(value):
- """Infer type from value."""
- if isinstance(value, string_types) or value is None:
- return Type.STRING
- elif isinstance(value, (int, float)):
- return Type.NUMBER
- elif isinstance(value, bool):
- return Type.BOOLEAN
- raise exceptions.Error(
- 'Value of unknown type: {value}'.format(value=value))
- class Connection(object):
- """Connection to a Druid database."""
- def __init__(
- self,
- host='localhost',
- port=8082,
- path='/druid/v2/sql/',
- scheme='http',
- ):
- netloc = '{host}:{port}'.format(host=host, port=port)
- self.url = parse.urlunparse(
- (scheme, netloc, path, None, None, None))
- self.closed = False
- self.cursors = []
- @check_closed
- def close(self):
- """Close the connection now."""
- self.closed = True
- for cursor in self.cursors:
- try:
- cursor.close()
- except exceptions.Error:
- pass # already closed
- @check_closed
- def commit(self):
- """
- Commit any pending transaction to the database.
- Not supported.
- """
- pass
- @check_closed
- def cursor(self):
- """Return a new Cursor Object using the connection."""
- cursor = Cursor(self.url)
- self.cursors.append(cursor)
- return cursor
- @check_closed
- def execute(self, operation, parameters=None):
- cursor = self.cursor()
- return cursor.execute(operation, parameters)
- def __enter__(self):
- return self.cursor()
- def __exit__(self, *exc):
- self.close()
- class Cursor(object):
- """Connection cursor."""
- def __init__(self, url):
- self.url = url
- # This read/write attribute specifies the number of rows to fetch at a
- # time with .fetchmany(). It defaults to 1 meaning to fetch a single
- # row at a time.
- self.arraysize = 1
- self.closed = False
- # this is updated only after a query
- self.description = None
- # this is set to an iterator after a successfull query
- self._results = None
- @property
- @check_result
- @check_closed
- def rowcount(self):
- # consume the iterator
- results = list(self._results)
- n = len(results)
- self._results = iter(results)
- return n
- @check_closed
- def close(self):
- """Close the cursor."""
- self.closed = True
- @check_closed
- def execute(self, operation, parameters=None):
- query = apply_parameters(operation, parameters or {})
- # `_stream_query` returns a generator that produces the rows; we need
- # to consume the first row so that `description` is properly set, so
- # let's consume it and insert it back.
- results = self._stream_query(query)
- try:
- first_row = next(results)
- self._results = itertools.chain([first_row], results)
- except StopIteration:
- self._results = iter([])
- return self
- @check_closed
- def executemany(self, operation, seq_of_parameters=None):
- raise exceptions.NotSupportedError(
- '`executemany` is not supported, use `execute` instead')
- @check_result
- @check_closed
- def fetchone(self):
- """
- Fetch the next row of a query result set, returning a single sequence,
- or `None` when no more data is available.
- """
- try:
- return self.next()
- except StopIteration:
- return None
- @check_result
- @check_closed
- def fetchmany(self, size=None):
- """
- Fetch the next set of rows of a query result, returning a sequence of
- sequences (e.g. a list of tuples). An empty sequence is returned when
- no more rows are available.
- """
- size = size or self.arraysize
- return list(itertools.islice(self, size))
- @check_result
- @check_closed
- def fetchall(self):
- """
- Fetch all (remaining) rows of a query result, returning them as a
- sequence of sequences (e.g. a list of tuples). Note that the cursor's
- arraysize attribute can affect the performance of this operation.
- """
- return list(self)
- @check_closed
- def setinputsizes(self, sizes):
- # not supported
- pass
- @check_closed
- def setoutputsizes(self, sizes):
- # not supported
- pass
- @check_closed
- def __iter__(self):
- return self
- @check_closed
- def __next__(self):
- return next(self._results)
- next = __next__
- def _stream_query(self, query):
- """
- Stream rows from a query.
- This method will yield rows as the data is returned in chunks from the
- server.
- """
- self.description = None
- headers = {'Content-Type': 'application/json'}
- payload = {'query': query}
- r = requests.post(self.url, stream=True, headers=headers, json=payload)
- if r.encoding is None:
- r.encoding = 'utf-8'
- # raise any error messages
- if r.status_code != 200:
- payload = r.json()
- msg = (
- '{error} ({errorClass}): {errorMessage}'.format(**payload)
- )
- raise exceptions.ProgrammingError(msg)
- # Druid will stream the data in chunks of 8k bytes, splitting the JSON
- # between them; setting `chunk_size` to `None` makes it use the server
- # size
- chunks = r.iter_content(chunk_size=None, decode_unicode=True)
- Row = None
- for row in rows_from_chunks(chunks):
- # update description
- if self.description is None:
- self.description = get_description_from_row(row)
- # return row in namedtuple
- if Row is None:
- Row = namedtuple('Row', row.keys(), rename=True)
- yield Row(*row.values())
- def rows_from_chunks(chunks):
- """
- A generator that yields rows from JSON chunks.
- Druid will return the data in chunks, but they are not aligned with the
- JSON objects. This function will parse all complete rows inside each chunk,
- yielding them as soon as possible.
- """
- body = ''
- for chunk in chunks:
- if chunk:
- body = ''.join((body, chunk))
- # find last complete row
- boundary = 0
- brackets = 0
- in_string = False
- for i, char in enumerate(body):
- if char == '"':
- if not in_string:
- in_string = True
- elif body[i - 1] != '\\':
- in_string = False
- if in_string:
- continue
- if char == '{':
- brackets += 1
- elif char == '}':
- brackets -= 1
- if brackets == 0 and i > boundary:
- boundary = i + 1
- rows = body[:boundary].lstrip('[,')
- body = body[boundary:]
- for row in json.loads('[{rows}]'.format(rows=rows)):
- yield row
- def apply_parameters(operation, parameters):
- escaped_parameters = {
- key: escape(value) for key, value in parameters.items()
- }
- return operation % escaped_parameters
- def escape(value):
- if value == '*':
- return value
- elif isinstance(value, string_types):
- return "'{}'".format(value.replace("'", "''"))
- elif isinstance(value, (int, float)):
- return value
- elif isinstance(value, bool):
- return 'TRUE' if value else 'FALSE'
- elif isinstance(value, (list, tuple)):
- return ', '.join(escape(element) for element in value)
|