This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new cae1728f127 Add EmrServerlessStopSessionOperator to Amazon provider
(#74870)
cae1728f127 is described below
commit cae1728f1274e6a89bede3ad534b1ee1f8ea4995
Author: Vincent Gromakowski <[email protected]>
AuthorDate: Mon Oct 19 16:45:29 2026 +0200
Add EmrServerlessStopSessionOperator to Amazon provider (#74870)
Dags that start an EMR Serverless interactive session need a way to end
it as soon as their work is done rather than relying on the session's
idle timeout.
Stopping follows the same rules as starting: completion waiting decides
whether the task waits, deferrable mode only picks how it waits, and the
triggerer polls with the task's region, SSL verification and botocore
settings.
Co-authored-by: vgkowski <[email protected]>
---
.../amazon/docs/operators/emr/emr_serverless.rst | 14 +++
.../airflow/providers/amazon/aws/operators/emr.py | 92 +++++++++++++++++
.../airflow/providers/amazon/aws/triggers/emr.py | 48 +++++++++
.../amazon/aws/waiters/emr-serverless.json | 19 ++++
.../amazon/aws/example_emr_serverless_session.py | 10 ++
.../aws/operators/test_emr_serverless_session.py | 115 +++++++++++++++++++--
.../tests/unit/amazon/aws/triggers/test_emr.py | 60 +++++++++++
7 files changed, 347 insertions(+), 11 deletions(-)
diff --git a/providers/amazon/docs/operators/emr/emr_serverless.rst
b/providers/amazon/docs/operators/emr/emr_serverless.rst
index 7a7a9d749e5..11070ac56dc 100644
--- a/providers/amazon/docs/operators/emr/emr_serverless.rst
+++ b/providers/amazon/docs/operators/emr/emr_serverless.rst
@@ -193,6 +193,20 @@ and anyone who can read the task's XCom can use it until
it expires.
:start-after: [START howto_operator_emr_serverless_get_session_endpoint]
:end-before: [END howto_operator_emr_serverless_get_session_endpoint]
+.. _howto/operator:EmrServerlessStopSessionOperator:
+
+Stop an EMR Serverless interactive session
+==========================================
+
+To terminate an interactive session, use
+:class:`~airflow.providers.amazon.aws.operators.emr.EmrServerlessStopSessionOperator`.
+
+.. exampleinclude::
/../../amazon/tests/system/amazon/aws/example_emr_serverless_session.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_emr_serverless_stop_session]
+ :end-before: [END howto_operator_emr_serverless_stop_session]
+
Reference
---------
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 d128ec5b7dd..5c2356db030 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py
@@ -49,6 +49,7 @@ from airflow.providers.amazon.aws.triggers.emr import (
EmrServerlessStartApplicationTrigger,
EmrServerlessStartJobTrigger,
EmrServerlessStopApplicationTrigger,
+ EmrServerlessStopSessionTrigger,
EmrTerminateJobFlowTrigger,
)
from airflow.providers.amazon.aws.utils import validate_execute_complete_event
@@ -2080,3 +2081,94 @@ class
EmrServerlessGetSessionEndpointOperator(AwsBaseOperator[EmrServerlessHook]
"auth_token": auth_token,
"auth_token_expires_at": response.get("authTokenExpiresAt"),
}
+
+
+class EmrServerlessStopSessionOperator(AwsBaseOperator[EmrServerlessHook]):
+ """
+ Terminate an EMR Serverless interactive session.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:EmrServerlessStopSessionOperator`
+
+ :param application_id: ID of the EMR Serverless application.
+ :param session_id: ID of the interactive session to terminate.
+ :param wait_for_completion: If True, wait for the session to terminate
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 terminate. 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",
+ "session_id",
+ )
+
+ def __init__(
+ self,
+ *,
+ application_id: str,
+ session_id: str,
+ 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.session_id = session_id
+ 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) -> None:
+ self.hook._check_interactive_session_support()
+ self.log.info("Terminating EMR Serverless session %s", self.session_id)
+ self.hook.conn.terminate_session(applicationId=self.application_id,
sessionId=self.session_id)
+
+ if self.wait_for_completion:
+ if self.deferrable:
+ self.defer(
+ trigger=EmrServerlessStopSessionTrigger(
+ application_id=self.application_id,
+ session_id=self.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_terminated"),
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ args={"applicationId": self.application_id, "sessionId":
self.session_id},
+ failure_message="EMR Serverless session failed to
terminate",
+ status_message="EMR Serverless session status is",
+ status_args=["session.state", "session.stateDetails"],
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event["status"] != "success":
+ raise RuntimeError(f"Error terminating EMR Serverless session:
{validated_event}")
+ self.log.info("EMR Serverless session %s terminated",
validated_event["session_id"])
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 24fecaf8bce..387a315f11a 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py
@@ -861,3 +861,51 @@ class EmrServerlessSessionTrigger(AwsBaseWaiterTrigger):
verify=verify,
botocore_config=botocore_config,
)
+
+
+class EmrServerlessStopSessionTrigger(AwsBaseWaiterTrigger):
+ """
+ Poll an EMR Serverless interactive session until it is terminated.
+
+ :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_terminated",
+ waiter_args={"applicationId": application_id, "sessionId":
session_id},
+ failure_message="EMR Serverless session failed to terminate",
+ status_message="EMR Serverless session status is",
+ status_queries=["session.state"],
+ return_key="session_id",
+ return_value=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 ceaa231fc68..64ef8306519 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
@@ -189,6 +189,25 @@
"state": "failure"
}
]
+ },
+ "serverless_session_terminated": {
+ "operation": "GetSession",
+ "delay": 10,
+ "maxAttempts": 60,
+ "acceptors": [
+ {
+ "matcher": "path",
+ "argument": "session.state",
+ "expected": "TERMINATED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "session.state",
+ "expected": "FAILED",
+ "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
index b2b247df618..a24786a5cf4 100644
--- a/providers/amazon/tests/system/amazon/aws/example_emr_serverless_session.py
+++ b/providers/amazon/tests/system/amazon/aws/example_emr_serverless_session.py
@@ -24,6 +24,7 @@ from airflow.providers.amazon.aws.operators.emr import (
EmrServerlessGetSessionEndpointOperator,
EmrServerlessStartSessionOperator,
EmrServerlessStopApplicationOperator,
+ EmrServerlessStopSessionOperator,
)
from airflow.providers.common.compat.sdk import DAG, TriggerRule, chain
@@ -77,6 +78,14 @@ with DAG(
)
# [END howto_operator_emr_serverless_get_session_endpoint]
+ # [START howto_operator_emr_serverless_stop_session]
+ stop_session = EmrServerlessStopSessionOperator(
+ task_id="stop_session",
+ application_id=application_id,
+ session_id=session_id,
+ )
+ # [END howto_operator_emr_serverless_stop_session]
+
stop_app = EmrServerlessStopApplicationOperator(
task_id="stop_app",
application_id=application_id,
@@ -97,6 +106,7 @@ with DAG(
# TEST BODY
start_session,
get_endpoint,
+ stop_session,
# TEST TEARDOWN
stop_app,
delete_app,
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
index 50606e97ada..6e29371f6b1 100644
---
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
@@ -25,8 +25,12 @@ from airflow.providers.amazon.aws.hooks.emr import
EmrServerlessHook
from airflow.providers.amazon.aws.operators.emr import (
EmrServerlessGetSessionEndpointOperator,
EmrServerlessStartSessionOperator,
+ EmrServerlessStopSessionOperator,
+)
+from airflow.providers.amazon.aws.triggers.emr import (
+ EmrServerlessSessionTrigger,
+ EmrServerlessStopSessionTrigger,
)
-from airflow.providers.amazon.aws.triggers.emr import
EmrServerlessSessionTrigger
APP_ID = "app-123"
SESSION_ID = "sess-abc"
@@ -37,6 +41,17 @@ BOTOCORE_CONFIG = {"retries": {"max_attempts": 7}}
WAIT = "airflow.providers.amazon.aws.operators.emr.wait"
[email protected]
+def supported_botocore():
+ # The session operators 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 mock.patch(
+ "airflow.providers.amazon.aws.hooks.emr.get_botocore_version",
+ return_value=(1, 43, 0),
+ ):
+ yield
+
+
class TestEmrServerlessStartSessionOperator:
@mock.patch(WAIT)
@mock.patch.object(EmrServerlessHook, "get_waiter")
@@ -140,6 +155,7 @@ class TestEmrServerlessStartSessionOperator:
op.execute_complete({}, {"status": "failure", "session_id":
SESSION_ID})
[email protected]("supported_botocore")
class TestEmrServerlessGetSessionEndpointOperator:
RAW_RESPONSE = {
"applicationId": APP_ID,
@@ -149,16 +165,6 @@ class TestEmrServerlessGetSessionEndpointOperator:
"authTokenExpiresAt": None,
}
- @pytest.fixture(autouse=True)
- def _supported_botocore(self):
- # The operator gates on botocore >= 1.43.0. Pin a supported version so
these
- # tests exercise the boto call regardless of the botocore installed in
CI.
- with mock.patch(
- "airflow.providers.amazon.aws.hooks.emr.get_botocore_version",
- return_value=(1, 43, 0),
- ):
- yield
-
@mock.patch.object(EmrServerlessHook, "conn")
def test_returns_transformed_endpoint(self, mock_conn):
mock_conn.get_session_endpoint.return_value = self.RAW_RESPONSE
@@ -195,3 +201,90 @@ class TestEmrServerlessGetSessionEndpointOperator:
with pytest.raises(RuntimeError, match="botocore >= 1.43.0"):
op.execute({})
mock_conn.get_session_endpoint.assert_not_called()
+
+
[email protected]("supported_botocore")
+class TestEmrServerlessStopSessionOperator:
+ @mock.patch(WAIT)
+ @mock.patch.object(EmrServerlessHook, "get_waiter")
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_terminate_and_wait(self, mock_conn, get_waiter, wait_mock):
+ op = EmrServerlessStopSessionOperator(task_id="stop",
application_id=APP_ID, session_id=SESSION_ID)
+ op.execute({})
+
+
mock_conn.terminate_session.assert_called_once_with(applicationId=APP_ID,
sessionId=SESSION_ID)
+ wait_mock.assert_called_once()
+ get_waiter.assert_called_once_with("serverless_session_terminated")
+
+ @mock.patch(WAIT)
+ @mock.patch.object(EmrServerlessHook, "get_waiter")
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_terminate_no_wait(self, mock_conn, get_waiter, wait_mock):
+ op = EmrServerlessStopSessionOperator(
+ task_id="stop",
+ application_id=APP_ID,
+ session_id=SESSION_ID,
+ wait_for_completion=False,
+ )
+ op.execute({})
+ wait_mock.assert_not_called()
+
+ @mock.patch(WAIT)
+ @mock.patch.object(EmrServerlessHook, "get_waiter")
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_deferrable_defers(self, mock_conn, get_waiter, wait_mock):
+ op = EmrServerlessStopSessionOperator(
+ task_id="stop",
+ application_id=APP_ID,
+ session_id=SESSION_ID,
+ deferrable=True,
+ region_name=REGION_NAME,
+ verify=VERIFY,
+ botocore_config=BOTOCORE_CONFIG,
+ )
+ with pytest.raises(TaskDeferred) as deferred:
+ op.execute({})
+
+
mock_conn.terminate_session.assert_called_once_with(applicationId=APP_ID,
sessionId=SESSION_ID)
+ trigger = deferred.value.trigger
+ assert isinstance(trigger, EmrServerlessStopSessionTrigger)
+ assert trigger.region_name == REGION_NAME
+ assert trigger.verify == VERIFY
+ assert trigger.botocore_config == BOTOCORE_CONFIG
+ wait_mock.assert_not_called()
+
+ @mock.patch(WAIT)
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_wait_for_completion_false_does_not_defer(self, mock_conn,
wait_mock):
+ op = EmrServerlessStopSessionOperator(
+ task_id="stop",
+ application_id=APP_ID,
+ session_id=SESSION_ID,
+ wait_for_completion=False,
+ deferrable=True,
+ )
+
+ op.execute({})
+
+
mock_conn.terminate_session.assert_called_once_with(applicationId=APP_ID,
sessionId=SESSION_ID)
+ wait_mock.assert_not_called()
+
+ @mock.patch.object(EmrServerlessHook, "conn")
+ def test_gates_on_old_botocore(self, mock_conn):
+ op = EmrServerlessStopSessionOperator(task_id="stop",
application_id=APP_ID, session_id=SESSION_ID)
+ with mock.patch(
+ "airflow.providers.amazon.aws.hooks.emr.get_botocore_version",
+ return_value=(1, 41, 0),
+ ):
+ with pytest.raises(RuntimeError, match="botocore >= 1.43.0"):
+ op.execute({})
+ mock_conn.terminate_session.assert_not_called()
+
+ def test_execute_complete_success(self):
+ op = EmrServerlessStopSessionOperator(task_id="stop",
application_id=APP_ID, session_id=SESSION_ID)
+ assert op.execute_complete({}, {"status": "success", "session_id":
SESSION_ID}) is None
+
+ def test_execute_complete_failure_raises(self):
+ op = EmrServerlessStopSessionOperator(task_id="stop",
application_id=APP_ID, session_id=SESSION_ID)
+ 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 0de41a78fb7..bb8f62a3ccb 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_emr.py
@@ -35,6 +35,7 @@ from airflow.providers.amazon.aws.triggers.emr import (
EmrServerlessStartApplicationTrigger,
EmrServerlessStartJobTrigger,
EmrServerlessStopApplicationTrigger,
+ EmrServerlessStopSessionTrigger,
EmrStepSensorTrigger,
EmrTerminateJobFlowTrigger,
)
@@ -712,3 +713,62 @@ class TestEmrServerlessSessionTrigger:
verify="/path/to/ca.pem",
config={"retries": {"max_attempts": 7}},
)
+
+
+class TestEmrServerlessStopSessionTrigger:
+ def test_serialization(self):
+ trigger = EmrServerlessStopSessionTrigger(
+ 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.EmrServerlessStopSessionTrigger"
+ 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 = EmrServerlessStopSessionTrigger(
+ 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 EmrServerlessStopSessionTrigger.aws_hook_class is
EmrServerlessHook
+
+ def test_hook_receives_configuration(self):
+ trigger = EmrServerlessStopSessionTrigger(
+ 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(EmrServerlessStopSessionTrigger,
"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}},
+ )
