This is an automated email from the ASF dual-hosted git repository.
pankajkoti pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0de5587894 Add TimeoutError to be a retryable error in databricks
provider (#43137)
0de5587894 is described below
commit 0de5587894286a1a9455084943dbee57891ea016
Author: Kalyan R <[email protected]>
AuthorDate: Fri Oct 18 12:50:04 2024 +0530
Add TimeoutError to be a retryable error in databricks provider (#43137)
closes #43128
---
.../src/airflow/providers/databricks/hooks/databricks_base.py | 3 ++-
providers/tests/databricks/hooks/test_databricks.py | 11 +++++++++++
2 files changed, 13 insertions(+), 1 deletion(-)
diff --git
a/providers/src/airflow/providers/databricks/hooks/databricks_base.py
b/providers/src/airflow/providers/databricks/hooks/databricks_base.py
index 08a6eb8d40..1a4ccb6e98 100644
--- a/providers/src/airflow/providers/databricks/hooks/databricks_base.py
+++ b/providers/src/airflow/providers/databricks/hooks/databricks_base.py
@@ -28,6 +28,7 @@ from __future__ import annotations
import copy
import platform
import time
+from asyncio.exceptions import TimeoutError
from functools import cached_property
from typing import TYPE_CHECKING, Any
from urllib.parse import urlsplit
@@ -679,7 +680,7 @@ class BaseDatabricksHook(BaseHook):
if exception.status >= 500 or exception.status == 429:
return True
- if isinstance(exception, ClientConnectorError):
+ if isinstance(exception, (ClientConnectorError, TimeoutError)):
return True
return False
diff --git a/providers/tests/databricks/hooks/test_databricks.py
b/providers/tests/databricks/hooks/test_databricks.py
index bec238f70e..94a8cfb9c4 100644
--- a/providers/tests/databricks/hooks/test_databricks.py
+++ b/providers/tests/databricks/hooks/test_databricks.py
@@ -21,6 +21,7 @@ import itertools
import json
import ssl
import time
+from asyncio.exceptions import TimeoutError
from unittest import mock
from unittest.mock import AsyncMock
@@ -1551,6 +1552,16 @@ class TestDatabricksHookAsyncMethods:
await self.hook._a_do_api_call(GET_RUN_ENDPOINT, {})
assert mock_errors.call_count == DEFAULT_RETRY_NUMBER
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
+ async def test_do_api_call_retries_with_client_timeout_error(self,
mock_get):
+ mock_get.side_effect = TimeoutError()
+ with mock.patch.object(self.hook.log, "error") as mock_errors:
+ async with self.hook:
+ with pytest.raises(AirflowException):
+ await self.hook._a_do_api_call(GET_RUN_ENDPOINT, {})
+ assert mock_errors.call_count == DEFAULT_RETRY_NUMBER
+
@pytest.mark.asyncio
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
async def test_do_api_call_retries_with_retryable_error(self, mock_get):
