From 2ea48ae1bf834819e583164cd1d9ba6811d99ad2 Mon Sep 17 00:00:00 2001 From: avcu Date: Thu, 23 Jul 2026 15:27:01 -0400 Subject: [PATCH 1/2] Move template-field validation out of init for Amazon DataSync operator --- .../amazon/aws/operators/datasync.py | 28 +++++++------- .../amazon/aws/operators/test_datasync.py | 37 ++++++++++++------- 2 files changed, 38 insertions(+), 27 deletions(-) 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 9345278ad5370..1b97ec399afae 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py @@ -174,19 +174,6 @@ def __init__( self.task_execution_kwargs = task_execution_kwargs or {} self.delete_task_after_execution = delete_task_after_execution - # Validations - valid = False - if self.task_arn: - valid = True - if self.source_location_uri and self.destination_location_uri: - valid = True - if not valid: - raise AirflowException( - f"Either specify task_arn or both source_location_uri and destination_location_uri. " - f"task_arn={task_arn!r}, source_location_uri={source_location_uri!r}, " - f"destination_location_uri={destination_location_uri!r}" - ) - # Candidates - these are found in AWS as possible things # for us to use self.candidate_source_location_arns: list[str] | None = None @@ -201,7 +188,22 @@ def __init__( def _hook_parameters(self) -> dict[str, Any]: return {**super()._hook_parameters, "wait_interval_seconds": self.wait_interval_seconds} + def validate_inputs(self) -> None: + valid = False + if self.task_arn: + valid = True + if self.source_location_uri and self.destination_location_uri: + valid = True + if not valid: + raise AirflowException( + 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}" + ) + def execute(self, context: Context): + self.validate_inputs() + # If task_arn was not specified then try to # find 0, 1 or many candidate DataSync Tasks to run if not self.task_arn: 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 2addb0ce539e4..40a9f8788f95d 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py @@ -206,17 +206,20 @@ def test_init(self, mock_get_conn): # ### Check mocks: mock_get_conn.assert_not_called() - def test_init_fails(self, mock_get_conn): + def test_execute_fails(self, mock_get_conn): # ### Set up mocks: mock_get_conn.return_value = self.client # ### Begin tests: + self.set_up_operator(task_id="task_1", source_location_uri=None) with pytest.raises(AirflowException): - self.set_up_operator(source_location_uri=None) + self.datasync.execute(None) + self.set_up_operator(task_id="task_2", destination_location_uri=None) with pytest.raises(AirflowException): - self.set_up_operator(destination_location_uri=None) + self.datasync.execute(None) + self.set_up_operator(task_id="task_3", source_location_uri=None, destination_location_uri=None) with pytest.raises(AirflowException): - self.set_up_operator(source_location_uri=None, destination_location_uri=None) + self.datasync.execute(None) # ### Check mocks: mock_get_conn.assert_not_called() @@ -430,17 +433,20 @@ def test_init(self, mock_get_conn): # ### Check mocks: mock_get_conn.assert_not_called() - def test_init_fails(self, mock_get_conn): + def test_execute_fails(self, mock_get_conn): # ### Set up mocks: mock_get_conn.return_value = self.client # ### Begin tests: + self.set_up_operator(task_id="task_1", source_location_uri=None) with pytest.raises(AirflowException): - self.set_up_operator(source_location_uri=None) + self.datasync.execute(None) + self.set_up_operator(task_id="task_2", destination_location_uri=None) with pytest.raises(AirflowException): - self.set_up_operator(destination_location_uri=None) + self.datasync.execute(None) + self.set_up_operator(task_id="task_3", source_location_uri=None, destination_location_uri=None) with pytest.raises(AirflowException): - self.set_up_operator(source_location_uri=None, destination_location_uri=None) + self.datasync.execute(None) # ### Check mocks: mock_get_conn.assert_not_called() @@ -640,13 +646,14 @@ def test_init(self, mock_get_conn): # ### Check mocks: mock_get_conn.assert_not_called() - def test_init_fails(self, mock_get_conn): + def test_execute_fails(self, mock_get_conn): # ### Set up mocks: mock_get_conn.return_value = self.client # ### Begin tests: + self.set_up_operator(task_arn=None) with pytest.raises(AirflowException): - self.set_up_operator(task_arn=None) + self.datasync.execute(None) # ### Check mocks: mock_get_conn.assert_not_called() @@ -761,13 +768,14 @@ def test_init(self, mock_get_conn): # ### Check mocks: mock_get_conn.assert_not_called() - def test_init_fails(self, mock_get_conn): + def test_execute_fails(self, mock_get_conn): # ### Set up mocks: mock_get_conn.return_value = self.client # ### Begin tests: + self.set_up_operator(task_arn=None) with pytest.raises(AirflowException): - self.set_up_operator(task_arn=None) + self.datasync.execute(None) # ### Check mocks: mock_get_conn.assert_not_called() @@ -973,13 +981,14 @@ def test_init(self, mock_get_conn): # ### Check mocks: mock_get_conn.assert_not_called() - def test_init_fails(self, mock_get_conn): + def test_execute_fails(self, mock_get_conn): # ### Set up mocks: mock_get_conn.return_value = self.client # ### Begin tests: + self.set_up_operator(task_arn=None) with pytest.raises(AirflowException): - self.set_up_operator(task_arn=None) + self.datasync.execute(None) # ### Check mocks: mock_get_conn.assert_not_called() From 5e325d1c070c78908cd695f5deec4b92f456e418 Mon Sep 17 00:00:00 2001 From: avcu Date: Thu, 23 Jul 2026 18:38:41 -0400 Subject: [PATCH 2/2] Update validate operators init exemptions to remove DataSyncOperator --- scripts/ci/prek/validate_operators_init_exemptions.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt b/scripts/ci/prek/validate_operators_init_exemptions.txt index 01fd5bc56dbb7..22a82b96ff594 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -9,7 +9,6 @@ providers/amazon/src/airflow/providers/amazon/aws/operators/appflow.py::AppflowBaseOperator providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py::BedrockCreateKnowledgeBaseOperator providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py::BedrockRaGOperator -providers/amazon/src/airflow/providers/amazon/aws/operators/datasync.py::DataSyncOperator providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py::DmsModifyTaskOperator providers/amazon/src/airflow/providers/amazon/aws/operators/dms.py::DmsStartReplicationOperator providers/amazon/src/airflow/providers/amazon/aws/operators/ecs.py::EcsRunTaskOperator