This is an automated email from the ASF dual-hosted git repository.
eladkal 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 950fc37bf6c Fix Data Fusion polling after transient 404 responses
(#74237)
950fc37bf6c is described below
commit 950fc37bf6c23267912f203a43e05bd6846c7936
Author: Ulada Zakharava <[email protected]>
AuthorDate: Thu Oct 8 08:41:09 2026 +0200
Fix Data Fusion polling after transient 404 responses (#74237)
---
generated/known_airflow_exceptions.txt | 2 +-
.../providers/google/cloud/hooks/datafusion.py | 7 +++-
.../providers/google/cloud/sensors/datafusion.py | 11 +++----
.../unit/google/cloud/hooks/test_datafusion.py | 38 ++++++++++++++++++++++
.../unit/google/cloud/sensors/test_datafusion.py | 35 ++++++++++++++++----
5 files changed, 78 insertions(+), 15 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index d1200225357..b26359f86f8 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -279,7 +279,7 @@
providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py::4
providers/google/src/airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py::1
providers/google/src/airflow/providers/google/cloud/sensors/dataflow.py::8
providers/google/src/airflow/providers/google/cloud/sensors/dataform.py::3
-providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py::2
+providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py::1
providers/google/src/airflow/providers/google/cloud/sensors/dataplex.py::7
providers/google/src/airflow/providers/google/cloud/sensors/dataproc.py::6
providers/google/src/airflow/providers/google/cloud/sensors/dataproc_metastore.py::2
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/datafusion.py
b/providers/google/src/airflow/providers/google/cloud/hooks/datafusion.py
index 0da08fe4ea0..6472f397907 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/datafusion.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/datafusion.py
@@ -28,7 +28,7 @@ from typing import Any
from urllib.parse import quote, urlencode, urljoin
import google.auth
-from aiohttp import ClientSession
+from aiohttp import ClientResponseError, ClientSession
from gcloud.aio.auth import AioSession, Token
from google.api_core.retry import exponential_sleep_generator
from googleapiclient.discovery import Resource, build
@@ -614,6 +614,11 @@ class DataFusionAsyncHook(GoogleBaseAsyncHook):
try:
pipeline = await session_aio.get(url=url, headers=headers)
break
+ except ClientResponseError as exc:
+ if exc.status == 404:
+ await asyncio.sleep(time_to_wait)
+ else:
+ raise
except ValueError as exc:
if "404" in str(exc):
await asyncio.sleep(time_to_wait)
diff --git
a/providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py
b/providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py
index d4ba2969ca1..e6d3370b395 100644
--- a/providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py
+++ b/providers/google/src/airflow/providers/google/cloud/sensors/datafusion.py
@@ -22,7 +22,9 @@ from __future__ import annotations
from collections.abc import Iterable, Sequence
from typing import TYPE_CHECKING
-from airflow.providers.common.compat.sdk import AirflowException,
AirflowNotFoundException, BaseSensorOperator
+from requests.exceptions import HTTPError
+
+from airflow.providers.common.compat.sdk import AirflowException,
BaseSensorOperator
from airflow.providers.google.cloud.hooks.datafusion import DataFusionHook
from airflow.providers.google.common.hooks.base_google import
PROVIDE_PROJECT_ID
@@ -113,11 +115,8 @@ class
CloudDataFusionPipelineStateSensor(BaseSensorOperator):
namespace=self.namespace,
)
pipeline_status = pipeline_workflow.get("status")
- except AirflowNotFoundException:
- message = "Specified Pipeline ID was not found."
- raise AirflowException(message)
- except AirflowException:
- pass # Because the pipeline may not be visible in system yet
+ except HTTPError:
+ pass # A newly started pipeline run may not be visible in CDAP yet
if pipeline_status is not None:
if self.failure_statuses and pipeline_status in
self.failure_statuses:
message = (
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_datafusion.py
b/providers/google/tests/unit/google/cloud/hooks/test_datafusion.py
index ce4b24bc1b6..eea02e731a5 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_datafusion.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_datafusion.py
@@ -21,6 +21,7 @@ import logging
from unittest import mock
import pytest
+from aiohttp import ClientResponseError
from requests.exceptions import HTTPError
from airflow.providers.google.cloud.hooks.datafusion import
DataFusionAsyncHook, DataFusionHook
@@ -645,6 +646,43 @@ class TestDataFusionHook:
class TestDataFusionHookAsynch:
+ @pytest.mark.asyncio
+ @mock.patch(HOOK_STR.format("asyncio.sleep"), new_callable=mock.AsyncMock)
+ @mock.patch(HOOK_STR.format("AioSession"))
+ @mock.patch(HOOK_STR.format("Token"))
+ async def test_get_link_retries_after_404(self, mock_token,
mock_aio_session, mock_sleep, hook_async):
+ mock_token.return_value.__aenter__.return_value.get =
mock.AsyncMock(return_value="token")
+ response = MockAiohttpClientResponse(payload={"status": "RUNNING"})
+ mock_aio_session.return_value.get = mock.AsyncMock(
+ side_effect=[
+ ClientResponseError(request_info=mock.Mock(), history=(),
status=404),
+ response,
+ ]
+ )
+
+ result = await hook_async._get_link(url=CONSTRUCTED_PIPELINE_URL,
session=session)
+
+ assert result is response
+ assert mock_aio_session.return_value.get.await_count == 2
+ mock_sleep.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ @mock.patch(HOOK_STR.format("asyncio.sleep"), new_callable=mock.AsyncMock)
+ @mock.patch(HOOK_STR.format("AioSession"))
+ @mock.patch(HOOK_STR.format("Token"))
+ async def test_get_link_propagates_non_404_error(
+ self, mock_token, mock_aio_session, mock_sleep, hook_async
+ ):
+ mock_token.return_value.__aenter__.return_value.get =
mock.AsyncMock(return_value="token")
+ error = ClientResponseError(request_info=mock.Mock(), history=(),
status=500)
+ mock_aio_session.return_value.get = mock.AsyncMock(side_effect=error)
+
+ with pytest.raises(ClientResponseError) as ctx:
+ await hook_async._get_link(url=CONSTRUCTED_PIPELINE_URL,
session=session)
+
+ assert ctx.value is error
+ mock_sleep.assert_not_awaited()
+
@pytest.mark.asyncio
@mock.patch(HOOK_STR.format("DataFusionAsyncHook._get_link"))
async def test_async_get_pipeline_should_execute_successfully(self,
mocked_link, hook_async):
diff --git
a/providers/google/tests/unit/google/cloud/sensors/test_datafusion.py
b/providers/google/tests/unit/google/cloud/sensors/test_datafusion.py
index b5f70cb0c71..3aa11598fee 100644
--- a/providers/google/tests/unit/google/cloud/sensors/test_datafusion.py
+++ b/providers/google/tests/unit/google/cloud/sensors/test_datafusion.py
@@ -20,8 +20,9 @@ from __future__ import annotations
from unittest import mock
import pytest
+from requests.exceptions import HTTPError, RequestException
-from airflow.providers.common.compat.sdk import AirflowException,
AirflowNotFoundException
+from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.google.cloud.hooks.datafusion import PipelineStates
from airflow.providers.google.cloud.sensors.datafusion import
CloudDataFusionPipelineStateSensor
@@ -99,9 +100,9 @@ class TestCloudDataFusionPipelineStateSensor:
task.poke(mock.MagicMock())
@mock.patch("airflow.providers.google.cloud.sensors.datafusion.DataFusionHook")
- def test_not_found_exception(self, mock_hook):
+ def test_pipeline_not_visible_yet(self, mock_hook):
mock_hook.return_value.get_instance.return_value = {"apiEndpoint":
INSTANCE_URL}
- mock_hook.return_value.get_pipeline_workflow.side_effect =
AirflowNotFoundException()
+ mock_hook.return_value.get_pipeline_workflow.side_effect = HTTPError()
task = CloudDataFusionPipelineStateSensor(
task_id="test_task_id",
@@ -116,8 +117,28 @@ class TestCloudDataFusionPipelineStateSensor:
impersonation_chain=IMPERSONATION_CHAIN,
)
- with pytest.raises(
- AirflowException,
- match="Specified Pipeline ID was not found.",
- ):
+ assert task.poke(mock.MagicMock()) is False
+
+
@mock.patch("airflow.providers.google.cloud.sensors.datafusion.DataFusionHook")
+ def test_other_request_error_is_not_suppressed(self, mock_hook):
+ mock_hook.return_value.get_instance.return_value = {"apiEndpoint":
INSTANCE_URL}
+ error = RequestException("Retrieving a pipeline state failed with code
500")
+ mock_hook.return_value.get_pipeline_workflow.side_effect = error
+
+ task = CloudDataFusionPipelineStateSensor(
+ task_id="test_task_id",
+ pipeline_name=PIPELINE_NAME,
+ pipeline_id=PIPELINE_ID,
+ project_id=PROJECT_ID,
+ expected_statuses={PipelineStates.COMPLETED},
+ failure_statuses=FAILURE_STATUSES,
+ instance_name=INSTANCE_NAME,
+ location=LOCATION,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(RequestException) as ctx:
task.poke(mock.MagicMock())
+
+ assert ctx.value is error
