Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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",
)
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -1787,6 +1790,8 @@ class EmrServerlessDeleteApplicationOperator(EmrServerlessStopApplicationOperato
"application_id",
)

stop_complete_method_name: str = "stop_complete"

def __init__(
self,
application_id: str,
Expand Down Expand Up @@ -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)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down