kyuubi/python/pyhive/presto.py

"""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'])