Retry on 502 bad gateway responses

Author: d-cole

tests/unit/test_client.py

@@ -753,6 +753,7 @@ def test_authentication_fail_retry(monkeypatch):
 
 
 @pytest.mark.parametrize("status_code, attempts", [
+    (502, 3),
     (503, 3),
     (504, 3),
 ])

trino/client.py

@@ -340,9 +340,9 @@ def max_attempts(self, value) -> None:
             self._handle_retry,
             exceptions=self._exceptions,
             conditions=(
-                # need retry when there is no exception but the status code is 503 or 504
+                # need retry when there is no exception but the status code is 502, 503, or 504
                 lambda response: getattr(response, "status_code", None)
-                in (503, 504),
+                in (502, 503, 504),
             ),
             max_attempts=self._max_attempts,
         )
@@ -416,6 +416,9 @@ def _process_error(self, error, query_id):
         return exceptions.TrinoQueryError(error, query_id)
 
     def raise_response_error(self, http_response):
+        if http_response.status_code == 502:
+            raise exceptions.Http502Error("error 502: bad gateway")
+
         if http_response.status_code == 503:
             raise exceptions.Http503Error("error 503: service unavailable")
 

trino/exceptions.py

@@ -29,6 +29,10 @@ class HttpError(Exception):
     pass
 
 
+class Http502Error(Exception):
+    pass
+
+
 class Http503Error(HttpError):
     pass