Merge pull request #2128 from artwr/artwr-improve_presto_hook_error_when_cluster_is_unavailable

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)