This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 6815d040ed1 Add dedicated exceptions for terminal waiter failures and
max-attempt exhaustion, and allow AWS triggers to customize exception-to-event
translation without reimplementing the shared waiter lifecycle. (#72455)
6815d040ed1 is described below
commit 6815d040ed1b170477904f3f1c31316d00b855b4
Author: SameerMesiah97 <[email protected]>
AuthorDate: Fri Sep 4 13:47:41 2026 +0100
Add dedicated exceptions for terminal waiter failures and max-attempt
exhaustion, and allow AWS triggers to customize exception-to-event translation
without reimplementing the shared waiter lifecycle. (#72455)
---
generated/known_airflow_exceptions.txt | 2 +-
.../src/airflow/providers/amazon/aws/exceptions.py | 14 +++++
.../airflow/providers/amazon/aws/triggers/base.py | 13 ++++-
.../amazon/aws/utils/waiter_with_logging.py | 15 ++++--
.../tests/unit/amazon/aws/triggers/test_base.py | 12 +++++
.../amazon/aws/utils/test_waiter_with_logging.py | 63 +++++++++++++++++++++-
6 files changed, 109 insertions(+), 10 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index bfb87136d2e..607b038da83 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -125,7 +125,7 @@
providers/amazon/src/airflow/providers/amazon/aws/triggers/sqs.py::1
providers/amazon/src/airflow/providers/amazon/aws/utils/__init__.py::1
providers/amazon/src/airflow/providers/amazon/aws/utils/connection_wrapper.py::2
providers/amazon/src/airflow/providers/amazon/aws/utils/waiter.py::1
-providers/amazon/src/airflow/providers/amazon/aws/utils/waiter_with_logging.py::6
+providers/amazon/src/airflow/providers/amazon/aws/utils/waiter_with_logging.py::2
providers/amazon/tests/unit/amazon/aws/operators/test_ecs.py::1
providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py::8
providers/apache/beam/src/airflow/providers/apache/beam/operators/beam.py::4
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
b/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
index 9760067ce5b..5feb0637e8b 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
@@ -17,6 +17,8 @@
# under the License.
from __future__ import annotations
+from typing import Any
+
from airflow.providers.common.compat.sdk import AirflowException
# Note: Any AirflowException raised is expected to cause the TaskInstance
@@ -106,3 +108,15 @@ class DataSyncTaskCreationError(AirflowException):
class DataSyncTaskExecutionFailedError(AirflowException):
"""Raised when a DataSync task execution could not be started or did not
complete successfully."""
+
+
+class WaiterTerminalFailure(AirflowException):
+ """Raised when an AWS waiter reaches a terminal failure state."""
+
+ def __init__(self, message: str, last_response: dict[str, Any]):
+ super().__init__(message)
+ self.last_response = last_response
+
+
+class WaiterMaxAttemptsError(AirflowException):
+ """Raised when an AWS waiter exhausts its configured attempts."""
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py
index 999b5e5bfd1..b952e27fbba 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py
@@ -141,6 +141,15 @@ class AwsBaseWaiterTrigger(BaseTrigger):
def hook(self) -> AwsGenericHook:
"""Override in subclasses to return the right hook."""
+ def _event_from_exception(self, error: AirflowException) -> TriggerEvent:
+ return TriggerEvent(
+ {
+ "status": "error",
+ "message": str(error),
+ self.return_key: self.return_value,
+ }
+ )
+
async def run(self) -> AsyncIterator[TriggerEvent]:
hook = self.hook()
async with await hook.get_async_conn() as client:
@@ -160,7 +169,7 @@ class AwsBaseWaiterTrigger(BaseTrigger):
self.status_message,
self.status_queries,
)
- except AirflowException as e:
- yield TriggerEvent({"status": "error", "message": str(e),
self.return_key: self.return_value})
+ except AirflowException as error:
+ yield self._event_from_exception(error)
else:
yield TriggerEvent({"status": "success", self.return_key:
self.return_value})
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/utils/waiter_with_logging.py
b/providers/amazon/src/airflow/providers/amazon/aws/utils/waiter_with_logging.py
index 60e20bbf2cd..6ae8ecf7d33 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/utils/waiter_with_logging.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/utils/waiter_with_logging.py
@@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, Any
import jmespath
from botocore.exceptions import NoCredentialsError, WaiterError
+from airflow.providers.amazon.aws.exceptions import WaiterMaxAttemptsError,
WaiterTerminalFailure
from airflow.providers.common.compat.sdk import AirflowException
if TYPE_CHECKING:
@@ -105,7 +106,10 @@ def wait(
if "terminal failure" in error_reason:
log.error("%s: %s", failure_message,
_LazyStatusFormatter(status_args, last_response))
- raise AirflowException(f"{failure_message}: {error}")
+ raise WaiterTerminalFailure(
+ f"{failure_message}: {error}",
+ last_response=last_response,
+ )
if (
"An error occurred" in error_reason
@@ -130,7 +134,7 @@ def wait(
break
attempt += 1
else:
- raise AirflowException("Waiter error: max attempts reached")
+ raise WaiterMaxAttemptsError("Waiter error: max attempts reached")
async def async_wait(
@@ -184,8 +188,9 @@ async def async_wait(
last_response = error.last_response
if "terminal failure" in error_reason:
- raise AirflowException(
- f"{failure_message}: {_LazyStatusFormatter(status_args,
last_response)}\n{error}"
+ raise WaiterTerminalFailure(
+ f"{failure_message}: {_LazyStatusFormatter(status_args,
last_response)}\n{error}",
+ last_response=last_response,
)
if (
@@ -211,7 +216,7 @@ async def async_wait(
break
attempt += 1
else:
- raise AirflowException("Waiter error: max attempts reached")
+ raise WaiterMaxAttemptsError("Waiter error: max attempts reached")
class _LazyStatusFormatter:
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py
index 6423ff9c38c..0c2ea69149a 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_base.py
@@ -144,3 +144,15 @@ class TestAwsBaseWaiterTrigger:
assert res.payload["status"] == "error"
assert "AWS Glue job failed." in res.payload["message"]
assert res.payload["hello"] == "world"
+
+ def test_event_from_exception(self):
+ self.trigger.return_key = "hello"
+ self.trigger.return_value = "world"
+
+ event = self.trigger._event_from_exception(AirflowException("AWS Glue
job failed."))
+
+ assert event.payload == {
+ "status": "error",
+ "message": "AWS Glue job failed.",
+ "hello": "world",
+ }
diff --git
a/providers/amazon/tests/unit/amazon/aws/utils/test_waiter_with_logging.py
b/providers/amazon/tests/unit/amazon/aws/utils/test_waiter_with_logging.py
index 3414403578e..817be27df2d 100644
--- a/providers/amazon/tests/unit/amazon/aws/utils/test_waiter_with_logging.py
+++ b/providers/amazon/tests/unit/amazon/aws/utils/test_waiter_with_logging.py
@@ -25,6 +25,10 @@ from unittest.mock import AsyncMock
import pytest
from botocore.exceptions import WaiterError
+from airflow.providers.amazon.aws.exceptions import (
+ WaiterMaxAttemptsError,
+ WaiterTerminalFailure,
+)
from airflow.providers.amazon.aws.utils.waiter_with_logging import
_LazyStatusFormatter, async_wait, wait
from airflow.providers.common.compat.sdk import AirflowException
@@ -141,7 +145,7 @@ class TestWaiter:
last_response=generate_response("Pending"),
)
mock_waiter.wait.side_effect = [error, error, error]
- with pytest.raises(AirflowException) as exc:
+ with pytest.raises(WaiterMaxAttemptsError) as exc:
wait(
waiter=mock_waiter,
waiter_delay=123,
@@ -162,6 +166,30 @@ class TestWaiter:
assert mock_waiter.wait.call_count == 2
mock_sleep.assert_called_with(123)
+ @pytest.mark.asyncio
+ async def test_async_wait_max_attempts_exceeded(self):
+ mock_waiter = mock.MagicMock()
+ error = WaiterError(
+ name="test_waiter",
+ reason="test_reason",
+ last_response=generate_response("Pending"),
+ )
+ mock_waiter.wait = AsyncMock(side_effect=error)
+
+ with pytest.raises(WaiterMaxAttemptsError) as exc:
+ await async_wait(
+ waiter=mock_waiter,
+ waiter_delay=0,
+ waiter_max_attempts=2,
+ args={"test_arg": "test_value"},
+ failure_message="test failure message",
+ status_message="test status message",
+ status_args=["Status.State"],
+ )
+
+ assert "Waiter error: max attempts reached" in str(exc.value)
+ assert mock_waiter.wait.call_count == 2
+
@mock.patch("time.sleep")
def test_wait_with_failure(self, mock_sleep):
mock_sleep.return_value = True
@@ -178,7 +206,7 @@ class TestWaiter:
)
mock_waiter.wait.side_effect = [error, error, error, failure_error]
- with pytest.raises(AirflowException) as exc:
+ with pytest.raises(WaiterTerminalFailure) as exc:
wait(
waiter=mock_waiter,
waiter_delay=123,
@@ -198,6 +226,37 @@ class TestWaiter:
)
assert mock_waiter.wait.call_count == 4
+ @pytest.mark.asyncio
+ async def test_async_wait_with_failure(self):
+ mock_waiter = mock.MagicMock()
+ error = WaiterError(
+ name="test_waiter",
+ reason="test_reason",
+ last_response=generate_response("Pending"),
+ )
+ last_response = generate_response("Failure")
+ failure_error = WaiterError(
+ name="test_waiter",
+ reason="terminal failure in waiter",
+ last_response=last_response,
+ )
+ mock_waiter.wait = AsyncMock(side_effect=[error, error, failure_error])
+
+ with pytest.raises(WaiterTerminalFailure) as exc:
+ await async_wait(
+ waiter=mock_waiter,
+ waiter_delay=0,
+ waiter_max_attempts=10,
+ args={"test_arg": "test_value"},
+ failure_message="test failure message",
+ status_message="test status message",
+ status_args=["Status.State"],
+ )
+
+ assert "test failure message" in str(exc.value)
+ assert exc.value.last_response == last_response
+ assert mock_waiter.wait.call_count == 3
+
@mock.patch("time.sleep")
def test_wait_with_unknown_failure(self, mock_sleep):
mock_sleep.return_value = True
