Presto query executor may cause exception when rendering error messages

[visualage]

When I misconfigured my Presto connector (i.e. left the catalog field empty), when testing the connection, it would cause the Redash server to throw an exception. The exception is complaining about unable to invoke get method on unicode object.

This is the pull request to fix it: https://github.com/getredash/redash/pull/2143

Technical details:

• Redash Version: 2.0
• Browser/OS: Firefox/Mac (does not matter)
• How did you install Redash: bootstrap script on clean Ubuntu LTS 16.04

[ljluestc]

try:
    import prestodb
    from prestodb import constants
    enabled = True
except ImportError:
    enabled = False

from redash.query_runner import BaseQueryRunner, register
from redash.utils import json_dumps


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

    @classmethod
    def enabled(cls):
        return enabled

    @classmethod
    def name(cls):
        return "Presto"

    def __init__(self, configuration):
        super(Presto, self).__init__(configuration)
        self._connection = None

    def connect(self):
        if self._connection is None:
            # Ensure catalog is not empty or None; fallback gracefully if misconfigured
            catalog = self.configuration.get('catalog') or 'default'
            self._connection = prestodb.dbapi.connect(
                host=self.configuration.get('host', 'localhost'),
                port=self.configuration.get('port', 8080),
                user=self.configuration.get('username', 'admin'),
                catalog=catalog,
                schema=self.configuration.get('schema'),
                http_scheme=constants.HTTPS if self.configuration.get('https') else constants.HTTP,
            )
        return self._connection

    def run_query(self, query, user):
        connection = self.connect()
        cursor = connection.cursor()

        try:
            cursor.execute(query)
            results = cursor.fetchall()
            columns = [desc[0] for desc in cursor.description]
            data = [dict(zip(columns, row)) for row in results]

            error = None
            json_data = json_dumps(data)
        except Exception as e:
            error = str(e)
            json_data = None
            if self._connection:
                self._connection.close()
                self._connection = None
        finally:
            cursor.close()

        return json_data, error

    def get_schema(self, get_stats=False):
        # Simplified schema retrieval; handles empty catalog gracefully
        query = """
            SELECT table_schema, table_name, column_name
            FROM information_schema.columns
            WHERE table_schema NOT IN ('information_schema', 'pg_catalog')
        """
        results, error = self.run_query(query, None)

        if error is not None:
            raise Exception("Failed getting schema: {}".format(error))

        schema = {}
        for row in results:
            table_name = "{}.{}".format(row['table_schema'], row['table_name'])
            if table_name not in schema:
                schema[table_name] = {"name": table_name, "columns": []}
            schema[table_name]["columns"].append(row['column_name'])

        return list(schema.values())


if enabled:
    register(Presto)
​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​