Add EmrServerlessStartSessionOperator to Amazon provider

providers/amazon/docs/operators/emr/emr_serverless.rst

@@ -153,6 +153,29 @@ To monitor the state of an EMR Serverless Application you can use
    :start-after: [START howto_sensor_emr_serverless_application]
    :end-before: [END howto_sensor_emr_serverless_application]
 
+.. _howto/operator:EmrServerlessStartSessionOperator:
+
+Start an EMR Serverless interactive session
+===========================================
+
+To start an EMR Serverless interactive session that a Spark Connect client can attach to, use
+:class:`~airflow.providers.amazon.aws.operators.emr.EmrServerlessStartSessionOperator`.
+Set ``deferrable=True`` to release the worker slot while the session warms up.
+
+.. note::
+    Interactive sessions require Amazon EMR release ``emr-7.13.0`` or later, and the session APIs
+    are only available in ``botocore>=1.43.0``. Deferrable mode additionally needs
+    ``aiobotocore>=3.6.0``, the first release whose ``botocore`` pin allows 1.43.0. The Amazon
+    provider keeps a lower minimum for these libraries, so install compatible versions to use
+    interactive sessions; the operator raises a clear error at runtime if the installed
+    ``botocore`` is too old.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_emr_serverless_session.py
+   :language: python
+   :dedent: 4
+   :start-after: [START howto_operator_emr_serverless_start_session]
+   :end-before: [END howto_operator_emr_serverless_start_session]
+
 Reference
 ---------
 

providers/amazon/src/airflow/providers/amazon/aws/hooks/emr.py

@@ -27,6 +27,7 @@
 from tenacity import retry_if_exception, stop_after_attempt, wait_fixed
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.utils import get_botocore_version
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
 from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
 
@@ -263,10 +264,30 @@ class EmrServerlessHook(AwsBaseHook):
     APPLICATION_FAILURE_STATES = {"STOPPED", "TERMINATED"}
     APPLICATION_SUCCESS_STATES = {"CREATED", "STARTED"}
 
+    SESSION_INTERMEDIATE_STATES = {"SUBMITTED", "STARTING"}
+    SESSION_FAILURE_STATES = {"FAILED", "TERMINATING", "TERMINATED"}
+    SESSION_SUCCESS_STATES = {"STARTED", "IDLE"}
+
+    # botocore version that first shipped the EMR Serverless interactive session APIs.
+    # The provider keeps a lower botocore floor, so the session methods gate on this at
+    # runtime instead of forcing every user onto a newer botocore.
+    SESSION_MIN_BOTOCORE_VERSION = (1, 43, 0)
+
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         kwargs["client_type"] = "emr-serverless"
         super().__init__(*args, **kwargs)
 
+    def _check_interactive_session_support(self) -> None:
+        """Raise a clear error if the installed botocore is too old for interactive sessions."""
+        if get_botocore_version() < self.SESSION_MIN_BOTOCORE_VERSION:
+            required = ".".join(map(str, self.SESSION_MIN_BOTOCORE_VERSION))
+            installed = ".".join(map(str, get_botocore_version()))
+            raise RuntimeError(
+                f"EMR Serverless interactive sessions require botocore >= {required}, "
+                f"but botocore {installed} is installed. Upgrade botocore (and aiobotocore >= 3.6.0 "
+                "for deferrable mode) to use this feature."
+            )
+
     def cancel_running_jobs(
         self, application_id: str, waiter_config: dict | None = None, wait_for_completion: bool = True
     ) -> int:
