Skip to content

Commit

Permalink
feat(databricks): include 403 as retryable error (#1376)
Browse files Browse the repository at this point in the history
according to user feedbacks, 403 sometimes works after retry
  • Loading branch information
Lee-W authored Dec 6, 2023
1 parent 76e58d7 commit 9e9d8a7
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 3 deletions.
5 changes: 4 additions & 1 deletion astronomer/providers/databricks/hooks/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
5 changes: 3 additions & 2 deletions tests/databricks/hooks/test_databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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,
)

Expand Down

0 comments on commit 9e9d8a7

Please sign in to comment.