From 43d40c714f5c6be1ed9f84938b62c42553ed63c7 Mon Sep 17 00:00:00 2001 From: Guilherme de Amorim Date: Tue, 8 Jul 2025 17:27:55 -0300 Subject: [PATCH 1/2] :necktie: allow retries for Bad Gateway --- cubejs/client.py | 5 ++++- cubejs/errors.py | 24 ++++++++++++++---------- tests/test_client.py | 4 ++++ 3 files changed, 22 insertions(+), 11 deletions(-) diff --git a/cubejs/client.py b/cubejs/client.py index 16dadc1..bed74f5 100644 --- a/cubejs/client.py +++ b/cubejs/client.py @@ -8,6 +8,7 @@ AuthorizationError, ContinueWaitError, RequestError, + RetryableError, ServerError, UnexpectedResponseError, ) @@ -32,6 +33,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 RetryableError() if response.status_code == 500: raise ServerError(response.text) if response.status_code != 200: @@ -39,7 +42,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..48b8e80 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,17 @@ 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..." diff --git a/tests/test_client.py b/tests/test_client.py index 7d7c1ef..fb9dabe 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.RetryableError) as retryable_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 "we can attempt a retry" in str(retryable_error.value) From e49db1aeb3fdf0cba22ee03fc1fbbd8f2ad11465 Mon Sep 17 00:00:00 2001 From: Guilherme de Amorim Date: Tue, 8 Jul 2025 17:46:58 -0300 Subject: [PATCH 2/2] :recycle: create new BadGatewayError --- cubejs/client.py | 3 ++- cubejs/errors.py | 7 +++++++ tests/test_client.py | 4 ++-- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/cubejs/client.py b/cubejs/client.py index bed74f5..4a90de4 100644 --- a/cubejs/client.py +++ b/cubejs/client.py @@ -6,6 +6,7 @@ from cubejs.errors import ( AuthorizationError, + BadGatewayError, ContinueWaitError, RequestError, RetryableError, @@ -34,7 +35,7 @@ def _error_handler(response: httpx.Response) -> None: if "Continue wait" in response.text: raise ContinueWaitError() if response.status_code == 502: - raise RetryableError() + raise BadGatewayError() if response.status_code == 500: raise ServerError(response.text) if response.status_code != 200: diff --git a/cubejs/errors.py b/cubejs/errors.py index 48b8e80..ddef535 100644 --- a/cubejs/errors.py +++ b/cubejs/errors.py @@ -53,3 +53,10 @@ class ContinueWaitError(RetryableError): 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 fb9dabe..eb56f42 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -99,7 +99,7 @@ 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.RetryableError) as retryable_error: + 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: @@ -117,4 +117,4 @@ def test_error_handler(): ) assert str(server_error.value) == "CubeJS server error: " assert str(unexpected_response_error.value) == "CubeJS unexpected response: " - assert "we can attempt a retry" in str(retryable_error.value) + assert "attempting a retry" in str(bad_gateway_error.value)