diff --git a/providers/amazon/docs/operators/emr/emr_serverless.rst b/providers/amazon/docs/operators/emr/emr_serverless.rst index 004c825f8e242..a5e379c8970f2 100644 --- a/providers/amazon/docs/operators/emr/emr_serverless.rst +++ b/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 --------- diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/emr.py index 87c931d134f40..9f134211c4a11 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/emr.py +++ b/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 ( diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py index 4537bf23fd1ff..017c00d9e65ae 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py +++ b/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 diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py index 8c557f494b1ff..24fecaf8bce0c 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py +++ b/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, + ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/waiters/emr-serverless.json b/providers/amazon/src/airflow/providers/amazon/aws/waiters/emr-serverless.json index 4066109382a6a..ceaa231fc6848 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/waiters/emr-serverless.json +++ b/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" + } + ] } } } diff --git a/providers/amazon/tests/system/amazon/aws/example_emr_serverless_session.py b/providers/amazon/tests/system/amazon/aws/example_emr_serverless_session.py new file mode 100644 index 0000000000000..ff55ad40a635a --- /dev/null +++ b/providers/amazon/tests/system/amazon/aws/example_emr_serverless_session.py @@ -0,0 +1,103 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from datetime import datetime + +from airflow.providers.amazon.aws.operators.emr import ( + EmrServerlessCreateApplicationOperator, + EmrServerlessDeleteApplicationOperator, + EmrServerlessStartSessionOperator, + EmrServerlessStopApplicationOperator, +) +from airflow.providers.common.compat.sdk import DAG, TriggerRule, chain + +from system.amazon.aws.utils import SystemTestContextBuilder + +DAG_ID = "example_emr_serverless_session" + +# Externally fetched variables: +ROLE_ARN_KEY = "ROLE_ARN" + +sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build() + +with DAG( + dag_id=DAG_ID, + schedule="@once", + start_date=datetime(2021, 1, 1), + catchup=False, + tags=["example", "emr-serverless"], +) as dag: + test_context = sys_test_context_task() + role_arn = test_context[ROLE_ARN_KEY] + + create_app = EmrServerlessCreateApplicationOperator( + task_id="create_app", + release_label="emr-7.13.0", + job_type="SPARK", + config={ + "name": "session-systest", + # Interactive sessions (required for the StartSession API) are only available + # on emr-7.13.0+ and must be explicitly enabled on the application. + "interactiveConfiguration": {"sessionEnabled": True}, + }, + ) + application_id = create_app.output + + # [START howto_operator_emr_serverless_start_session] + start_session = EmrServerlessStartSessionOperator( + task_id="start_session", + application_id=application_id, + execution_role_arn=role_arn, + idle_timeout_minutes=5, + ) + # [END howto_operator_emr_serverless_start_session] + + stop_app = EmrServerlessStopApplicationOperator( + task_id="stop_app", + application_id=application_id, + force_stop=True, + trigger_rule=TriggerRule.ALL_DONE, + ) + + delete_app = EmrServerlessDeleteApplicationOperator( + task_id="delete_app", + application_id=application_id, + trigger_rule=TriggerRule.ALL_DONE, + ) + + chain( + # TEST SETUP + test_context, + create_app, + # TEST BODY + start_session, + # TEST TEARDOWN + stop_app, + delete_app, + ) + + from tests_common.test_utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_emr_serverless.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_emr_serverless.py index 75e4c5f897425..d6bf0d0794de9 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_emr_serverless.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_emr_serverless.py @@ -18,6 +18,8 @@ from unittest.mock import MagicMock, PropertyMock, patch +import pytest + from airflow.providers.amazon.aws.hooks.emr import EmrServerlessHook task_id = "test_emr_serverless_create_application_operator" @@ -75,3 +77,58 @@ def test_cancel_jobs_but_no_jobs(self, conn_mock: MagicMock): # nothing very interesting should happen conn_mock.assert_called_once() + + +class TestEmrServerlessHookSession: + @pytest.fixture(autouse=True) + def _supported_botocore(self): + # The session methods gate on botocore >= 1.43.0. Pin a supported version so these + # tests exercise the boto calls regardless of the botocore installed in CI. + with patch( + "airflow.providers.amazon.aws.hooks.emr.get_botocore_version", + return_value=(1, 43, 0), + ): + yield + + @patch.object(EmrServerlessHook, "conn", new_callable=PropertyMock) + def test_start_session_minimal(self, conn_mock: MagicMock): + conn_mock().start_session.return_value = {"sessionId": "sess-1"} + hook = EmrServerlessHook(aws_conn_id="aws_default") + + session_id = hook.start_session(application_id="app", execution_role_arn="role") + + assert session_id == "sess-1" + conn_mock().start_session.assert_called_once_with(applicationId="app", executionRoleArn="role") + + @patch.object(EmrServerlessHook, "conn", new_callable=PropertyMock) + def test_start_session_with_optional_params(self, conn_mock: MagicMock): + conn_mock().start_session.return_value = {"sessionId": "sess-2"} + hook = EmrServerlessHook(aws_conn_id="aws_default") + + session_id = hook.start_session( + application_id="app", + execution_role_arn="role", + name="my-session", + idle_timeout_minutes=15, + configuration_overrides={"applicationConfiguration": []}, + ) + + assert session_id == "sess-2" + conn_mock().start_session.assert_called_once_with( + applicationId="app", + executionRoleArn="role", + name="my-session", + idleTimeoutMinutes=15, + configurationOverrides={"applicationConfiguration": []}, + ) + + @patch.object(EmrServerlessHook, "conn", new_callable=PropertyMock) + def test_start_session_gates_on_old_botocore(self, conn_mock: MagicMock): + hook = EmrServerlessHook(aws_conn_id="aws_default") + with patch( + "airflow.providers.amazon.aws.hooks.emr.get_botocore_version", + return_value=(1, 41, 0), + ): + with pytest.raises(RuntimeError, match="botocore >= 1.43.0"): + hook.start_session(application_id="app", execution_role_arn="role") + conn_mock().start_session.assert_not_called() diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless_session.py b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless_session.py new file mode 100644 index 0000000000000..0a81beb2840c9 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_serverless_session.py @@ -0,0 +1,137 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.exceptions import TaskDeferred +from airflow.providers.amazon.aws.hooks.emr import EmrServerlessHook +from airflow.providers.amazon.aws.operators.emr import EmrServerlessStartSessionOperator +from airflow.providers.amazon.aws.triggers.emr import EmrServerlessSessionTrigger + +APP_ID = "app-123" +SESSION_ID = "sess-abc" +ROLE = "arn:aws:iam::111122223333:role/emr-exec" +REGION_NAME = "eu-west-1" +VERIFY = "/path/to/ca.pem" +BOTOCORE_CONFIG = {"retries": {"max_attempts": 7}} +WAIT = "airflow.providers.amazon.aws.operators.emr.wait" + + +class TestEmrServerlessStartSessionOperator: + @mock.patch(WAIT) + @mock.patch.object(EmrServerlessHook, "get_waiter") + @mock.patch.object(EmrServerlessHook, "start_session") + def test_start_and_wait(self, start_session, get_waiter, wait_mock): + start_session.return_value = SESSION_ID + op = EmrServerlessStartSessionOperator( + task_id="start", + application_id=APP_ID, + execution_role_arn=ROLE, + idle_timeout_minutes=15, + ) + result = op.execute({}) + + start_session.assert_called_once_with( + application_id=APP_ID, + execution_role_arn=ROLE, + name=None, + idle_timeout_minutes=15, + configuration_overrides=None, + ) + wait_mock.assert_called_once() + get_waiter.assert_called_once_with("serverless_session_ready") + assert result == {"application_id": APP_ID, "session_id": SESSION_ID} + + @mock.patch(WAIT) + @mock.patch.object(EmrServerlessHook, "get_waiter") + @mock.patch.object(EmrServerlessHook, "start_session") + def test_no_wait(self, start_session, get_waiter, wait_mock): + start_session.return_value = SESSION_ID + op = EmrServerlessStartSessionOperator( + task_id="start", + application_id=APP_ID, + execution_role_arn=ROLE, + wait_for_completion=False, + ) + op.execute({}) + wait_mock.assert_not_called() + + @mock.patch.object(EmrServerlessHook, "get_waiter") + @mock.patch.object(EmrServerlessHook, "start_session") + def test_deferrable_defers(self, start_session, get_waiter): + start_session.return_value = SESSION_ID + op = EmrServerlessStartSessionOperator( + task_id="start", + application_id=APP_ID, + execution_role_arn=ROLE, + deferrable=True, + region_name=REGION_NAME, + verify=VERIFY, + botocore_config=BOTOCORE_CONFIG, + ) + with pytest.raises(TaskDeferred) as deferred: + op.execute({}) + + trigger = deferred.value.trigger + assert isinstance(trigger, EmrServerlessSessionTrigger) + assert trigger.return_key == "session_details" + assert trigger.return_value == {"application_id": APP_ID, "session_id": SESSION_ID} + assert trigger.region_name == REGION_NAME + assert trigger.verify == VERIFY + assert trigger.botocore_config == BOTOCORE_CONFIG + get_waiter.assert_not_called() + + @mock.patch.object(EmrServerlessHook, "get_waiter") + @mock.patch.object(EmrServerlessHook, "start_session") + def test_wait_for_completion_false_does_not_defer(self, start_session, get_waiter): + start_session.return_value = SESSION_ID + op = EmrServerlessStartSessionOperator( + task_id="start", + application_id=APP_ID, + execution_role_arn=ROLE, + wait_for_completion=False, + deferrable=True, + ) + + result = op.execute({}) + + assert op.wait_for_completion is False + assert result == {"application_id": APP_ID, "session_id": SESSION_ID} + get_waiter.assert_not_called() + + def test_execute_complete_success_uses_only_event_values(self): + op = EmrServerlessStartSessionOperator( + task_id="start", application_id="different-app", execution_role_arn=ROLE + ) + session_details = {"application_id": APP_ID, "session_id": SESSION_ID} + + result = op.execute_complete( + {}, + {"status": "success", "session_details": session_details}, + ) + + assert result == session_details + + def test_execute_complete_failure_raises(self): + op = EmrServerlessStartSessionOperator( + task_id="start", application_id=APP_ID, execution_role_arn=ROLE + ) + with pytest.raises(RuntimeError): + op.execute_complete({}, {"status": "failure", "session_id": SESSION_ID}) diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py b/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py index 4e497fc52d592..0de41a78fb7e4 100644 --- a/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py +++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py @@ -31,6 +31,7 @@ EmrServerlessCreateApplicationTrigger, EmrServerlessDeleteApplicationTrigger, EmrServerlessJobSensorTrigger, + EmrServerlessSessionTrigger, EmrServerlessStartApplicationTrigger, EmrServerlessStartJobTrigger, EmrServerlessStopApplicationTrigger, @@ -652,3 +653,62 @@ def test_serialization(self): "waiter_max_attempts": 60, "aws_conn_id": "aws_default", } + + +class TestEmrServerlessSessionTrigger: + def test_serialization(self): + trigger = EmrServerlessSessionTrigger( + application_id="test_application_id", + session_id="test_session_id", + waiter_delay=10, + waiter_max_attempts=60, + aws_conn_id="aws_default", + ) + classpath, kwargs = trigger.serialize() + assert classpath == "airflow.providers.amazon.aws.triggers.emr.EmrServerlessSessionTrigger" + assert kwargs == { + "application_id": "test_application_id", + "session_id": "test_session_id", + "waiter_delay": 10, + "waiter_max_attempts": 60, + "aws_conn_id": "aws_default", + } + + def test_serialization_with_hook_configuration(self): + trigger = EmrServerlessSessionTrigger( + application_id="test_application_id", + session_id="test_session_id", + region_name="eu-west-1", + verify="/path/to/ca.pem", + botocore_config={"retries": {"max_attempts": 7}}, + ) + + _, kwargs = trigger.serialize() + + assert kwargs["region_name"] == "eu-west-1" + assert kwargs["verify"] == "/path/to/ca.pem" + assert kwargs["botocore_config"] == {"retries": {"max_attempts": 7}} + + def test_hook_class(self): + assert EmrServerlessSessionTrigger.aws_hook_class is EmrServerlessHook + + def test_hook_receives_configuration(self): + trigger = EmrServerlessSessionTrigger( + application_id="test_application_id", + session_id="test_session_id", + aws_conn_id="test_conn", + region_name="eu-west-1", + verify="/path/to/ca.pem", + botocore_config={"retries": {"max_attempts": 7}}, + ) + + with mock.patch.object(EmrServerlessSessionTrigger, "aws_hook_class") as hook_class: + hook = trigger.hook() + + assert hook is hook_class.return_value + hook_class.assert_called_once_with( + aws_conn_id="test_conn", + region_name="eu-west-1", + verify="/path/to/ca.pem", + config={"retries": {"max_attempts": 7}}, + )