diff --git a/astronomer/providers/databricks/hooks/databricks.py b/astronomer/providers/databricks/hooks/databricks.py index 255b1a56a..a105170ee 100644 --- a/astronomer/providers/databricks/hooks/databricks.py +++ b/astronomer/providers/databricks/hooks/databricks.py @@ -162,11 +162,14 @@ def _retryable_error_async(exception: ClientConnectorError | ClientResponseError - requests_exceptions.ConnectionError - requests_exceptions.Timeout - anything with a status code >= 500 + - status code == 403 Most retryable errors are covered by status code >= 500. :return: if the status is retryable :rtype: bool """ if isinstance(exception, ClientResponseError): - return exception.status >= 500 + status_code = exception.status + # according to user feedback, 403 sometimes works after retry + return status_code >= 500 or status_code == 403 return True diff --git a/tests/databricks/hooks/test_databricks.py b/tests/databricks/hooks/test_databricks.py index d68a7e0a0..1a99659fd 100644 --- a/tests/databricks/hooks/test_databricks.py +++ b/tests/databricks/hooks/test_databricks.py @@ -179,8 +179,9 @@ async def test_do_api_call_async_non_retryable_error(self, aioresponse): with pytest.raises(AirflowException): await hook._do_api_call_async(GET_RUN_ENDPOINT, params) + @pytest.mark.parametrize("status_code", (500, 503, 403)) @pytest.mark.asyncio - async def test_do_api_call_async_retryable_error(self, aioresponse): + async def test_do_api_call_async_retryable_error(self, aioresponse, status_code): """ Asserts that the Databricks hook will attempt another API call as many times as the retry_limit when a retryable error is returned by the API. @@ -194,7 +195,7 @@ async def test_do_api_call_async_retryable_error(self, aioresponse): aioresponse.get( f"https://localhost/api/{api_version}/jobs/runs/get?run_id=unit_test_run_id", - status=500, + status=status_code, repeat=True, )