Skip to content

Commit

Permalink
fix(databricks): fix _retryable_error_async type
Browse files Browse the repository at this point in the history
  • Loading branch information
Lee-W committed Oct 23, 2023
1 parent c3fd458 commit d83e751
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions astronomer/providers/databricks/hooks/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from typing import Any, Dict, Tuple, cast

import aiohttp
from aiohttp import ClientConnectorError, ClientError, ClientResponseError
from aiohttp import ClientConnectorError, ClientResponseError
from airflow import __version__
from airflow.exceptions import AirflowException
from airflow.providers.databricks.hooks.databricks import (
Expand Down Expand Up @@ -150,7 +150,7 @@ async def _do_api_call_async(
attempt_num += 1
await asyncio.sleep(self.retry_delay)

def _retryable_error_async(self, exception: ClientError) -> bool:
def _retryable_error_async(self, exception: ClientConnectorError | ClientResponseError) -> bool:
"""
Determines whether or not an exception that was thrown might be successful
on a subsequent attempt.
Expand All @@ -164,4 +164,6 @@ def _retryable_error_async(self, exception: ClientError) -> bool:
:return: if the status is retryable
:rtype: bool
"""
return exception.status >= 500
if isinstance(exception, ClientResponseError):
return exception.status >= 500
return True

0 comments on commit d83e751

Please sign in to comment.