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 4fb12b3af09 Include DMS stop reason in DmsTaskBaseSensor failure
message (#75240)
4fb12b3af09 is described below
commit 4fb12b3af09892a7a40d18dd0ae5bfd0dbba08e4
Author: Niko Oliveira <[email protected]>
AuthorDate: Wed Oct 21 06:49:44 2026 -0700
Include DMS stop reason in DmsTaskBaseSensor failure message (#75240)
When a replication task reached a termination status the sensor only
reported the status name ("Unexpected status: stopped"), even though the
DescribeReplicationTasks response it already reads on every poke carries
StopReason and LastFailureMessage. Once the task is deleted (e.g. by a
teardown task) that information is gone, so a task that stopped
unexpectedly could not be diagnosed after the fact.
Status and reason are read from the same describe call so they always
describe the same snapshot, and the number of API calls per poke stays
the same.
---
.../src/airflow/providers/amazon/aws/hooks/dms.py | 24 ++++++++++----
.../airflow/providers/amazon/aws/sensors/dms.py | 11 +++++--
.../amazon/tests/unit/amazon/aws/hooks/test_dms.py | 6 ++++
.../tests/unit/amazon/aws/sensors/test_dms.py | 37 +++++++++++++++-------
4 files changed, 57 insertions(+), 21 deletions(-)
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/dms.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/dms.py
index b08dc68621c..4f45d1ad016 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/dms.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/dms.py
@@ -118,12 +118,12 @@ class DmsHook(AwsBaseHook):
return tasks
- def get_task_status(self, replication_task_arn: str) -> str | None:
+ def get_task_description(self, replication_task_arn: str) -> dict | None:
"""
- Retrieve task status.
+ Retrieve the replication task description, without task settings.
:param replication_task_arn: Replication task ARN
- :return: Current task status
+ :return: The task as returned by ``DescribeReplicationTasks``, or
``None`` if not found
"""
replication_tasks = self.find_replication_tasks_by_arn(
replication_task_arn=replication_task_arn,
@@ -131,12 +131,24 @@ class DmsHook(AwsBaseHook):
)
if len(replication_tasks) == 1:
- status = replication_tasks[0]["Status"]
- self.log.info('Replication task with ARN(%s) has status "%s".',
replication_task_arn, status)
- return status
+ task = replication_tasks[0]
+ self.log.info(
+ 'Replication task with ARN(%s) has status "%s".',
replication_task_arn, task["Status"]
+ )
+ return task
self.log.info("Replication task with ARN(%s) is not found.",
replication_task_arn)
return None
+ def get_task_status(self, replication_task_arn: str) -> str | None:
+ """
+ Retrieve task status.
+
+ :param replication_task_arn: Replication task ARN
+ :return: Current task status
+ """
+ task = self.get_task_description(replication_task_arn)
+ return task["Status"] if task else None
+
def create_replication_task(
self,
replication_task_id: str,
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/dms.py
b/providers/amazon/src/airflow/providers/amazon/aws/sensors/dms.py
index 47eb56f10e9..ba5e4057263 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/dms.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/dms.py
@@ -39,7 +39,8 @@ class DmsTaskBaseSensor(AwsBaseSensor[DmsHook]):
:param target_statuses: the target statuses, sensor waits until
the task reaches any of these states
:param termination_statuses: the termination statuses, sensor fails when
- the task reaches any of these states
+ the task reaches any of these states. The error includes the task's
+ ``StopReason`` and ``LastFailureMessage`` when DMS reports them.
:param aws_conn_id: The Airflow connection used for AWS credentials.
If this is ``None`` or empty then the default boto3 behaviour is used.
If
running Airflow in a distributed manner and aws_conn_id is None or
@@ -68,7 +69,8 @@ class DmsTaskBaseSensor(AwsBaseSensor[DmsHook]):
self.termination_statuses: Iterable[str] = termination_statuses or []
def poke(self, context: Context):
- if not (status :=
self.hook.get_task_status(self.replication_task_arn)):
+ task = self.hook.get_task_description(self.replication_task_arn) or {}
+ if not (status := task.get("Status")):
raise AirflowException(
f"Failed to read task status, task with ARN
{self.replication_task_arn} not found"
)
@@ -79,7 +81,10 @@ class DmsTaskBaseSensor(AwsBaseSensor[DmsHook]):
return True
if status in self.termination_statuses:
- raise AirflowException(f"Unexpected status: {status}")
+ details = ", ".join(
+ f"{field}: {task[field]}" for field in ("StopReason",
"LastFailureMessage") if task.get(field)
+ )
+ raise AirflowException(f"Unexpected status: {status}" + (f"
({details})" if details else ""))
return False
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_dms.py
b/providers/amazon/tests/unit/amazon/aws/hooks/test_dms.py
index c1dc96bd9f4..9c807df25a8 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_dms.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_dms.py
@@ -309,6 +309,12 @@ class TestDmsHook:
assert len(tasks) == 1
assert tasks[0]["ReplicationTaskArn"] == MOCK_TASK_ARN
+ @mock.patch.object(DmsHook, "get_conn")
+ def test_get_task_description(self, mock_conn):
+ mock_conn.return_value.describe_replication_tasks.return_value =
MOCK_DESCRIBE_RESPONSE
+
+ assert self.dms.get_task_description(MOCK_TASK_ARN) ==
MOCK_TASK_RESPONSE_DATA
+
@mock.patch.object(DmsHook, "get_conn")
def test_get_task_status(self, mock_conn):
mock_conn.return_value.describe_replication_tasks.return_value =
MOCK_DESCRIBE_RESPONSE
diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_dms.py
b/providers/amazon/tests/unit/amazon/aws/sensors/test_dms.py
index 31c8edb6cca..7491e307123 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_dms.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_dms.py
@@ -26,8 +26,8 @@ from airflow.providers.common.compat.sdk import
AirflowException
@pytest.fixture
-def mocked_get_task_status():
- with mock.patch.object(DmsHook, "get_task_status") as m:
+def mocked_get_task_description():
+ with mock.patch.object(DmsHook, "get_task_description", autospec=True) as
m:
yield m
@@ -65,13 +65,13 @@ class TestDmsTaskCompletedSensor:
assert sensor.hook._config is None
@pytest.mark.parametrize("status", ["stopped"])
- def test_poke_completed(self, mocked_get_task_status, status):
- mocked_get_task_status.return_value = status
+ def test_poke_completed(self, mocked_get_task_description, status):
+ mocked_get_task_description.return_value = {"Status": status}
assert DmsTaskCompletedSensor(**self.default_op_kwargs).poke({})
@pytest.mark.parametrize("status", ["running", "starting"])
- def test_poke_not_completed(self, mocked_get_task_status, status):
- mocked_get_task_status.return_value = status
+ def test_poke_not_completed(self, mocked_get_task_description, status):
+ mocked_get_task_description.return_value = {"Status": status}
assert not DmsTaskCompletedSensor(**self.default_op_kwargs).poke({})
@pytest.mark.parametrize(
@@ -87,13 +87,26 @@ class TestDmsTaskCompletedSensor:
"testing",
],
)
- def test_poke_terminated_status(self, mocked_get_task_status, status):
- mocked_get_task_status.return_value = status
- error_message = f"Unexpected status: {status}"
- with pytest.raises(AirflowException, match=error_message):
+ def test_poke_terminated_status(self, mocked_get_task_description, status):
+ mocked_get_task_description.return_value = {"Status": status}
+ with pytest.raises(AirflowException) as ctx:
DmsTaskCompletedSensor(**self.default_op_kwargs).poke({})
+ assert str(ctx.value) == f"Unexpected status: {status}"
- def test_poke_none_status(self, mocked_get_task_status):
- mocked_get_task_status.return_value = None
+ def test_poke_terminated_status_includes_reason(self,
mocked_get_task_description):
+ mocked_get_task_description.return_value = {
+ "Status": "failed",
+ "StopReason": "Stop Reason FATAL_ERROR",
+ "LastFailureMessage": "Last Error",
+ }
+ with pytest.raises(AirflowException) as ctx:
+ DmsTaskCompletedSensor(**self.default_op_kwargs).poke({})
+ assert (
+ str(ctx.value)
+ == "Unexpected status: failed (StopReason: Stop Reason
FATAL_ERROR, LastFailureMessage: Last Error)"
+ )
+
+ def test_poke_none_status(self, mocked_get_task_description):
+ mocked_get_task_description.return_value = None
with pytest.raises(AirflowException, match="task with ARN .* not
found"):
DmsTaskCompletedSensor(**self.default_op_kwargs).poke({})
