diff --git a/airflow/hooks/presto_hook.py b/airflow/hooks/presto_hook.py index f8f6ac887f..768ff3f836 100644 --- a/airflow/hooks/presto_hook.py +++ b/airflow/hooks/presto_hook.py @@ -54,6 +54,19 @@ class PrestoHook(DbApiHook): def _strip_sql(sql): return sql.strip().rstrip(';') + def _get_pretty_exception_message(self, e): + """ + Parses some DatabaseError to provide a better error message + """ + if (hasattr(e, 'message') + and 'errorName' in e.message + and 'message' in e.message): + return ('{name}: {message}'.format( + name=e.message['errorName'], + message=e.message['message'])) + else: + return str(e) + def get_records(self, hql, parameters=None): """ Get a set of records from Presto @@ -62,14 +75,7 @@ class PrestoHook(DbApiHook): return super(PrestoHook, self).get_records( self._strip_sql(hql), parameters) except DatabaseError as e: - if (hasattr(e, 'message') and - 'errorName' in e.message and - 'message' in e.message): - # Use the structured error data in the raised exception - raise PrestoException('{name}: {message}'.format( - name=e.message['errorName'], message=e.message['message'])) - else: - raise PrestoException(str(e)) + raise PrestoException(self._parse_exception_message(e)) def get_first(self, hql, parameters=None): """ @@ -80,7 +86,7 @@ class PrestoHook(DbApiHook): return super(PrestoHook, self).get_first( self._strip_sql(hql), parameters) except DatabaseError as e: - raise PrestoException(e[0]['message']) + raise PrestoException(self._parse_exception_message(e)) def get_pandas_df(self, hql, parameters=None): """ @@ -92,7 +98,7 @@ class PrestoHook(DbApiHook): cursor.execute(self._strip_sql(hql), parameters) data = cursor.fetchall() except DatabaseError as e: - raise PrestoException(e[0]['message']) + raise PrestoException(self._parse_exception_message(e)) column_descriptions = cursor.description if data: df = pandas.DataFrame(data)