This is an automated email from the ASF dual-hosted git repository.
potiuk 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 ab452960931 Fix pickle round-trip of Amazon provider exceptions and
add their tests (#72824)
ab452960931 is described below
commit ab45296093135e96be4df9c10b13194374d9d040
Author: Bingqin Wang <[email protected]>
AuthorDate: Thu Sep 10 02:12:15 2026 -0500
Fix pickle round-trip of Amazon provider exceptions and add their tests
(#72824)
EcsTaskFailToStart.__reduce__ returned a bare string instead of a
one-element tuple, so pickling it has raised PicklingError since #22002.
WaiterTerminalFailure kept last_response outside args, so unpickling it failed
with a missing-argument TypeError. Add a __reduce__ for it in the style of the
ECS exceptions, add the missing test module for aws/exceptions.py, and drop its
OVERLOOKED_TESTS entry.
---
.../tests/unit/always/test_project_structure.py | 1 -
.../src/airflow/providers/amazon/aws/exceptions.py | 7 +-
.../tests/unit/amazon/aws/test_exceptions.py | 137 +++++++++++++++++++++
3 files changed, 143 insertions(+), 2 deletions(-)
diff --git a/airflow-core/tests/unit/always/test_project_structure.py
b/airflow-core/tests/unit/always/test_project_structure.py
index f59a1f77823..509243e21e4 100644
--- a/airflow-core/tests/unit/always/test_project_structure.py
+++ b/airflow-core/tests/unit/always/test_project_structure.py
@@ -72,7 +72,6 @@ class TestProjectStructure:
"providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker.py",
"providers/amazon/tests/unit/amazon/aws/sensors/test_emr.py",
"providers/amazon/tests/unit/amazon/aws/sensors/test_sagemaker.py",
- "providers/amazon/tests/unit/amazon/aws/test_exceptions.py",
"providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker_unified_studio.py",
"providers/amazon/tests/unit/amazon/aws/utils/test_rds.py",
"providers/amazon/tests/unit/amazon/aws/utils/test_sagemaker.py",
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
b/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
index 5feb0637e8b..64c15f86189 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
@@ -34,7 +34,7 @@ class EcsTaskFailToStart(Exception):
def __reduce__(self):
"""Return ECSTask state and its message."""
- return EcsTaskFailToStart, (self.message)
+ return EcsTaskFailToStart, (self.message,)
class EcsOperatorError(Exception):
@@ -115,8 +115,13 @@ class WaiterTerminalFailure(AirflowException):
def __init__(self, message: str, last_response: dict[str, Any]):
super().__init__(message)
+ self.message = message
self.last_response = last_response
+ def __reduce__(self):
+ """Return the waiter failure state as its message and the last waiter
response."""
+ return WaiterTerminalFailure, (self.message, self.last_response)
+
class WaiterMaxAttemptsError(AirflowException):
"""Raised when an AWS waiter exhausts its configured attempts."""
diff --git a/providers/amazon/tests/unit/amazon/aws/test_exceptions.py
b/providers/amazon/tests/unit/amazon/aws/test_exceptions.py
new file mode 100644
index 00000000000..3714019eedb
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/test_exceptions.py
@@ -0,0 +1,137 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import pickle
+
+import pytest
+
+from airflow.providers.amazon.aws.exceptions import (
+ DataSyncLocationNotFoundError,
+ DataSyncMultipleLocationsError,
+ DataSyncMultipleTasksError,
+ DataSyncTaskCreationError,
+ DataSyncTaskExecutionFailedError,
+ DataSyncTaskNotFoundError,
+ EcsOperatorError,
+ EcsTaskFailToStart,
+ GlueJobRunStoppedError,
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+ S3HookPathTraversalError,
+ S3HookUriParseFailure,
+ WaiterMaxAttemptsError,
+ WaiterTerminalFailure,
+)
+from airflow.providers.common.compat.sdk import AirflowException
+
+ECS_FAILURES = [
+ {
+ "arn": "arn:aws:ecs:us-east-1:123456789012:container-instance/abc123",
+ "reason": "RESOURCE:MEMORY",
+ }
+]
+LAST_RESPONSE = {"Cluster": {"Status": "FAILED", "StatusReason": "Insufficient
capacity"}}
+
+AIRFLOW_EXCEPTION_SUBCLASSES = [
+ S3HookUriParseFailure,
+ S3HookPathTraversalError,
+ NeptuneGraphCreationFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ GlueJobRunStoppedError,
+ DataSyncTaskNotFoundError,
+ DataSyncMultipleTasksError,
+ DataSyncMultipleLocationsError,
+ DataSyncLocationNotFoundError,
+ DataSyncTaskCreationError,
+ DataSyncTaskExecutionFailedError,
+ WaiterMaxAttemptsError,
+]
+
+
+class TestEcsTaskFailToStart:
+ def test_message(self):
+ exc = EcsTaskFailToStart("The task failed to start due to:
OutOfMemoryError")
+
+ assert exc.message == "The task failed to start due to:
OutOfMemoryError"
+ assert str(exc) == "The task failed to start due to: OutOfMemoryError"
+
+ def test_pickle_round_trip(self):
+ exc = EcsTaskFailToStart("The task failed to start due to:
OutOfMemoryError")
+
+ restored = pickle.loads(pickle.dumps(exc))
+
+ assert isinstance(restored, EcsTaskFailToStart)
+ assert restored.message == exc.message
+ assert str(restored) == str(exc)
+
+
+class TestEcsOperatorError:
+ def test_failures_and_message(self):
+ exc = EcsOperatorError(ECS_FAILURES, "ECS could not run the task")
+
+ assert exc.failures == ECS_FAILURES
+ assert exc.message == "ECS could not run the task"
+ assert str(exc) == "ECS could not run the task"
+
+ def test_pickle_round_trip(self):
+ exc = EcsOperatorError(ECS_FAILURES, "ECS could not run the task")
+
+ restored = pickle.loads(pickle.dumps(exc))
+
+ assert isinstance(restored, EcsOperatorError)
+ assert restored.failures == ECS_FAILURES
+ assert restored.message == "ECS could not run the task"
+
+
+class TestWaiterTerminalFailure:
+ def test_message_and_last_response(self):
+ exc = WaiterTerminalFailure("Waiter reached a terminal failure state",
LAST_RESPONSE)
+
+ assert isinstance(exc, AirflowException)
+ assert str(exc) == "Waiter reached a terminal failure state"
+ assert exc.last_response == LAST_RESPONSE
+
+ def test_pickle_round_trip(self):
+ exc = WaiterTerminalFailure("Waiter reached a terminal failure state",
LAST_RESPONSE)
+
+ restored = pickle.loads(pickle.dumps(exc))
+
+ assert isinstance(restored, WaiterTerminalFailure)
+ assert str(restored) == "Waiter reached a terminal failure state"
+ assert restored.last_response == LAST_RESPONSE
+
+
[email protected]("exc_cls", AIRFLOW_EXCEPTION_SUBCLASSES)
+def test_airflow_exception_subclasses_keep_message(exc_cls):
+ exc = exc_cls("something went wrong")
+
+ assert isinstance(exc, AirflowException)
+ assert str(exc) == "something went wrong"
+
+ restored = pickle.loads(pickle.dumps(exc))
+
+ assert isinstance(restored, exc_cls)
+ assert str(restored) == "something went wrong"
