This is an automated email from the ASF dual-hosted git repository.
o-nikolas 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 6b81cda9b23 Narrow AirflowException to specific exceptions in DataSync
operator (#70152)
6b81cda9b23 is described below
commit 6b81cda9b23d677382879ae91a521792abe8f8fd
Author: Haseeb Malik <[email protected]>
AuthorDate: Fri Aug 14 20:29:43 2026 -0400
Narrow AirflowException to specific exceptions in DataSync operator (#70152)
---
generated/known_airflow_exceptions.txt | 1 -
.../src/airflow/providers/amazon/aws/exceptions.py | 24 ++++++
.../providers/amazon/aws/operators/datasync.py | 39 ++++++----
.../unit/amazon/aws/operators/test_datasync.py | 91 ++++++++++++++++++----
4 files changed, 127 insertions(+), 28 deletions(-)
diff --git a/generated/known_airflow_exceptions.txt
b/generated/known_airflow_exceptions.txt
index 8be3fafad4a..202df42602e 100644
--- a/generated/known_airflow_exceptions.txt
+++ b/generated/known_airflow_exceptions.txt
@@ -68,7 +68,6 @@
providers/amazon/src/airflow/providers/amazon/aws/operators/athena.py::3
providers/amazon/src/airflow/providers/amazon/aws/operators/batch.py::5
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py::6
providers/amazon/src/airflow/providers/amazon/aws/operators/comprehend.py::2
-providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py::10
providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py::5
providers/amazon/src/airflow/providers/amazon/aws/operators/ec2.py::1
providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py::7
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
b/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
index fc6d668ae74..9760067ce5b 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/exceptions.py
@@ -82,3 +82,27 @@ class NeptuneImportTaskFailedError(AirflowException):
class GlueJobRunStoppedError(AirflowException):
"""Raised when a Glue job run finishes in a state that is not a real
success."""
+
+
+class DataSyncTaskNotFoundError(AirflowException):
+ """Raised when a DataSync task could not be identified or created for the
requested locations."""
+
+
+class DataSyncMultipleTasksError(AirflowException):
+ """Raised when multiple DataSync tasks match and random task choice is not
allowed."""
+
+
+class DataSyncMultipleLocationsError(AirflowException):
+ """Raised when multiple DataSync locations match and random location
choice is not allowed."""
+
+
+class DataSyncLocationNotFoundError(AirflowException):
+ """Raised when a DataSync location could not be determined or created."""
+
+
+class DataSyncTaskCreationError(AirflowException):
+ """Raised when DataSync task creation did not return a task ARN."""
+
+
+class DataSyncTaskExecutionFailedError(AirflowException):
+ """Raised when a DataSync task execution could not be started or did not
complete successfully."""
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py
index 1b97ec399af..936949b7ab1 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py
@@ -23,6 +23,14 @@ import random
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
+from airflow.providers.amazon.aws.exceptions import (
+ DataSyncLocationNotFoundError,
+ DataSyncMultipleLocationsError,
+ DataSyncMultipleTasksError,
+ DataSyncTaskCreationError,
+ DataSyncTaskExecutionFailedError,
+ DataSyncTaskNotFoundError,
+)
from airflow.providers.amazon.aws.hooks.datasync import DataSyncHook
from airflow.providers.amazon.aws.links.datasync import
DataSyncTaskExecutionLink, DataSyncTaskLink
from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
@@ -102,13 +110,16 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore
client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
- :raises AirflowException: If ``task_arn`` was not specified, or if
+ :raises ValueError: If ``task_arn`` was not specified, or if
either ``source_location_uri`` or ``destination_location_uri`` were
not specified.
- :raises AirflowException: If source or destination Location were not found
+ :raises DataSyncLocationNotFoundError: If source or destination Location
were not found
and could not be created.
- :raises AirflowException: If ``choose_task`` or ``choose_location`` fails.
- :raises AirflowException: If Task creation, update, execution or delete
fails.
+ :raises DataSyncMultipleTasksError: If ``choose_task`` fails.
+ :raises DataSyncMultipleLocationsError: If ``choose_location`` fails.
+ :raises DataSyncTaskNotFoundError: If a task could not be identified or
created.
+ :raises DataSyncTaskCreationError: If Task creation fails.
+ :raises DataSyncTaskExecutionFailedError: If Task execution fails.
"""
aws_hook_class = DataSyncHook
@@ -195,7 +206,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
if self.source_location_uri and self.destination_location_uri:
valid = True
if not valid:
- raise AirflowException(
+ raise ValueError(
f"Either specify task_arn or both source_location_uri and
destination_location_uri. "
f"task_arn={self.task_arn!r},
source_location_uri={self.source_location_uri!r}, "
f"destination_location_uri={self.destination_location_uri!r}"
@@ -218,7 +229,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
self._create_datasync_task()
if not self.task_arn:
- raise AirflowException("DataSync TaskArn could not be identified
or created.")
+ raise DataSyncTaskNotFoundError("DataSync TaskArn could not be
identified or created.")
task_id = self.task_arn.split("/")[-1]
@@ -247,7 +258,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
self._execute_datasync_task(context=context)
if not self.task_execution_arn:
- raise AirflowException("Nothing was executed")
+ raise DataSyncTaskExecutionFailedError("Nothing was executed")
# Delete the DataSyncTask
if self.delete_task_after_execution:
@@ -288,7 +299,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
# from AWS and might lead to confusion. Rather explicitly
# choose a random one
return random.choice(task_arn_list)
- raise AirflowException(f"Unable to choose a Task from {task_arn_list}")
+ raise DataSyncMultipleTasksError(f"Unable to choose a Task from
{task_arn_list}")
def choose_location(self, location_arn_list: list[str] | None) -> str |
None:
"""Select 1 DataSync LocationArn from a list."""
@@ -302,7 +313,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
# from AWS and might lead to confusion. Rather explicitly
# choose a random one
return random.choice(location_arn_list)
- raise AirflowException(f"Unable to choose a Location from
{location_arn_list}")
+ raise DataSyncMultipleLocationsError(f"Unable to choose a Location
from {location_arn_list}")
def _create_datasync_task(self) -> None:
"""Create a AWS DataSyncTask."""
@@ -313,7 +324,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
self.source_location_uri, **self.create_source_location_kwargs
)
if not self.source_location_arn:
- raise AirflowException(
+ raise DataSyncLocationNotFoundError(
"Unable to determine source LocationArn. Does a suitable
DataSync Location exist?"
)
@@ -328,7 +339,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
self.destination_location_uri,
**self.create_destination_location_kwargs
)
if not self.destination_location_arn:
- raise AirflowException(
+ raise DataSyncLocationNotFoundError(
"Unable to determine destination LocationArn. Does a suitable
DataSync Location exist?"
)
@@ -337,7 +348,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
self.source_location_arn, self.destination_location_arn,
**self.create_task_kwargs
)
if not self.task_arn:
- raise AirflowException("Task could not be created")
+ raise DataSyncTaskCreationError("Task could not be created")
self.log.info("Created a Task with TaskArn %s", self.task_arn)
def _update_datasync_task(self) -> None:
@@ -352,7 +363,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
def _execute_datasync_task(self, context: Context) -> None:
"""Create and monitor an AWS DataSync TaskExecution for a Task."""
if not self.task_arn:
- raise AirflowException("Missing TaskArn")
+ raise ValueError("Missing TaskArn")
# Create a task execution:
self.log.info("Starting execution for TaskArn %s", self.task_arn)
@@ -406,7 +417,7 @@ class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
self.log.log(level, "%s=%s", k, v)
if not result:
- raise AirflowException(f"Failed TaskExecutionArn
{self.task_execution_arn}")
+ raise DataSyncTaskExecutionFailedError(f"Failed TaskExecutionArn
{self.task_execution_arn}")
def _cancel_datasync_task_execution(self):
"""Cancel the submitted DataSync task."""
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py
index 40a9f8788f9..959f8d609b2 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py
@@ -23,10 +23,17 @@ import pytest
from moto import mock_aws
from airflow.models import DAG, DagRun, TaskInstance
+from airflow.providers.amazon.aws.exceptions import (
+ DataSyncLocationNotFoundError,
+ DataSyncMultipleLocationsError,
+ DataSyncMultipleTasksError,
+ DataSyncTaskCreationError,
+ DataSyncTaskExecutionFailedError,
+ DataSyncTaskNotFoundError,
+)
from airflow.providers.amazon.aws.hooks.datasync import DataSyncHook
from airflow.providers.amazon.aws.links.datasync import DataSyncTaskLink
from airflow.providers.amazon.aws.operators.datasync import DataSyncOperator
-from airflow.providers.common.compat.sdk import AirflowException
from airflow.utils.state import DagRunState
from airflow.utils.types import DagRunType
@@ -212,13 +219,13 @@ class TestDataSyncOperatorCreate(DataSyncTestCaseBase):
# ### Begin tests:
self.set_up_operator(task_id="task_1", source_location_uri=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
self.set_up_operator(task_id="task_2", destination_location_uri=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
self.set_up_operator(task_id="task_3", source_location_uri=None,
destination_location_uri=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
# ### Check mocks:
mock_get_conn.assert_not_called()
@@ -320,7 +327,7 @@ class TestDataSyncOperatorCreate(DataSyncTestCaseBase):
self.client.create_location_smb(**MOCK_DATA["create_source_location_kwargs"])
self.set_up_operator(task_id="datasync_task1")
- with pytest.raises(AirflowException):
+ with pytest.raises(DataSyncMultipleLocationsError):
self.datasync.execute(None)
# Delete all tasks:
@@ -333,6 +340,64 @@ class TestDataSyncOperatorCreate(DataSyncTestCaseBase):
# ### Check mocks:
mock_get_conn.assert_called()
+ def test_no_task_identified_or_created(self, mock_get_conn):
+ # ### Set up mocks:
+ mock_get_conn.return_value = self.client
+ # ### Begin tests:
+
+ # Delete all tasks so none can be matched.
+ tasks = self.client.list_tasks()
+ for task in tasks["Tasks"]:
+ self.client.delete_task(TaskArn=task["TaskArn"])
+
+ # Without create_task_kwargs there is nothing to run or create.
+ self.datasync = DataSyncOperator(
+ task_id="datasync_no_task",
+ dag=self.dag,
+ source_location_uri=SOURCE_LOCATION_URI,
+ destination_location_uri=DESTINATION_LOCATION_URI,
+ wait_interval_seconds=0,
+ )
+ with pytest.raises(DataSyncTaskNotFoundError):
+ self.datasync.execute(None)
+ # ### Check mocks:
+ mock_get_conn.assert_called()
+
+ def test_create_task_no_source_location(self, mock_get_conn):
+ # ### Set up mocks:
+ mock_get_conn.return_value = self.client
+ # ### Begin tests:
+
+ self.datasync = DataSyncOperator(
+ task_id="datasync_no_source_location",
+ dag=self.dag,
+ source_location_uri="smb://nowhere/subdir",
+ destination_location_uri=DESTINATION_LOCATION_URI,
+ create_task_kwargs={"Options": {"VerifyMode": "NONE"}},
+ wait_interval_seconds=0,
+ )
+ with pytest.raises(DataSyncLocationNotFoundError):
+ self.datasync.execute(None)
+ # ### Check mocks:
+ mock_get_conn.assert_called()
+
+ @mock.patch.object(DataSyncHook, "create_task", return_value=None)
+ def test_create_task_without_task_arn(self, mock_create_task,
mock_get_conn):
+ # ### Set up mocks:
+ mock_get_conn.return_value = self.client
+ # ### Begin tests:
+
+ # Delete all tasks so the operator falls through to creation.
+ tasks = self.client.list_tasks()
+ for task in tasks["Tasks"]:
+ self.client.delete_task(TaskArn=task["TaskArn"])
+
+ self.set_up_operator()
+ with pytest.raises(DataSyncTaskCreationError):
+ self.datasync.execute(None)
+ # ### Check mocks:
+ mock_get_conn.assert_called()
+
def test_execute_specific_task(self, mock_get_conn):
# ### Set up mocks:
mock_get_conn.return_value = self.client
@@ -439,13 +504,13 @@ class TestDataSyncOperatorGetTasks(DataSyncTestCaseBase):
# ### Begin tests:
self.set_up_operator(task_id="task_1", source_location_uri=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
self.set_up_operator(task_id="task_2", destination_location_uri=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
self.set_up_operator(task_id="task_3", source_location_uri=None,
destination_location_uri=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
# ### Check mocks:
mock_get_conn.assert_not_called()
@@ -543,7 +608,7 @@ class TestDataSyncOperatorGetTasks(DataSyncTestCaseBase):
assert len(locations["Locations"]) == 2
# Execute the task
- with pytest.raises(AirflowException):
+ with pytest.raises(DataSyncMultipleTasksError):
self.datasync.execute(None)
# Assert 0 additional task and 0 additional locations
@@ -652,7 +717,7 @@ class TestDataSyncOperatorUpdate(DataSyncTestCaseBase):
# ### Begin tests:
self.set_up_operator(task_arn=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
# ### Check mocks:
mock_get_conn.assert_not_called()
@@ -774,7 +839,7 @@ class TestDataSyncOperator(DataSyncTestCaseBase):
# ### Begin tests:
self.set_up_operator(task_arn=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
# ### Check mocks:
mock_get_conn.assert_not_called()
@@ -863,7 +928,7 @@ class TestDataSyncOperator(DataSyncTestCaseBase):
self.set_up_operator()
# Execute the task
- with pytest.raises(AirflowException):
+ with pytest.raises(DataSyncTaskExecutionFailedError):
self.datasync.execute(None)
# ### Check mocks:
mock_get_conn.assert_called()
@@ -987,7 +1052,7 @@ class TestDataSyncOperatorDelete(DataSyncTestCaseBase):
# ### Begin tests:
self.set_up_operator(task_arn=None)
- with pytest.raises(AirflowException):
+ with pytest.raises(ValueError, match="Either specify task_arn"):
self.datasync.execute(None)
# ### Check mocks:
mock_get_conn.assert_not_called()
