Fix automatic termination issue in EmrOperator by ensuring `waiter_…

…max_attempts` is set for deferrable triggers (apache#38658)
romsharon98 · Jul 26, 2024 · d55342e · d55342e
1 parent c37ad0f
commit d55342e
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 0 deletions.
diff --git a/airflow/providers/amazon/aws/operators/emr.py b/airflow/providers/amazon/aws/operators/emr.py
@@ -617,6 +617,14 @@ def execute(self, context: Context) -> str | None:
                     job_id=self.job_id,
                     aws_conn_id=self.aws_conn_id,
                     waiter_delay=self.poll_interval,
+                    waiter_max_attempts=self.max_polling_attempts,
+                )
+                if self.max_polling_attempts
+                else EmrContainerTrigger(
+                    virtual_cluster_id=self.virtual_cluster_id,
+                    job_id=self.job_id,
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=self.poll_interval,
                 ),
                 method_name="execute_complete",
             )

diff --git a/airflow/providers/amazon/aws/sensors/emr.py b/airflow/providers/amazon/aws/sensors/emr.py
@@ -354,6 +354,14 @@ def execute(self, context: Context):
                     job_id=self.job_id,
                     aws_conn_id=self.aws_conn_id,
                     waiter_delay=self.poll_interval,
+                    waiter_max_attempts=self.max_retries,
+                )
+                if self.max_retries
+                else EmrContainerTrigger(
+                    virtual_cluster_id=self.virtual_cluster_id,
+                    job_id=self.job_id,
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=self.poll_interval,
                 ),
                 method_name="execute_complete",
             )

diff --git a/tests/providers/amazon/aws/operators/test_emr_containers.py b/tests/providers/amazon/aws/operators/test_emr_containers.py
@@ -144,6 +144,22 @@ def test_operator_defer(self, mock_submit_job, mock_check_query_status):
             exc.value.trigger, EmrContainerTrigger
         ), f"{exc.value.trigger} is not a EmrContainerTrigger"
 
+    @mock.patch.object(EmrContainerHook, "submit_job")
+    @mock.patch.object(
+        EmrContainerHook, "check_query_status", return_value=EmrContainerHook.INTERMEDIATE_STATES[0]
+    )
+    def test_operator_defer_with_timeout(self, mock_submit_job, mock_check_query_status):
+        self.emr_container.deferrable = True
+        self.emr_container.max_polling_attempts = 1000
+
+        with pytest.raises(TaskDeferred) as e:
+            self.emr_container.execute(context=None)
+
+        trigger = e.value.trigger
+        assert isinstance(trigger, EmrContainerTrigger), f"{trigger} is not a EmrContainerTrigger"
+        assert trigger.waiter_delay == self.emr_container.poll_interval
+        assert trigger.attempts == self.emr_container.max_polling_attempts
+
 
 class TestEmrEksCreateClusterOperator:
     def setup_method(self):

diff --git a/tests/providers/amazon/aws/sensors/test_emr_containers.py b/tests/providers/amazon/aws/sensors/test_emr_containers.py
@@ -84,3 +84,17 @@ def test_sensor_defer(self, mock_poke):
         assert isinstance(
             e.value.trigger, EmrContainerTrigger
         ), f"{e.value.trigger} is not a EmrContainerTrigger"
+
+    @mock.patch("airflow.providers.amazon.aws.sensors.emr.EmrContainerSensor.poke")
+    def test_sensor_defer_with_timeout(self, mock_poke):
+        self.sensor.deferrable = True
+        mock_poke.return_value = False
+        self.sensor.max_retries = 1000
+
+        with pytest.raises(TaskDeferred) as e:
+            self.sensor.execute(context=None)
+
+        trigger = e.value.trigger
+        assert isinstance(trigger, EmrContainerTrigger), f"{trigger} is not a EmrContainerTrigger"
+        assert trigger.waiter_delay == self.sensor.poll_interval
+        assert trigger.attempts == self.sensor.max_retries