diff --git a/cubejs/client.py b/cubejs/client.py index 16dadc1..4a90de4 100644 --- a/cubejs/client.py +++ b/cubejs/client.py @@ -6,8 +6,10 @@ from cubejs.errors import ( AuthorizationError, + BadGatewayError, ContinueWaitError, RequestError, + RetryableError, ServerError, UnexpectedResponseError, ) @@ -32,6 +34,8 @@ 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: @@ -39,7 +43,7 @@ def _error_handler(response: httpx.Response) -> None: @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), ) diff --git a/cubejs/errors.py b/cubejs/errors.py index b817fcd..ddef535 100644 --- a/cubejs/errors.py +++ b/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..." diff --git a/tests/test_client.py b/tests/test_client.py index 7d7c1ef..eb56f42 100644 --- a/tests/test_client.py +++ b/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)