SameerMesiah97 commented on code in PR #73719:
URL: https://github.com/apache/airflow/pull/73719#discussion_r4108097130
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py:
##########
@@ -1490,18 +1490,30 @@ def execute(self, context: Context, event: dict[str,
Any] | None = None) -> str
return self.job_id
- def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> str | None:
validated_event = validate_execute_complete_event(event)
+ # On resume execute() has not run, so self.job_id is unset; read the
run details the trigger
+ # sends in the event (present on both success and failure) and fall
back to self as a guard.
+ job_details = validated_event.get("job_details") or {}
+ job_id = job_details.get("job_id") or self.job_id
+ application_id = job_details.get("application_id") or
self.application_id
+
if validated_event["status"] == "success":
self.log.info("Serverless job completed")
- return validated_event["job_details"]["job_id"]
- self.log.info("Cancelling EMR Serverless job %s", self.job_id)
- self.hook.conn.cancel_job_run(
- applicationId=validated_event["job_details"]["application_id"],
- jobRunId=validated_event["job_details"]["job_id"],
+ return job_id
+ if job_id:
+ # The waiter fails on both a job failure and a poll timeout; on
timeout the run may still
+ # be active, so try to cancel it to avoid an orphan. Guard it so a
cancel error (e.g. an
+ # already-terminal run) does not mask the real failure reason
surfaced below.
+ self.log.info("Cancelling EMR Serverless job %s after failure or
timeout", job_id)
+ try:
+ self.hook.conn.cancel_job_run(applicationId=application_id,
jobRunId=job_id)
+ except Exception:
Review Comment:
Could we narrow this `Exception `to the errors that are expected to arise?
##########
providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless.py:
##########
@@ -693,6 +693,58 @@ def
test_execute_complete_deferrable_failure_triggers_cancel(
log_msgs = [call.args[0] for call in mock_log.call_args_list]
assert any("Cancelling EMR Serverless job" in msg for msg in
log_msgs)
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_execute_complete_failure_surfaces_message(self, mock_conn):
+ operator = EmrServerlessStartJobOperator(
+ task_id="test_task",
+ application_id="app-id",
+ execution_role_arn="arn",
+ job_driver={"sparkSubmit": {"entryPoint": "s3://x"}},
+ deferrable=True,
+ )
+ failed_event = {
+ "status": "failure",
+ "message": "Serverless Job failed: boom",
+ "job_details": {"application_id": "app-id", "job_id": "job-id"},
+ }
+ with pytest.raises(AirflowException, match="boom"):
+ operator.execute_complete(mock.MagicMock(), failed_event)
+
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def
test_execute_complete_failure_without_job_details_does_not_keyerror(self,
mock_conn):
+ """A trigger failure event without job_details must surface the
reason, not raise KeyError."""
+ operator = EmrServerlessStartJobOperator(
+ task_id="test_task",
+ application_id="app-id",
+ execution_role_arn="arn",
+ job_driver={"sparkSubmit": {"entryPoint": "s3://x"}},
+ deferrable=True,
+ )
+ failed_event = {"status": "failure", "message": "Serverless Job
failed: boom"}
+ with pytest.raises(AirflowException, match="boom"):
+ operator.execute_complete(mock.MagicMock(), failed_event)
+ # No run id available on resume, so nothing to cancel.
+ mock_conn.cancel_job_run.assert_not_called()
+
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_execute_complete_failure_cancel_error_does_not_mask_reason(self,
mock_conn):
+ mock_conn.cancel_job_run.side_effect = Exception("run is not in a
cancellable state")
+ operator = EmrServerlessStartJobOperator(
+ task_id="test_task",
+ application_id="app-id",
+ execution_role_arn="arn",
+ job_driver={"sparkSubmit": {"entryPoint": "s3://x"}},
+ deferrable=True,
+ )
+ failed_event = {
+ "status": "failure",
+ "message": "Serverless Job failed: boom",
+ "job_details": {"application_id": "app-id", "job_id": "job-id"},
+ }
+ with pytest.raises(AirflowException, match="boom"):
+ operator.execute_complete(mock.MagicMock(), failed_event)
+ mock_conn.cancel_job_run.assert_called_once()
Review Comment:
I would assert the cancellation arguments too. Please see the below:
```
mock_conn.cancel_job_run.assert_called_once_with(
applicationId="app-id", jobRunId="job-id"
)
```
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py:
##########
@@ -1490,18 +1490,30 @@ def execute(self, context: Context, event: dict[str,
Any] | None = None) -> str
return self.job_id
- def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> str | None:
validated_event = validate_execute_complete_event(event)
+ # On resume execute() has not run, so self.job_id is unset; read the
run details the trigger
+ # sends in the event (present on both success and failure) and fall
back to self as a guard.
Review Comment:
Comment is a bit unclear. The below is more clear:
`#After deferral, self.job_id may be unavailable. Use the IDs supplied by
the trigger event for cancellation`
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py:
##########
@@ -1490,18 +1490,30 @@ def execute(self, context: Context, event: dict[str,
Any] | None = None) -> str
return self.job_id
- def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> str | None:
validated_event = validate_execute_complete_event(event)
+ # On resume execute() has not run, so self.job_id is unset; read the
run details the trigger
+ # sends in the event (present on both success and failure) and fall
back to self as a guard.
+ job_details = validated_event.get("job_details") or {}
+ job_id = job_details.get("job_id") or self.job_id
+ application_id = job_details.get("application_id") or
self.application_id
+
if validated_event["status"] == "success":
self.log.info("Serverless job completed")
- return validated_event["job_details"]["job_id"]
- self.log.info("Cancelling EMR Serverless job %s", self.job_id)
- self.hook.conn.cancel_job_run(
- applicationId=validated_event["job_details"]["application_id"],
- jobRunId=validated_event["job_details"]["job_id"],
+ return job_id
+ if job_id:
+ # The waiter fails on both a job failure and a poll timeout; on
timeout the run may still
+ # be active, so try to cancel it to avoid an orphan. Guard it so a
cancel error (e.g. an
+ # already-terminal run) does not mask the real failure reason
surfaced below.
+ self.log.info("Cancelling EMR Serverless job %s after failure or
timeout", job_id)
+ try:
+ self.hook.conn.cancel_job_run(applicationId=application_id,
jobRunId=job_id)
+ except Exception:
+ self.log.exception("Failed to cancel EMR Serverless job %s",
job_id)
+ raise AirflowException(
+ f"EMR Serverless job failed or timed out in deferrable mode:
{validated_event.get('message', '')}"
)
Review Comment:
1) `AirflowException` is to avoided in provider code. Perhaps, you can use
`RuntimeError` instead.
2) Could we distinguish a terminal job failure from exhausted polling
attempts in the trigger event, then raise a message specific to the outcome
here? At present both arrive as `"failure"`, so “failed or timed out” leaves
users unable to tell what happened.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless.py:
##########
@@ -693,6 +693,58 @@ def
test_execute_complete_deferrable_failure_triggers_cancel(
log_msgs = [call.args[0] for call in mock_log.call_args_list]
assert any("Cancelling EMR Serverless job" in msg for msg in
log_msgs)
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_execute_complete_failure_surfaces_message(self, mock_conn):
+ operator = EmrServerlessStartJobOperator(
+ task_id="test_task",
+ application_id="app-id",
+ execution_role_arn="arn",
+ job_driver={"sparkSubmit": {"entryPoint": "s3://x"}},
+ deferrable=True,
+ )
+ failed_event = {
+ "status": "failure",
+ "message": "Serverless Job failed: boom",
+ "job_details": {"application_id": "app-id", "job_id": "job-id"},
+ }
+ with pytest.raises(AirflowException, match="boom"):
+ operator.execute_complete(mock.MagicMock(), failed_event)
+
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def
test_execute_complete_failure_without_job_details_does_not_keyerror(self,
mock_conn):
+ """A trigger failure event without job_details must surface the
reason, not raise KeyError."""
+ operator = EmrServerlessStartJobOperator(
+ task_id="test_task",
+ application_id="app-id",
+ execution_role_arn="arn",
+ job_driver={"sparkSubmit": {"entryPoint": "s3://x"}},
+ deferrable=True,
+ )
+ failed_event = {"status": "failure", "message": "Serverless Job
failed: boom"}
+ with pytest.raises(AirflowException, match="boom"):
+ operator.execute_complete(mock.MagicMock(), failed_event)
+ # No run id available on resume, so nothing to cancel.
+ mock_conn.cancel_job_run.assert_not_called()
+
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_execute_complete_failure_cancel_error_does_not_mask_reason(self,
mock_conn):
+ mock_conn.cancel_job_run.side_effect = Exception("run is not in a
cancellable state")
+ operator = EmrServerlessStartJobOperator(
+ task_id="test_task",
+ application_id="app-id",
+ execution_role_arn="arn",
+ job_driver={"sparkSubmit": {"entryPoint": "s3://x"}},
+ deferrable=True,
+ )
+ failed_event = {
+ "status": "failure",
+ "message": "Serverless Job failed: boom",
+ "job_details": {"application_id": "app-id", "job_id": "job-id"},
+ }
+ with pytest.raises(AirflowException, match="boom"):
Review Comment:
Match the full raised message, or at least the meaningful failure reason,
rather than just "boom". This feedback applies to all applicable tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
