diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py index aab675201497c..9524c7fc7aa4f 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py @@ -1643,6 +1643,9 @@ class EmrServerlessStopApplicationOperator(AwsBaseOperator[EmrServerlessHook]): """ aws_hook_class = EmrServerlessHook + #: Method to resume at once the application has stopped. Subclasses may override this + #: to chain further steps instead of finishing the task. + stop_complete_method_name: str = "execute_complete" template_fields: Sequence[str] = aws_template_fields( "application_id", ) @@ -1710,7 +1713,7 @@ def execute(self, context: Context) -> None: waiter_max_attempts=self.waiter_max_attempts, ), timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay), - method_name="execute_complete", + method_name=self.stop_complete_method_name, ) if self.wait_for_completion: waiter = self.hook.get_waiter("serverless_app_stopped") @@ -1740,7 +1743,7 @@ def stop_application(self, context: Context, event: dict[str, Any] | None = None waiter_max_attempts=self.waiter_max_attempts, ), timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay), - method_name="execute_complete", + method_name=self.stop_complete_method_name, ) def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None: @@ -1787,6 +1790,8 @@ class EmrServerlessDeleteApplicationOperator(EmrServerlessStopApplicationOperato "application_id", ) + stop_complete_method_name: str = "stop_complete" + def __init__( self, application_id: str, @@ -1851,6 +1856,30 @@ def execute(self, context: Context) -> None: self.log.info("EMR serverless application deleted") + def stop_complete(self, context: Context, event: dict[str, Any] | None = None) -> None: + validated_event = validate_execute_complete_event(event) + + if validated_event["status"] != "success": + raise AirflowException(f"Error stopping EMR Serverless application: {validated_event}") + + self.log.info("Now deleting application: %s", self.application_id) + response = self.hook.conn.delete_application(applicationId=self.application_id) + + if response["ResponseMetadata"]["HTTPStatusCode"] != 200: + raise AirflowException(f"Application deletion failed: {response}") + + if self.deferrable: + self.defer( + trigger=EmrServerlessDeleteApplicationTrigger( + application_id=self.application_id, + aws_conn_id=self.aws_conn_id, + waiter_delay=self.waiter_delay, + waiter_max_attempts=self.waiter_max_attempts, + ), + timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay), + method_name="execute_complete", + ) + def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None: validated_event = validate_execute_complete_event(event) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless.py b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless.py index 75fcf1bce3d24..eb3416a7b4be1 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless.py @@ -30,6 +30,10 @@ EmrServerlessStartJobOperator, EmrServerlessStopApplicationOperator, ) +from airflow.providers.amazon.aws.triggers.emr import ( + EmrServerlessDeleteApplicationTrigger, + EmrServerlessStopApplicationTrigger, +) from airflow.providers.amazon.version_compat import NOTSET from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred @@ -1400,6 +1404,38 @@ def test_delete_application_deferrable(self, mock_conn): with pytest.raises(TaskDeferred): operator.execute(None) + + @mock.patch.object(EmrServerlessHook, "conn") + def test_delete_application_deferrable_deletes_after_stop(self, mock_conn): + """Regression test for #72123: deferrable delete must call delete_application via chained deferral.""" + mock_conn.stop_application.return_value = {} + mock_conn.delete_application.return_value = {"ResponseMetadata": {"HTTPStatusCode": 200}} + + operator = EmrServerlessDeleteApplicationOperator( + task_id=task_id, + application_id=application_id, + deferrable=True, + ) + + # First defer defers on the STOP trigger, resuming at stop_complete (not execute_complete). + with pytest.raises(TaskDeferred) as exc_info: + operator.execute(None) + assert isinstance(exc_info.value.trigger, EmrServerlessStopApplicationTrigger) + assert exc_info.value.method_name == "stop_complete" + mock_conn.stop_application.assert_called_once() + mock_conn.delete_application.assert_not_called() + + # Stop trigger fires successfully -> operator now issues DeleteApplication and defers + # on the DELETE trigger, resuming at execute_complete. + with pytest.raises(TaskDeferred) as exc_info: + operator.stop_complete({}, {"status": "success"}) + assert isinstance(exc_info.value.trigger, EmrServerlessDeleteApplicationTrigger) + assert exc_info.value.method_name == "execute_complete" + mock_conn.delete_application.assert_called_once_with(applicationId=application_id) + + # Delete trigger fires successfully -> task completes. + operator.execute_complete({}, {"status": "success"}) + def test_execute_complete_error(self): operator = EmrServerlessDeleteApplicationOperator( task_id=task_id, application_id=application_id_delete_operator