Add query_id in PrestoQueryError

prestodb · Feb 23, 2018 · b352912 · b352912
1 parent 4bccbda
commit b352912
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 8 deletions.
diff --git a/prestodb/client.py b/prestodb/client.py
@@ -373,14 +373,14 @@ def delete(self, url):
             proxies=PROXIES,
         )
 
-    def _process_error(self, error):
+    def _process_error(self, error, query_id):
         error_type = error['errorType']
         if error_type == 'EXTERNAL':
-            raise exceptions.PrestoExternalError(error)
+            raise exceptions.PrestoExternalError(error, query_id)
         elif error_type == 'USER_ERROR':
-            return exceptions.PrestoUserError(error)
+            return exceptions.PrestoUserError(error, query_id)
 
-        return exceptions.PrestoQueryError(error)
+        return exceptions.PrestoQueryError(error, query_id)
 
     def raise_response_error(self, http_response):
         if http_response.status_code == 503:
@@ -402,7 +402,7 @@ def process(self, http_response):
         response = http_response.json()
         logger.debug('HTTP {}: {}'.format(http_response.status_code, response))
         if 'error' in response:
-            raise self._process_error(response['error'])
+            raise self._process_error(response['error'], response.get('id'))
 
         if constants.HEADER_CLEAR_SESSION in http_response.headers:
             for prop in get_header_values(
@@ -505,7 +505,10 @@ def execute(self):
         call fetch() until is_finished is true.
         """
         if self._cancelled:
-            raise exceptions.PrestoUserError("Query has been cancelled")
+            raise exceptions.PrestoUserError(
+                "Query has been cancelled",
+                self.query_id,
+            )
 
         response = self._request.post(self._sql)
         status = self._request.process(response)

diff --git a/prestodb/exceptions.py b/prestodb/exceptions.py
@@ -45,8 +45,9 @@ class DatabaseError(Exception):
 
 
 class PrestoQueryError(Exception):
-    def __init__(self, error):
+    def __init__(self, error, query_id=None):
         self._error = error
+        self._query_id = query_id
 
     @property
     def error_code(self):
@@ -80,12 +81,17 @@ def error_location(self):
         location = self._error['errorLocation']
         return (location['lineNumber'], location['columnNumber'])
 
+    @property
+    def query_id(self):
+        return self._query_id
+
     def __repr__(self):
-        return '{}(type={}, name={}, message="{}")'.format(
+        return '{}(type={}, name={}, message="{}", query_id={})'.format(
             self.__class__.__name__,
             self.error_type,
             self.error_name,
             self.message,
+            self.query_id,
         )
 
     def __str__(self):

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -452,6 +452,7 @@ def test_presto_fetch_error(monkeypatch):
     assert 'suppressed' in error.failure_info
     assert error.message == 'line 1:15: Schema must be specified when session schema is not set'
     assert error.error_location == (1, 15)
+    assert error.query_id == '20161116_205844_00002_xtnym'
 
 
 @pytest.mark.parametrize("error_code, error_type, error_message", [