@@ -311,6 +332,36 @@ def cancel_running_jobs(
 
         return count
 
+    def start_session(
+        self,
+        application_id: str,
+        execution_role_arn: str,
+        name: str | None = None,
+        idle_timeout_minutes: int | None = None,
+        configuration_overrides: dict | None = None,
+    ) -> str:
+        """
+        Start an EMR Serverless interactive session and return its id.
+
+        :param application_id: The id of the EMR Serverless application to run the session on.
+        :param execution_role_arn: The IAM role ARN 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.
+        """
+        self._check_interactive_session_support()
+        params: dict[str, Any] = {
+            "applicationId": application_id,
+            "executionRoleArn": execution_role_arn,
+        }
+        if name is not None:
+            params["name"] = name
+        if idle_timeout_minutes is not None:
+            params["idleTimeoutMinutes"] = idle_timeout_minutes
+        if configuration_overrides is not None:
+            params["configurationOverrides"] = configuration_overrides
+        return self.conn.start_session(**params)["sessionId"]
+
 
 def is_connection_being_updated_exception(exception: BaseException) -> bool:
     return (

providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py

@@ -45,6 +45,7 @@
     EmrServerlessCancelJobsTrigger,
     EmrServerlessCreateApplicationTrigger,
     EmrServerlessDeleteApplicationTrigger,
+    EmrServerlessSessionTrigger,
     EmrServerlessStartApplicationTrigger,
     EmrServerlessStartJobTrigger,
     EmrServerlessStopApplicationTrigger,
@@ -1910,3 +1911,114 @@ 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 and ``wait_for_completion`` is enabled, the operator will wait
+        asynchronously for the session to be ready. This mode requires aiobotocore 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.wait_for_completion:
+            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,
+                        region_name=self.region_name,
+                        verify=self.verify,
+                        botocore_config=self.botocore_config,
+                    ),
+                    timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay),
+                    method_name="execute_complete",
+                )
+            else:
+                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}")
+        session_details = validated_event["session_details"]
+        self.log.info("EMR Serverless session %s started", session_details["session_id"])
+        return session_details

providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py

@@ -813,3 +813,51 @@ def __init__(
     def hook_instance(self) -> AwsGenericHook:
         """This property is added for backward compatibility."""
         return self.hook()
+
+
+class EmrServerlessSessionTrigger(AwsBaseWaiterTrigger):
+    """
+    Poll an EMR Serverless interactive session until it reaches a ready state.
+
+    :param application_id: The ID of the EMR Serverless application.
+    :param session_id: The ID of the interactive session being polled.
+    :param waiter_delay: polling period in seconds to check for the status
+    :param waiter_max_attempts: The maximum number of attempts to be made
+    :param aws_conn_id: Reference to AWS connection id
+    :param region_name: The AWS region where the resources to watch are.
+    :param verify: Whether or not to verify SSL certificates.
+        See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = EmrServerlessHook
+
+    def __init__(
+        self,
+        *,
+        application_id: str,
+        session_id: str,
+        waiter_delay: int = 10,
+        waiter_max_attempts: int = 60,
+        aws_conn_id: str | None = "aws_default",
+        region_name: str | None = None,
+        verify: bool | str | None = None,
+        botocore_config: dict | None = None,
+    ) -> None:
+        super().__init__(
+            serialized_fields={"application_id": application_id, "session_id": session_id},
+            waiter_name="serverless_session_ready",
+            waiter_args={"applicationId": application_id, "sessionId": session_id},
+            failure_message="EMR Serverless session failed to start",
+            status_message="EMR Serverless session status is",
+            status_queries=["session.state"],
+            return_key="session_details",
+            return_value={"application_id": application_id, "session_id": session_id},
+            waiter_delay=waiter_delay,
+            waiter_max_attempts=waiter_max_attempts,
+            aws_conn_id=aws_conn_id,
+            region_name=region_name,
+            verify=verify,
+            botocore_config=botocore_config,
+        )

providers/amazon/src/airflow/providers/amazon/aws/waiters/emr-serverless.json

@@ -152,6 +152,43 @@
                     "state": "success"
                 }
             ]
+        },
+        "serverless_session_ready": {
+            "operation": "GetSession",
+            "delay": 10,
+            "maxAttempts": 60,
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "session.state",
+                    "expected": "STARTED",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "session.state",
+                    "expected": "IDLE",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "session.state",
+                    "expected": "FAILED",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "session.state",
+                    "expected": "TERMINATING",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "session.state",
+                    "expected": "TERMINATED",
+                    "state": "failure"
+                }
+            ]
         }
     }
 }