👔 allow retries for Bad Gateway

cubejs/client.py

@@ -6,8 +6,10 @@
 
 from cubejs.errors import (
     AuthorizationError,
+    BadGatewayError,
     ContinueWaitError,
     RequestError,
+    RetryableError,
     ServerError,
     UnexpectedResponseError,
 )
@@ -32,14 +34,16 @@ def _error_handler(response: httpx.Response) -> None:
         raise RequestError(response.text)
     if "Continue wait" in response.text:
         raise ContinueWaitError()
+    if response.status_code == 502:
+        raise BadGatewayError()
     if response.status_code == 500:
         raise ServerError(response.text)
     if response.status_code != 200:
         raise UnexpectedResponseError(response.text)
 
 
 @tenacity.retry(
-    retry=tenacity.retry_if_exception_type(ContinueWaitError),
+    retry=tenacity.retry_if_exception_type(RetryableError),
     wait=tenacity.wait_exponential(multiplier=2, min=1, max=30),
     stop=tenacity.stop_after_attempt(5),
 )

cubejs/errors.py

@@ -1,16 +1,6 @@
 """Errors expected from CubeJS server."""
 
 
-class ContinueWaitError(Exception):
-    """Raised when CubeJS responds with 'Continue wait' message."""
-
-    def __init__(self) -> None:
-        pass
-
-    def __str__(self) -> str:
-        return "CubeJS query is not ready yet, continue waiting..."
-
-
 class ServerError(Exception):
     """Raised when CubeJS responds with an error."""
 
@@ -49,3 +39,24 @@ def __init__(self, message: str) -> None:
 
     def __str__(self) -> str:
         return f"CubeJS unexpected response: {self.message}"
+
+
+class RetryableError(Exception):
+    """Raised when a retry can be performed."""
+
+    def __str__(self) -> str:
+        return "CubeJS failed but we can attempt a retry"
+
+
+class ContinueWaitError(RetryableError):
+    """Raised when CubeJS responds with 'Continue wait' message."""
+
+    def __str__(self) -> str:
+        return "CubeJS query is not ready yet, continue waiting..."
+
+
+class BadGatewayError(RetryableError):
+    """Raised when CubeJS responds with 'Bad Gateway'."""
+
+    def __str__(self) -> str:
+        return "CubeJS may be scaling instances, attempting a retry..."

tests/test_client.py

@@ -99,6 +99,9 @@ def test_error_handler():
     with pytest.raises(errors.ContinueWaitError) as continue_wait_error:
         _error_handler(Mock(status_code=200, text="Continue wait"))
 
+    with pytest.raises(errors.BadGatewayError) as bad_gateway_error:
+        _error_handler(Mock(status_code=502, text="Bad Gateway"))
+
     with pytest.raises(errors.ServerError) as server_error:
         _error_handler(Mock(status_code=500, text=""))
 
@@ -114,3 +117,4 @@ def test_error_handler():
     )
     assert str(server_error.value) == "CubeJS server error: "
     assert str(unexpected_response_error.value) == "CubeJS unexpected response: "
+    assert "attempting a retry" in str(bad_gateway_error.value)