vincbeck commented on code in PR #70763:
URL: https://github.com/apache/airflow/pull/70763#discussion_r3750777502
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py:
##########
@@ -1857,3 +1858,110 @@ def execute_complete(self, context: Context, event:
dict[str, Any] | None = None
if validated_event["status"] != "success":
raise AirflowException(f"Error deleting EMR Serverless
application: {validated_event}")
self.log.info("EMR serverless application %s deleted successfully",
self.application_id)
+
+
+class EmrServerlessStartSessionOperator(AwsBaseOperator[EmrServerlessHook]):
+ """
+ Start an EMR Serverless interactive session and wait until it is ready.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:EmrServerlessStartSessionOperator`
+
+ :param application_id: ID of the EMR Serverless application to run the
session on.
+ :param execution_role_arn: ARN of the IAM role the session assumes to
access data.
+ :param name: An optional name for the session.
+ :param idle_timeout_minutes: Auto-stop the session after this many idle
minutes.
+ :param configuration_overrides: Optional Spark/monitoring configuration
overrides.
+ :param wait_for_completion: If True, wait for the session to be ready
before returning.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used.
If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default
boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+ :param waiter_max_attempts: Number of times the waiter should poll the
session to check the state.
+ :param waiter_delay: Number of seconds between polling the state of the
session.
+ :param deferrable: If True, the operator will wait asynchronously for the
session to be ready.
+ This implies waiting for completion. This mode requires aiobotocore
module to be installed.
+ (default: False, but can be overridden in config file by setting
default_deferrable to True)
+ """
+
+ aws_hook_class = EmrServerlessHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "application_id",
+ "execution_role_arn",
+ "name",
+ "idle_timeout_minutes",
+ "configuration_overrides",
+ )
+
+ def __init__(
+ self,
+ *,
+ application_id: str,
+ execution_role_arn: str,
+ name: str | None = None,
+ idle_timeout_minutes: int | None = None,
+ configuration_overrides: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 10,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.application_id = application_id
+ self.execution_role_arn = execution_role_arn
+ self.name = name
+ self.idle_timeout_minutes = idle_timeout_minutes
+ self.configuration_overrides = configuration_overrides
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ session_id = self.hook.start_session(
+ application_id=self.application_id,
+ execution_role_arn=self.execution_role_arn,
+ name=self.name,
+ idle_timeout_minutes=self.idle_timeout_minutes,
+ configuration_overrides=self.configuration_overrides,
+ )
+ self.log.info("Started EMR Serverless session %s", session_id)
+
+ if self.deferrable:
+ self.defer(
+ trigger=EmrServerlessSessionTrigger(
+ application_id=self.application_id,
+ session_id=session_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ timeout=timedelta(seconds=self.waiter_max_attempts *
self.waiter_delay),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ wait(
+ waiter=self.hook.get_waiter("serverless_session_ready"),
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ args={"applicationId": self.application_id, "sessionId":
session_id},
+ failure_message="EMR Serverless session failed to start",
+ status_message="EMR Serverless session status is",
+ status_args=["session.state", "session.stateDetails"],
+ )
+ return {"application_id": self.application_id, "session_id":
session_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> dict:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event["status"] != "success":
+ raise RuntimeError(f"Error starting EMR Serverless session:
{validated_event}")
+ self.log.info("EMR Serverless session %s started",
validated_event["session_id"])
+ return {"application_id": self.application_id, "session_id":
validated_event["session_id"]}
Review Comment:
`execute_complete` should use only values from the event. Using
`self.application_id` in `execute_complete` is wrong. You cannot assume thr
same machine/worker will run `execute` and `execute_complete`. Therefore, if
you init `EmrServerlessStartSessionOperator` with `application_id=random()`,
this will fail
--
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]
