diff --git a/trino/client.py b/trino/client.py index 1697de65..a15a7c91 100644 --- a/trino/client.py +++ b/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 diff --git a/trino/exceptions.py b/trino/exceptions.py index d48fc9ef..cca7cf3d 100644 --- a/trino/exceptions.py +++ b/trino/exceptions.py @@ -67,6 +67,10 @@ class TrinoAuthError(OperationalError): pass +class TrinoConnectionError(OperationalError): + pass + + class TrinoDataError(NotSupportedError): pass