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 2671465495a Migrate EmrServerlessStartJobTrigger from CancelledError
to on_kill() (#65740)
2671465495a is described below
commit 2671465495a5165e183e046d658548d948a7e29f
Author: Mohammed Riyaan <[email protected]>
AuthorDate: Tue Apr 28 04:51:34 2026 +0530
Migrate EmrServerlessStartJobTrigger from CancelledError to on_kill()
(#65740)
Migrate EmrServerlessStartJobTrigger from the old asyncio.CancelledError
catch pattern to the new BaseTrigger.on_kill() hook (added via #65590).
Changes:
- Add async on_kill() method that cancels the EMR Serverless job
only on user-initiated task kills
- Remove the CancelledError handler from run(), letting the parent
class AwsBaseWaiterTrigger.run() handle the standard waiter flow
- Remove safe_to_cancel(), get_task_instance(), get_task_state()
which are no longer needed (on_kill only fires on user kills)
- Clean up unused imports (asyncio, AIRFLOW_V_3_0_PLUS,
TaskInstanceState, etc.)
- Update tests to directly verify on_kill() behavior
Closes: #65733 (EmrServerlessStartJobTrigger part)
---
.../airflow/providers/amazon/aws/triggers/emr.py | 23 +++++++++
.../tests/unit/amazon/aws/triggers/test_emr.py | 58 ++++++++++++++++++++++
2 files changed, 81 insertions(+)
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py
index 5fb5759a9de..b995192e4ee 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py
@@ -470,6 +470,10 @@ class EmrServerlessStartJobTrigger(AwsBaseWaiterTrigger):
)
yield TriggerEvent({"status": "success", self.return_key:
self.return_value})
except asyncio.CancelledError:
+ # TODO: Remove this CancelledError handler once the minimum
Airflow version
+ # supported by the provider is >= 3.3. In Airflow 3.3+,
BaseTrigger.on_kill()
+ # handles cancellation instead. Keeping this for backward
compatibility with
+ # older Airflow versions.
if self.job_id and self.cancel_on_kill and await
self.safe_to_cancel():
self.log.info(
"Task was cancelled. Cancelling EMR Serverless job.
Application ID: %s, Job ID: %s",
@@ -489,6 +493,25 @@ class EmrServerlessStartJobTrigger(AwsBaseWaiterTrigger):
except Exception as e:
yield TriggerEvent({"status": "failure", "message": str(e)})
+ async def on_kill(self) -> None:
+ """
+ Cancel the EMR Serverless job when the trigger is cancelled by a user
action.
+
+ This hook is available in Airflow 3.3+ via BaseTrigger.on_kill().
+ For older Airflow versions, the CancelledError handler in run()
provides
+ the same cancellation behavior.
+ """
+ if self.job_id and self.cancel_on_kill:
+ self.log.info(
+ "Cancelling EMR Serverless job. Application ID: %s, Job ID:
%s",
+ self.application_id,
+ self.job_id,
+ )
+ await sync_to_async(self.hook().conn.cancel_job_run)(
+ applicationId=self.application_id, jobRunId=self.job_id
+ )
+ self.log.info("EMR Serverless job %s cancelled successfully.",
self.job_id)
+
class EmrServerlessDeleteApplicationTrigger(AwsBaseWaiterTrigger):
"""
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py
index fef643dc927..002c5132aad 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py
@@ -403,6 +403,64 @@ class TestEmrServerlessStartJobTrigger:
mock_hook.conn.cancel_job_run.assert_not_called()
+ @pytest.mark.asyncio
+ async def test_on_kill_cancels_job(self):
+ """Test that on_kill() cancels the EMR Serverless job via the hook."""
+ trigger = EmrServerlessStartJobTrigger(
+ application_id="test_app",
+ job_id="test_job",
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ aws_conn_id="aws_default",
+ cancel_on_kill=True,
+ )
+
+ mock_hook = mock.MagicMock()
+ mock_hook.conn.cancel_job_run.return_value = {"ResponseMetadata":
{"HTTPStatusCode": 200}}
+
+ with mock.patch.object(trigger, "hook", return_value=mock_hook):
+ await trigger.on_kill()
+
+
mock_hook.conn.cancel_job_run.assert_called_once_with(applicationId="test_app",
jobRunId="test_job")
+
+ @pytest.mark.asyncio
+ async def test_on_kill_skips_when_cancel_disabled(self):
+ """Test that on_kill() does nothing when cancel_on_kill=False."""
+ trigger = EmrServerlessStartJobTrigger(
+ application_id="test_app",
+ job_id="test_job",
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ aws_conn_id="aws_default",
+ cancel_on_kill=False,
+ )
+
+ mock_hook = mock.MagicMock()
+
+ with mock.patch.object(trigger, "hook", return_value=mock_hook):
+ await trigger.on_kill()
+
+ mock_hook.conn.cancel_job_run.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_on_kill_skips_when_no_job_id(self):
+ """Test that on_kill() does nothing when job_id is None."""
+ trigger = EmrServerlessStartJobTrigger(
+ application_id="test_app",
+ job_id=None,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ aws_conn_id="aws_default",
+ cancel_on_kill=True,
+ )
+
+ mock_hook = mock.MagicMock()
+
+ with mock.patch.object(trigger, "hook", return_value=mock_hook):
+ await trigger.on_kill()
+
+ mock_hook.conn.cancel_job_run.assert_not_called()
+
class TestEmrServerlessDeleteApplicationTrigger:
def test_serialization(self):
