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

@@ -787,7 +787,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
@@ -818,7 +821,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)
@@ -836,7 +842,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

@@ -67,6 +67,10 @@ class TrinoAuthError(OperationalError):
     pass
 
 
+class TrinoConnectionError(OperationalError):
+    pass
+
+
 class TrinoDataError(NotSupportedError):
     pass