prestogres 0.1.0

data/pgsql/presto_client.py

@@ -0,0 +1,346 @@
+import os
+import httplib
+import time
+
+try: import simplejson as json
+except ImportError: import json
+
+VERSION = "0.1.0"
+
+class ClientSession(object):
+    def __init__(self, server, user, source=None, catalog=None, schema=None, debug=False):
+        self.server = server
+        self.user = user
+        self.source = source
+        self.catalog = catalog
+        self.schema = schema
+        self.debug = debug
+
+class StatementStats(object):
+    def __init__(self, state=None, scheduled=None, nodes=None, total_splits=None, queued_splits=None, running_splits=None, completed_splits=None, user_time_millis=None, cpu_time_millis=None, wall_time_millis=None, processed_rows=None, processed_bytes=None):
+        self.state = state
+        self.scheduled = scheduled
+        self.nodes = nodes
+        self.total_splits = total_splits
+        self.queued_splits = queued_splits
+        self.running_splits = running_splits
+        self.completed_splits = completed_splits
+        self.user_time_millis = user_time_millis
+        self.cpu_time_millis = cpu_time_millis
+        self.wall_time_millis = wall_time_millis
+        self.processed_rows = processed_rows
+        self.processed_bytes = processed_bytes
+        #self.root_stage = root_stage
+
+    @classmethod
+    def decode_dict(cls, dic):
+        return StatementStats(
+                state=dic.get("state"),
+                scheduled=dic.get("scheduled"),
+                nodes=dic.get("nodes"),
+                total_splits=dic.get("totalSplits"),
+                queued_splits=dic.get("queuedSplits"),
+                running_splits=dic.get("runningSplits"),
+                completed_splits=dic.get("completedSplits"),
+                user_time_millis=dic.get("userTimeMillis"),
+                cpu_time_millis=dic.get("cpuTimeMillis"),
+                wall_time_millis=dic.get("wallTimeMillis"),
+                processed_rows=dic.get("processedRows"),
+                processed_bytes=dic.get("processedBytes"),
+                #root_stage=StageStats.decode_dict(dic["rootStage")),
+                )
+
+class Column(object):
+    def __init__(self, name, type):
+        self.name = name
+        self.type = type
+
+    @classmethod
+    def decode_dict(cls, dic):
+        return Column(
+                name=dic.get("name"),
+                type=dic.get("type"),
+                )
+
+class ErrorLocation(object):
+    def __init__(self, line_number, column_number):
+        self.line_number = line_number
+        self.column_number = column_number
+
+    @classmethod
+    def decode_dict(cls, dic):
+        return ErrorLocation(
+                line_number=dic.get("lineNumber"),
+                column_number=dic.get("columnNumber"),
+                )
+
+class FailureInfo(object):
+    def __init__(self, type=None, message=None, cause=None, suppressed=None, stack=None, error_location=None):
+        self.type = type
+        self.message = message
+        self.cause = cause
+        self.suppressed = suppressed
+        self.stack = stack
+        self.error_location = error_location
+
+    @classmethod
+    def decode_dict(cls, dic):
+        return FailureInfo(
+                type=dic.get("type"),
+                message=dic.get("message"),
+                cause=dic.get("cause"),
+                suppressed=map(FailureInfo.decode_dict, dic["suppressed"]) if "suppressed" in dic else None,
+                stack=dic.get("stack"),
+                error_location=ErrorLocation.decode_dict(dic["errorLocation"]) if "errorLocation" in dic else None,
+                )
+
+class QueryError(object):
+    def __init__(self, message=None, sql_state=None, error_code=None, error_location=None, failure_info=None):
+        self.message = message
+        self.sql_state = sql_state
+        self.error_code = error_code
+        self.error_location = error_location
+        self.failure_info = failure_info
+
+    @classmethod
+    def decode_dict(cls, dic):
+        return QueryError(
+                message=dic.get("message"),
+                sql_state=dic.get("sqlState"),
+                error_code=dic.get("errorCode"),
+                error_location=ErrorLocation.decode_dict(dic["errorLocation"]) if "errorLocation" in dic else None,
+                failure_info=FailureInfo.decode_dict(dic["failureInfo"]) if "failureInfo" in dic else None,
+                )
+
+class QueryResults(object):
+    def __init__(self, id, info_uri=None, partial_cache_uri=None, next_uri=None, columns=None, data=None, stats=None, error=None):
+        self.id = id
+        self.info_uri = info_uri
+        self.partial_cache_uri = partial_cache_uri
+        self.next_uri = next_uri
+        self.columns = columns
+        self.data = data
+        self.stats = stats
+        self.error = error
+
+    @classmethod
+    def decode_dict(cls, dic):
+        return QueryResults(
+                id=dic.get("id"),
+                info_uri=dic.get("infoUri"),
+                partial_cache_uri=dic.get("partialCancelUri"),
+                next_uri=dic.get("nextUri"),
+                columns=map(Column.decode_dict, dic["columns"]) if "columns" in dic else None,
+                data=dic.get("data"),
+                stats=StatementStats.decode_dict(dic["stats"]) if "stats" in dic else None,
+                error=QueryError.decode_dict(dic["error"]) if "error" in dic else None,
+                )
+
+class PrestoException(Exception):
+    pass
+
+class PrestoHttpException(PrestoException):
+    def __init__(self, status, message):
+        PrestoException.__init__(self, message)
+        self.status = status
+
+class PrestoClientException(PrestoException):
+    pass
+
+class PrestoQueryException(PrestoException):
+    def __init__(self, message, query_id, error_code, failure_info):
+        PrestoException.__init__(self, message)
+        self.query_id = query_id
+        self.error_code = error_code
+        self.failure_info = failure_info
+
+class PrestoHeaders(object):
+    PRESTO_USER = "X-Presto-User"
+    PRESTO_SOURCE = "X-Presto-Source"
+    PRESTO_CATALOG = "X-Presto-Catalog"
+    PRESTO_SCHEMA = "X-Presto-Schema"
+
+    PRESTO_CURRENT_STATE = "X-Presto-Current-State"
+    PRESTO_MAX_WAIT = "X-Presto-Max-Wait"
+    PRESTO_MAX_SIZE = "X-Presto-Max-Size"
+    PRESTO_PAGE_SEQUENCE_ID = "X-Presto-Page-Sequence-Id"
+
+class StatementClient(object):
+    HEADERS = {
+            "User-Agent": "presto-python/%s" % VERSION
+            }
+
+    def __init__(self, http_client, session, query):
+        self.http_client = http_client
+        self.session = session
+        self.query = query
+
+        self.closed = False
+        self.exception = None
+        self.results = None
+        self._post_query_request()
+
+    def _post_query_request(self):
+        headers = StatementClient.HEADERS.copy()
+
+        if self.session.user is not None:
+            headers[PrestoHeaders.PRESTO_USER] = self.session.user
+        if self.session.source is not None:
+            headers[PrestoHeaders.PRESTO_SOURCE] = self.session.source
+        if self.session.catalog is not None:
+            headers[PrestoHeaders.PRESTO_CATALOG] = self.session.catalog
+        if self.session.schema is not None:
+            headers[PrestoHeaders.PRESTO_SCHEMA] = self.session.schema
+
+        self.http_client.request("POST", "/v1/statement", self.query, headers)
+        response = self.http_client.getresponse()
+        body = response.read()
+
+        if response.status != 200:
+            raise PrestoHttpException(response.status, "Failed to start query: %s" % body)
+
+        dic = json.loads(body)
+        self.results = QueryResults.decode_dict(dic)
+
+    @property
+    def is_query_failed(self):
+        return self.results.error is not None
+
+    @property
+    def is_query_succeeded(self):
+        return self.results.error is None and self.exception is None and not self.closed
+
+    @property
+    def has_next(self):
+        return self.results.next_uri is not None
+
+    def advance(self):
+        if self.closed or not self.has_next:
+            return False
+
+        uri = self.results.next_uri
+        start = time.time()
+        attempts = 0
+
+        while True:
+            try:
+                self.http_client.request("GET", uri)
+            except Exception as e:
+                self.exception = e
+                raise
+
+            response = self.http_client.getresponse()
+            body = response.read()
+
+            if response.status == 200 and body:
+                self.results = QueryResults.decode_dict(json.loads(body))
+                return True
+
+            if response.status != 503:  # retry on 503 Service Unavailable
+                # deterministic error
+                self.exception = PrestoHttpException(response.status, "Error fetching next at %s returned %s: %s" % (uri, response.status, body))  # TODO error class
+                raise self.exception
+
+            if (time.time() - start) > 2*60*60 or self.closed:
+                break
+
+        self.exception = PrestoHttpException(408, "Error fetching next")  # TODO error class
+        raise self.exception
+
+    def cancel_leaf_stage(self):
+        if self.results.next_uri is not None:
+            self.http_client.request("DELETE", self.results.next_uri)
+            response = self.http_client.getresponse()
+            response.read()
+            return response.status / 100 == 2
+        return False
+
+    def close(self):
+        if self.closed:
+            return
+
+        cancel_leaf_stage(self)
+
+        self.closed = True
+
+class Query(object):
+    @classmethod
+    def start(cls, session, query):
+        http_client = httplib.HTTPConnection(session.server)
+        return Query(StatementClient(http_client, session, query))
+
+    def __init__(self, client):
+        self.client = client
+
+    def _wait_for_data(self):
+        while self.client.has_next and self.client.results.data is None:
+            self.client.advance()
+
+    def columns(self):
+        self._wait_for_data()
+
+        if not self.client.is_query_succeeded:
+            self._raise_error()
+
+        return self.client.results.columns
+
+    def results(self):
+        client = self.client
+
+        if not client.is_query_succeeded:
+            self._raise_error()
+
+        if self.columns() is None:
+            raise PrestoException("Query %s has no columns" % client.results.id)
+
+        while True:
+            if client.results.data is None:
+                break
+
+            for row in client.results.data:
+                yield row
+
+            if not self.client.has_next:
+                break
+
+            client.advance()
+            if client.results.data is None:
+                break
+
+    def cancel(self):
+        self.client.cancel_leaf_stage()
+
+    def close(self):
+        self.client.cancel_leaf_stage()
+
+    def _raise_error(self):
+        if self.client.closed:
+            raise PrestoClientException("Query aborted by user")
+        elif self.client.exception is not None:
+            raise self.client.exception
+        elif self.client.is_query_failed:
+            results = self.client.results
+            error = results.error
+            if error is None:
+                raise PrestoQueryException("Query %s failed: (unknown reason)" % results.id, None, None)
+            raise PrestoQueryException("Query %s failed: %s" % (results.id, error.message), results.id, error.error_code, error.failure_info)
+
+class Client(object):
+    def __init__(self, **options):
+        self.session = ClientSession(**options)
+
+    def query(self, query):
+        return Query.start(self.session, query)
+
+    def run(self, query):
+        q = Query.start(self.session, query)
+        try:
+            columns = q.columns()
+            if columns is None:
+                return [], []
+            rows = []
+            map(rows.append, q.results())
+            return columns, rows
+        finally:
+            q.close()
+

