pankajastro commented on code in PR #32029:
URL: https://github.com/apache/airflow/pull/32029#discussion_r1243467709
##########
airflow/providers/amazon/aws/operators/emr.py:
##########
@@ -556,27 +560,57 @@ def execute(self, context: Context) -> str | None:
self.client_request_token,
self.tags,
)
+ if self.deferrable:
+ query_status = self.hook.check_query_status(job_id=self.job_id)
+ self.check_failure(query_status)
+ if query_status in EmrContainerHook.SUCCESS_STATES:
+ return self.job_id
+ timeout = (
+ timedelta(seconds=self.max_polling_attempts *
self.poll_interval)
+ if self.max_polling_attempts
+ else self.execution_timeout
+ )
+ self.defer(
+ timeout=timeout,
+ trigger=EmrContainerTrigger(
+ virtual_cluster_id=self.virtual_cluster_id,
+ job_id=self.job_id,
+ aws_conn_id=self.aws_conn_id,
+ poll_interval=self.poll_interval,
+ ),
+ method_name="execute_complete",
+ )
if self.wait_for_completion:
query_status = self.hook.poll_query_status(
self.job_id,
max_polling_attempts=self.max_polling_attempts,
poll_interval=self.poll_interval,
)
- if query_status in EmrContainerHook.FAILURE_STATES:
- error_message = self.hook.get_job_failure_reason(self.job_id)
- raise AirflowException(
- f"EMR Containers job failed. Final state is
{query_status}. "
- f"query_execution_id is {self.job_id}. Error:
{error_message}"
- )
- elif not query_status or query_status in
EmrContainerHook.INTERMEDIATE_STATES:
+ self.check_failure(query_status)
+ if not query_status or query_status in
EmrContainerHook.INTERMEDIATE_STATES:
raise AirflowException(
f"Final state of EMR Containers job is {query_status}. "
f"Max tries of poll status exceeded, query_execution_id is
{self.job_id}."
)
return self.job_id
+ def check_failure(self, query_status):
+ if query_status in EmrContainerHook.FAILURE_STATES:
+ error_message = self.hook.get_job_failure_reason(self.job_id)
+ raise AirflowException(
+ f"EMR Containers job failed. Final state is {query_status}. "
+ f"query_execution_id is {self.job_id}. Error: {error_message}"
+ )
+
+ def execute_complete(self, context, event=None):
+ if event["status"] != "success":
+ raise AirflowException(f"Error while running job: {event}")
+
+ self.log.info(event["message"])
Review Comment:
Thank you, I was not aware of this.
--
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]
