import json

from redash.utils import JSONEncoder
from redash.query_runner import *

import logging
logger = logging.getLogger(__name__)

from collections import defaultdict

try:
    from pyhive import presto
    enabled = True

except ImportError:
    enabled = False

PRESTO_TYPES_MAPPING = {
    "integer": TYPE_INTEGER,
    "tinyint": TYPE_INTEGER,
    "smallint": TYPE_INTEGER,
    "long": TYPE_INTEGER,
    "bigint": TYPE_INTEGER,
    "float": TYPE_FLOAT,
    "double": TYPE_FLOAT,
    "boolean": TYPE_BOOLEAN,
    "string": TYPE_STRING,
    "varchar": TYPE_STRING,
    "date": TYPE_DATE,
}


class Presto(BaseQueryRunner):
    @classmethod
    def configuration_schema(cls):
        return {
            'type': 'object',
            'properties': {
                'host': {
                    'type': 'string'
                },
                'port': {
                    'type': 'number'
                },
                'schema': {
                    'type': 'string'
                },
                'catalog': {
                    'type': 'string'
                },
                'username': {
                    'type': 'string'
                }
            },
            'required': ['host']
        }

    @classmethod
    def enabled(cls):
        return enabled

    @classmethod
    def type(cls):
        return "presto"

    def __init__(self, configuration):
        super(Presto, self).__init__(configuration)

    def run_query(self, query):
        connection = presto.connect(
                host=self.configuration.get('host', ''),
                port=self.configuration.get('port', 8080),
                username=self.configuration.get('username', 'redash'),
                catalog=self.configuration.get('catalog', 'hive'),
                schema=self.configuration.get('schema', 'default'))

        cursor = connection.cursor()


        try:
            cursor.execute(query)
            column_tuples = [(i[0], PRESTO_TYPES_MAPPING.get(i[1], None)) for i in cursor.description]
            columns = self.fetch_columns(column_tuples)
            rows = [dict(zip(([c['name'] for c in columns]), r)) for i, r in enumerate(cursor.fetchall())]
            data = {'columns': columns, 'rows': rows}
            json_data = json.dumps(data, cls=JSONEncoder)
            error = None
        except Exception, ex:
            json_data = None
            error = ex.message

        return json_data, error

register(Presto)

@@ -1444,16 +1444,899 @@
 ation)%0A%0A
+    def get_schema(self, get_stats=False):%0A        schema = %7B%7D%0A        query = %22%22%22%0A        SELECT table_schema, table_name, column_name%0A        FROM information_schema.columns%0A        WHERE table_schema NOT IN ('pg_catalog', 'information_schema')%0A        %22%22%22%0A%0A        results, error = self.run_query(query)%0A%0A        if error is not None:%0A            raise Exception(%22Failed getting schema.%22)%0A%0A        results = json.loads(results)%0A%0A        for row in results%5B'rows'%5D:%0A            if row%5B'table_schema'%5D != 'public':%0A                table_name = '%7B%7D.%7B%7D'.format(row%5B'table_schema'%5D, row%5B'table_name'%5D)%0A            else:%0A                table_name = row%5B'table_name'%5D%0A%0A            if table_name not in schema:%0A                schema%5Btable_name%5D = %7B'name': table_name, 'columns': %5B%5D%7D%0A%0A            schema%5Btable_name%5D%5B'columns'%5D.append(row%5B'column_name'%5D)%0A%0A        return schema.values()%0A%0A
     def 

add the ability to load the schema for the presto query runner