data/pgsql/prestogres.py

@@ -0,0 +1,156 @@
+import presto_client
+import plpy
+from collections import namedtuple
+import time
+
+def run_presto_as_temp_table(server, user, catalog, schema, table_name, query):
+    client = presto_client.Client(server=server, user=user, catalog=catalog, schema=schema)
+
+    create_sql = 'create temp table ' + plpy.quote_ident(table_name) + ' (\n  '
+    insert_sql = 'insert into ' + plpy.quote_ident(table_name) + ' (\n  '
+    values_types = []
+
+    q = client.query(query)
+    try:
+        columns = q.columns()
+        column_num = len(columns)
+
+        first = True
+        for column in columns:
+            if column.type == "varchar":
+                pg_column_type = "text"
+            elif column.type == "bigint":
+                pg_column_type = "bigint"
+            elif column.type == "boolean":
+                pg_column_type = "boolean"
+            elif column.type == "double":
+                pg_column_type = "double precision"
+            else:
+                raise Exception, "unknown column type: " + plpy.quote_ident(column.type)
+
+            if first:
+                first = False
+            else:
+                create_sql += ",\n  "
+                insert_sql += ",\n  "
+
+            create_sql += plpy.quote_ident(column.name) + ' ' + pg_column_type
+            insert_sql += plpy.quote_ident(column.name)
+            values_types.append(pg_column_type)
+
+        create_sql += '\n)'
+        #if trait:
+        #    create_sql += ' '
+        #    create_sql += trait
+        create_sql += ';'
+
+        insert_sql += '\n) values\n'
+        values_sql_format = '(' + ', '.join(['${}'] * column_num) + ')'
+
+        #plpy.execute("drop table if exists "+plpy.quote_ident(table_name))
+        plpy.execute(create_sql)
+
+        batch = []
+        for row in q.results():
+            batch.append(row)
+            if len(batch) > 10:
+                batch_len = len(batch)
+                # format string 'values ($1, $2), ($3, $4) ...'
+                values_sql = (', '.join([values_sql_format] * batch_len)).format(*range(1, batch_len * column_num + 1))
+                batch_insert_sql = insert_sql + values_sql
+                # flatten rows into an array
+                params = [item for sublist in batch for item in sublist]
+                plpy.execute(plpy.prepare(batch_insert_sql, values_types * batch_len), params)
+                del batch[:]
+
+        if batch:
+            batch_len = len(batch)
+            # format string 'values ($1, $2), ($3, $4) ...'
+            values_sql = (', '.join([values_sql_format] * batch_len)).format(*range(1, batch_len * column_num + 1))
+            batch_insert_sql = insert_sql + values_sql
+            # flatten rows into an array
+            params = [item for sublist in batch for item in sublist]
+            plpy.execute(plpy.prepare(batch_insert_sql, values_types * batch_len), params)
+
+    finally:
+        q.close()
+
+Column = namedtuple('Column', ('name', 'type', 'nullable'))
+
+cache_expire_times = {}
+
+def presto_create_tables(server, user, catalog):
+    client = presto_client.Client(server=server, user=user, catalog=catalog, schema="default")
+
+    cache_key = "%s:%s.%s" % (server, user, catalog)
+    expire_time = cache_expire_times.get(cache_key)
+    if expire_time is not None and time.time() - expire_time < 10:
+        # TODO scan cache_expire_times and remove expired cache entries if it is large
+        return
+
+    try:
+        schemas = {}
+
+        columns, rows = client.run("select table_schema, table_name, column_name, is_nullable, data_type from information_schema.columns")
+
+        if rows is None:
+            return
+
+        for row in rows:
+            schema_name = row[0]
+            table_name = row[1]
+            column_name = row[2]
+            is_nullable = row[3]
+            column_type = row[4]
+
+            tables = schemas.setdefault(schema_name, {})
+            columns = tables.setdefault(table_name, [])
+            columns.append(Column(column_name, column_type, is_nullable))
+
+        for schema_name, tables in schemas.items():
+            if schema_name == "sys" or schema_name == "information_schema":
+                # skip system schemas
+                continue
+
+            # create schema
+            try:
+                plpy.execute("create schema %s" % plpy.quote_ident(schema_name))
+            except:
+                # ignore error
+                pass
+
+            for table_name, columns in tables.items():
+                create_sql = "create table %s.%s (\n  " % (plpy.quote_ident(schema_name), plpy.quote_ident(table_name))
+
+                first = True
+                for column in columns:
+                    if column.type == "varchar":
+                        pg_column_type = "varchar(100)"
+                    elif column.type == "bigint":
+                        pg_column_type = "bigint"
+                    elif column.type == "boolean":
+                        pg_column_type = "boolean"
+                    elif column.type == "double":
+                        pg_column_type = "double precision"
+                    else:
+                        raise Exception("unknown column type: " + plpy.quote_ident(column.type))
+
+                    if first:
+                        first = False
+                    else:
+                        create_sql += ",\n  "
+
+                    create_sql += plpy.quote_ident(column.name) + " " + pg_column_type
+                    if not column.nullable:
+                        create_sql += " not null"
+
+                create_sql += "\n)"
+
+                plpy.execute("drop table if exists %s.%s" % (plpy.quote_ident(schema_name), plpy.quote_ident(table_name)))
+                plpy.execute(create_sql)
+
+        cache_expire_times[cache_key] = time.time()
+
+    except Exception as e:
+        plpy.error(str(e))
+