From e0cf1f2e1d38fe4898518255fbbe54954106ca6f Mon Sep 17 00:00:00 2001
From: shaheer <hshaheer99@gmail.com>
Date: Tue, 18 Jul 2023 23:48:04 +0530
Subject: [PATCH] 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.
---
 trino/client.py     | 15 ++++++++++++---
 trino/exceptions.py |  4 ++++
 2 files changed, 16 insertions(+), 3 deletions(-)

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