| """DB-API implementation backed by Presto |
| |
| See http://www.python.org/dev/peps/pep-0249/ |
| |
| Many docstrings in this file are based on the PEP, which is in the public domain. |
| """ |
| |
| from __future__ import absolute_import |
| from __future__ import unicode_literals |
| |
| from builtins import object |
| from decimal import Decimal |
| |
| from pyhive import common |
| from pyhive.common import DBAPITypeObject |
| # Make all exceptions visible in this module per DB-API |
| from pyhive.exc import * # noqa |
| import base64 |
| import getpass |
| import datetime |
| import logging |
| import requests |
| from requests.auth import HTTPBasicAuth |
| import os |
| |
| try: # Python 3 |
| import urllib.parse as urlparse |
| except ImportError: # Python 2 |
| import urlparse |
| |
| |
| # PEP 249 module globals |
| apilevel = '2.0' |
| threadsafety = 2 # Threads may share the module and connections. |
| paramstyle = 'pyformat' # Python extended format codes, e.g. ...WHERE name=%(name)s |
| |
| _logger = logging.getLogger(__name__) |
| |
| TYPES_CONVERTER = { |
| "decimal": Decimal, |
| # As of Presto 0.69, binary data is returned as the varbinary type in base64 format |
| "varbinary": base64.b64decode |
| } |
| |
| class PrestoParamEscaper(common.ParamEscaper): |
| def escape_datetime(self, item, format): |
| _type = "timestamp" if isinstance(item, datetime.datetime) else "date" |
| formatted = super(PrestoParamEscaper, self).escape_datetime(item, format, 3) |
| return "{} {}".format(_type, formatted) |
| |
| |
| _escaper = PrestoParamEscaper() |
| |
| |
| def connect(*args, **kwargs): |
| """Constructor for creating a connection to the database. See class :py:class:`Connection` for |
| arguments. |
| |
| :returns: a :py:class:`Connection` object. |
| """ |
| return Connection(*args, **kwargs) |
| |
| |
| class Connection(object): |
| """Presto does not have a notion of a persistent connection. |
| |
| Thus, these objects are small stateless factories for cursors, which do all the real work. |
| """ |
| |
| def __init__(self, *args, **kwargs): |
| self._args = args |
| self._kwargs = kwargs |
| |
| def close(self): |
| """Presto does not have anything to close""" |
| # TODO cancel outstanding queries? |
| pass |
| |
| def commit(self): |
| """Presto does not support transactions""" |
| pass |
| |
| def cursor(self): |
| """Return a new :py:class:`Cursor` object using the connection.""" |
| return Cursor(*self._args, **self._kwargs) |
| |
| def rollback(self): |
| raise NotSupportedError("Presto does not have transactions") # pragma: no cover |
| |
| |
| class Cursor(common.DBAPICursor): |
| """These objects represent a database cursor, which is used to manage the context of a fetch |
| operation. |
| |
| Cursors are not isolated, i.e., any changes done to the database by a cursor are immediately |
| visible by other cursors or connections. |
| """ |
| |
| def __init__(self, host, port='8080', username=None, principal_username=None, catalog='hive', |
| schema='default', poll_interval=1, source='pyhive', session_props=None, |
| protocol='http', password=None, requests_session=None, requests_kwargs=None, |
| KerberosRemoteServiceName=None, KerberosPrincipal=None, |
| KerberosConfigPath=None, KerberosKeytabPath=None, |
| KerberosCredentialCachePath=None, KerberosUseCanonicalHostname=None): |
| """ |
| :param host: hostname to connect to, e.g. ``presto.example.com`` |
| :param port: int -- port, defaults to 8080 |
| :param username: string -- defaults to system user name |
| :param principal_username: string -- defaults to ``username`` argument if it exists, |
| else defaults to system user name |
| :param catalog: string -- defaults to ``hive`` |
| :param schema: string -- defaults to ``default`` |
| :param poll_interval: float -- how often to ask the Presto REST interface for a progress |
| update, defaults to a second |
| :param source: string -- arbitrary identifier (shows up in the Presto monitoring page) |
| :param protocol: string -- network protocol, valid options are ``http`` and ``https``. |
| defaults to ``http`` |
| :param password: string -- Deprecated. Defaults to ``None``. |
| Using BasicAuth, requires ``https``. |
| Prefer ``requests_kwargs={'auth': HTTPBasicAuth(username, password)}``. |
| May not be specified with ``requests_kwargs['auth']``. |
| :param requests_session: a ``requests.Session`` object for advanced usage. If absent, this |
| class will use the default requests behavior of making a new session per HTTP request. |
| Caller is responsible for closing session. |
| :param requests_kwargs: Additional ``**kwargs`` to pass to requests |
| :param KerberosRemoteServiceName: string -- Presto coordinator Kerberos service name. |
| This parameter is required for Kerberos authentiation. |
| :param KerberosPrincipal: string -- The principal to use when authenticating to |
| the Presto coordinator. |
| :param KerberosConfigPath: string -- Kerberos configuration file. |
| (default: /etc/krb5.conf) |
| :param KerberosKeytabPath: string -- Kerberos keytab file. |
| :param KerberosCredentialCachePath: string -- Kerberos credential cache. |
| :param KerberosUseCanonicalHostname: boolean -- Use the canonical hostname of the |
| Presto coordinator for the Kerberos service principal by first resolving the |
| hostname to an IP address and then doing a reverse DNS lookup for that IP address. |
| This is enabled by default. |
| """ |
| super(Cursor, self).__init__(poll_interval) |
| # Config |
| self._host = host |
| self._port = port |
| """ |
| Presto User Impersonation: https://docs.starburstdata.com/latest/security/impersonation.html |
| |
| User impersonation allows the execution of queries in Presto based on principal_username |
| argument, instead of executing the query as the account which authenticated against Presto. |
| (Usually a service account) |
| |
| Allows for a service account to authenticate with Presto, and then leverage the |
| principal_username as the user Presto will execute the query as. This is required by |
| applications that leverage authentication methods like SAML, where the application has a |
| username, but not a password to still leverage user specific Presto Resource Groups and |
| Authorization rules that would not be applied when only using a shared service account. |
| This also allows auditing of who is executing a query in these environments, instead of |
| having all queryes run by the shared service account. |
| """ |
| self._username = principal_username or username or getpass.getuser() |
| self._catalog = catalog |
| self._schema = schema |
| self._arraysize = 1 |
| self._poll_interval = poll_interval |
| self._source = source |
| self._session_props = session_props if session_props is not None else {} |
| self.last_query_id = None |
| |
| if protocol not in ('http', 'https'): |
| raise ValueError("Protocol must be http/https, was {!r}".format(protocol)) |
| self._protocol = protocol |
| |
| self._requests_session = requests_session or requests |
| |
| requests_kwargs = dict(requests_kwargs) if requests_kwargs is not None else {} |
| |
| if KerberosRemoteServiceName is not None: |
| from requests_kerberos import HTTPKerberosAuth, OPTIONAL |
| |
| hostname_override = None |
| if KerberosUseCanonicalHostname is not None \ |
| and KerberosUseCanonicalHostname.lower() == 'false': |
| hostname_override = host |
| if KerberosConfigPath is not None: |
| os.environ['KRB5_CONFIG'] = KerberosConfigPath |
| if KerberosKeytabPath is not None: |
| os.environ['KRB5_CLIENT_KTNAME'] = KerberosKeytabPath |
| if KerberosCredentialCachePath is not None: |
| os.environ['KRB5CCNAME'] = KerberosCredentialCachePath |
| |
| requests_kwargs['auth'] = HTTPKerberosAuth(mutual_authentication=OPTIONAL, |
| principal=KerberosPrincipal, |
| service=KerberosRemoteServiceName, |
| hostname_override=hostname_override) |
| |
| else: |
| if password is not None and 'auth' in requests_kwargs: |
| raise ValueError("Cannot use both password and requests_kwargs authentication") |
| for k in ('method', 'url', 'data', 'headers'): |
| if k in requests_kwargs: |
| raise ValueError("Cannot override requests argument {}".format(k)) |
| if password is not None: |
| requests_kwargs['auth'] = HTTPBasicAuth(username, password) |
| if protocol != 'https': |
| raise ValueError("Protocol must be https when passing a password") |
| self._requests_kwargs = requests_kwargs |
| |
| self._reset_state() |
| |
| def _reset_state(self): |
| """Reset state about the previous query in preparation for running another query""" |
| super(Cursor, self)._reset_state() |
| self._nextUri = None |
| self._columns = None |
| |
| @property |
| def description(self): |
| """This read-only attribute is a sequence of 7-item sequences. |
| |
| Each of these sequences contains information describing one result column: |
| |
| - name |
| - type_code |
| - display_size (None in current implementation) |
| - internal_size (None in current implementation) |
| - precision (None in current implementation) |
| - scale (None in current implementation) |
| - null_ok (always True in current implementation) |
| |
| The ``type_code`` can be interpreted by comparing it to the Type Objects specified in the |
| section below. |
| """ |
| # Sleep until we're done or we got the columns |
| self._fetch_while( |
| lambda: self._columns is None and |
| self._state not in (self._STATE_NONE, self._STATE_FINISHED) |
| ) |
| if self._columns is None: |
| return None |
| return [ |
| # name, type_code, display_size, internal_size, precision, scale, null_ok |
| (col['name'], col['type'], None, None, None, None, True) |
| for col in self._columns |
| ] |
| |
| def execute(self, operation, parameters=None): |
| """Prepare and execute a database operation (query or command). |
| |
| Return values are not defined. |
| """ |
| headers = { |
| 'X-Presto-Catalog': self._catalog, |
| 'X-Presto-Schema': self._schema, |
| 'X-Presto-Source': self._source, |
| 'X-Presto-User': self._username, |
| } |
| |
| if self._session_props: |
| headers['X-Presto-Session'] = ','.join( |
| '{}={}'.format(propname, propval) |
| for propname, propval in self._session_props.items() |
| ) |
| |
| # Prepare statement |
| if parameters is None: |
| sql = operation |
| else: |
| sql = operation % _escaper.escape_args(parameters) |
| |
| self._reset_state() |
| |
| self._state = self._STATE_RUNNING |
| url = urlparse.urlunparse(( |
| self._protocol, |
| '{}:{}'.format(self._host, self._port), '/v1/statement', None, None, None)) |
| _logger.info('%s', sql) |
| _logger.debug("Headers: %s", headers) |
| response = self._requests_session.post( |
| url, data=sql.encode('utf-8'), headers=headers, **self._requests_kwargs) |
| self._process_response(response) |
| |
| def cancel(self): |
| if self._state == self._STATE_NONE: |
| raise ProgrammingError("No query yet") |
| if self._nextUri is None: |
| assert self._state == self._STATE_FINISHED, "Should be finished if nextUri is None" |
| return |
| |
| response = self._requests_session.delete(self._nextUri, **self._requests_kwargs) |
| if response.status_code != requests.codes.no_content: |
| fmt = "Unexpected status code after cancel {}\n{}" |
| raise OperationalError(fmt.format(response.status_code, response.content)) |
| |
| self._state = self._STATE_FINISHED |
| self._nextUri = None |
| |
| def poll(self): |
| """Poll for and return the raw status data provided by the Presto REST API. |
| |
| :returns: dict -- JSON status information or ``None`` if the query is done |
| :raises: ``ProgrammingError`` when no query has been started |
| |
| .. note:: |
| This is not a part of DB-API. |
| """ |
| if self._state == self._STATE_NONE: |
| raise ProgrammingError("No query yet") |
| if self._nextUri is None: |
| assert self._state == self._STATE_FINISHED, "Should be finished if nextUri is None" |
| return None |
| response = self._requests_session.get(self._nextUri, **self._requests_kwargs) |
| self._process_response(response) |
| return response.json() |
| |
| def _fetch_more(self): |
| """Fetch the next URI and update state""" |
| self._process_response(self._requests_session.get(self._nextUri, **self._requests_kwargs)) |
| |
| def _process_data(self, rows): |
| for i, col in enumerate(self.description): |
| col_type = col[1].split("(")[0].lower() |
| if col_type in TYPES_CONVERTER: |
| for row in rows: |
| if row[i] is not None: |
| row[i] = TYPES_CONVERTER[col_type](row[i]) |
| |
| def _process_response(self, response): |
| """Given the JSON response from Presto's REST API, update the internal state with the next |
| URI and any data from the response |
| """ |
| # TODO handle HTTP 503 |
| if response.status_code != requests.codes.ok: |
| fmt = "Unexpected status code {}\n{}" |
| raise OperationalError(fmt.format(response.status_code, response.content)) |
| |
| response_json = response.json() |
| _logger.debug("Got response %s", response_json) |
| assert self._state == self._STATE_RUNNING, "Should be running if processing response" |
| self._nextUri = response_json.get('nextUri') |
| self._columns = response_json.get('columns') |
| if 'id' in response_json: |
| self.last_query_id = response_json['id'] |
| if 'X-Presto-Clear-Session' in response.headers: |
| propname = response.headers['X-Presto-Clear-Session'] |
| self._session_props.pop(propname, None) |
| if 'X-Presto-Set-Session' in response.headers: |
| propname, propval = response.headers['X-Presto-Set-Session'].split('=', 1) |
| self._session_props[propname] = propval |
| if 'data' in response_json: |
| assert self._columns |
| new_data = response_json['data'] |
| self._process_data(new_data) |
| self._data += map(tuple, new_data) |
| if 'nextUri' not in response_json: |
| self._state = self._STATE_FINISHED |
| if 'error' in response_json: |
| raise DatabaseError(response_json['error']) |
| |
| |
| # |
| # Type Objects and Constructors |
| # |
| |
| |
| # See types in presto-main/src/main/java/com/facebook/presto/tuple/TupleInfo.java |
| FIXED_INT_64 = DBAPITypeObject(['bigint']) |
| VARIABLE_BINARY = DBAPITypeObject(['varchar']) |
| DOUBLE = DBAPITypeObject(['double']) |
| BOOLEAN = DBAPITypeObject(['boolean']) |
