Wrap all requests.exceptions.RequestException as TrinoConnectionError

Before this change if a user needed to catch connection errors (e.g.
connection refused) they needed to add a dependency on requests and
import the relevant exception from requests.exceptions module.

After this change all requests exception get rethrown as a
TrinoConnectionError instead.

Author: Shaheer-rossoneri14

trino/client.py

@@ -803,7 +803,10 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
         if self.cancelled:
             raise exceptions.TrinoUserError("Query has been cancelled", self.query_id)
 
-        response = self._request.post(self._query, additional_http_headers)
+        try:
+            response = self._request.post(self._query, additional_http_headers)
+        except requests.exceptions.RequestException as e:
+            raise trino.exceptions.TrinoConnectionError("failed to execute: {}".format(e))
         status = self._request.process(response)
         self._info_uri = status.info_uri
         self._query_id = status.id
@@ -834,7 +837,10 @@ def _update_state(self, status):
 
     def fetch(self) -> List[List[Any]]:
         """Continue fetching data for the current query_id"""
-        response = self._request.get(self._request.next_uri)
+        try:
+            response = self._request.get(self._request.next_uri)
+        except requests.exceptions.RequestException as e:
+            raise trino.exceptions.TrinoConnectionError("failed to fetch: {}".format(e))
         status = self._request.process(response)
         self._update_state(status)
         logger.debug(status)
@@ -852,7 +858,10 @@ def cancel(self) -> None:
             return
 
         logger.debug("cancelling query: %s", self.query_id)
-        response = self._request.delete(self._next_uri)
+        try:
+            response = self._request.delete(self._next_uri)
+        except requests.exceptions.RequestException as e:
+            raise trino.exceptions.TrinoConnectionError("failed to cancel query: {}".format(e))
         logger.debug(response)
         if response.status_code == requests.codes.no_content:
             self._cancelled = True

trino/exceptions.py

@@ -149,3 +149,7 @@ class Http503Error(HttpError):
 
 class Http504Error(HttpError):
     pass
+
+
+class TrinoConnectionError(HttpError):
+    pass