Add full support for AWS SSM Run Command in Airflow

docs/spelling_wordlist.txt

@@ -1156,6 +1156,7 @@ mssql
 mTLS
 mtls
 muldelete
+multi-cloud
 multimodal
 Multinamespace
 mutex

providers/amazon/docs/operators/ssm.rst

@@ -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.
+
+===================================
+Amazon Simple Systems Manager (SSM)
+===================================
+
+`Amazon Simple Systems Manager (Amazon SSM) <https://aws.amazon.com/systems-manager/>`__ is a service
+that helps centrally view, manage, and operate nodes at scale in AWS, on-premises, and multi-cloud
+environments. Systems Manager consolidates various tools to help complete common node tasks across AWS
+accounts and Regions.
+To use Systems Manager, nodes must be managed, which means SSM Agent is installed on the machine and
+the agent can communicate with the Systems Manager service.
+
+Prerequisite Tasks
+------------------
+
+.. include:: ../_partials/prerequisite_tasks.rst
+
+Generic Parameters
+------------------
+
+.. include:: ../_partials/generic_parameters.rst
+
+Operators
+---------
+
+.. _howto/operator:SsmRunCommandOperator:
+
+Runs commands on one or more managed nodes
+==========================================
+
+To run SSM run command, you can use
+:class:`~airflow.providers.amazon.aws.operators.ssm.SsmRunCommandOperator`.
+
+To monitor the state of the command for a specific instance, you can use the "command_executed"
+Waiter. Additionally, you can use the following components to track the status of the command execution:
+:class:`~airflow.providers.amazon.aws.sensors.ssm.SsmRunCommandCompletedSensor` Sensor,
+or the :class:`~airflow.providers.amazon.aws.triggers.ssm.SsmRunCommandTrigger` Trigger.
+
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_ssm.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_run_command]
+    :end-before: [END howto_operator_run_command]
+
+Sensors
+-------
+
+.. _howto/sensor:SsmRunCommandCompletedSensor:
+
+Wait for an Amazon SSM run command
+==================================
+
+To wait on the state of an Amazon SSM run command job until it reaches a terminal state you can use
+:class:`~airflow.providers.amazon.aws.sensors.SSM.SsmRunCommandCompletedSensor`
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_ssm.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_run_command]
+    :end-before: [END howto_sensor_run_command]
+
+IAM Permissions
+---------------
+
+You need to ensure the following IAM permissions are granted to allow Airflow to run and monitor SSM Run Command executions:
+
+.. code-block::
+
+    {
+      "Effect": "Allow",
+      "Action": [
+        "ssm:SendCommand",
+        "ssm:ListCommandInvocations",
+        "ssm:GetCommandInvocation"
+      ],
+      "Resource": "*"
+    }
+
+This policy allows access to all SSM documents and managed instances. For production environments,
+it is recommended to restrict the ``Resource`` field to specific SSM document ARNs and, if applicable,
+to the ARNs of intended target resources (such as EC2 instances), in accordance with the principle of least privilege.
+
+Reference
+---------
+
+* `AWS boto3 library documentation for Amazon SSM <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ssm.html>`__

providers/amazon/provider.yaml

@@ -281,6 +281,8 @@ integrations:
   - integration-name: Amazon Systems Manager (SSM)
     external-doc-url: https://aws.amazon.com/systems-manager/
     logo: /docs/integration-logos/AWS-Systems-Manager_light-bg@4x.png
+    how-to-guide:
+      - /docs/apache-airflow-providers-amazon/operators/ssm.rst
     tags: [aws]
   - integration-name: Amazon Web Services
     external-doc-url: https://aws.amazon.com/
@@ -430,6 +432,9 @@ operators:
   - integration-name: AWS Step Functions
     python-modules:
       - airflow.providers.amazon.aws.operators.step_function
+  - integration-name: Amazon Systems Manager (SSM)
+    python-modules:
+      - airflow.providers.amazon.aws.operators.ssm
   - integration-name: Amazon RDS
     python-modules:
       - airflow.providers.amazon.aws.operators.rds
@@ -531,6 +536,9 @@ sensors:
   - integration-name: AWS Step Functions
     python-modules:
       - airflow.providers.amazon.aws.sensors.step_function
+  - integration-name: Amazon Systems Manager (SSM)
+    python-modules:
+      - airflow.providers.amazon.aws.sensors.ssm
   - integration-name: Amazon QuickSight
     python-modules:
       - airflow.providers.amazon.aws.sensors.quicksight
@@ -737,6 +745,9 @@ triggers:
   - integration-name: Amazon Simple Storage Service (S3)
     python-modules:
       - airflow.providers.amazon.aws.triggers.s3
+  - integration-name: Amazon Systems Manager (SSM)
+    python-modules:
+      - airflow.providers.amazon.aws.triggers.ssm
   - integration-name: Amazon EMR
     python-modules:
       - airflow.providers.amazon.aws.triggers.emr

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

@@ -0,0 +1,131 @@
+# 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 collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.ssm import SsmHook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.ssm import SsmRunCommandTrigger
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
+    """
+    Executes the SSM Run Command to perform actions on managed instances.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:SsmRunCommandOperator`
+
+    :param document_name: The name of the Amazon Web Services Systems Manager document (SSM document) to run.
+    :param run_command_kwargs: Optional parameters to pass to the send_command API.
+
+    :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 120)
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 75)
+    :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :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 botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = SsmHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "document_name",
+        "run_command_kwargs",
+    )
+
+    def __init__(
+        self,
+        *,
+        document_name: str,
+        run_command_kwargs: dict[str, Any] | None = None,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 120,
+        waiter_max_attempts: int = 75,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
+
+        self.document_name = document_name
+        self.run_command_kwargs = run_command_kwargs or {}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
+            raise AirflowException(f"Error while running run command: {event}")
+
+        self.log.info("SSM run command `%s` completed.", event["command_id"])
+        return event["command_id"]
+
+    def execute(self, context: Context):
+        response = self.hook.conn.send_command(
+            DocumentName=self.document_name,
+            **self.run_command_kwargs,
+        )
+
+        command_id = response["Command"]["CommandId"]
+        task_description = f"SSM run command {command_id} to complete."
+
+        if self.deferrable:
+            self.log.info("Deferring for %s", task_description)
+            self.defer(
+                trigger=SsmRunCommandTrigger(
+                    command_id=command_id,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+            )
+
+        elif self.wait_for_completion:
+            self.log.info("Waiting for %s", task_description)
+            waiter = self.hook.get_waiter("command_executed")
+
+            instance_ids = response["Command"]["InstanceIds"]
+            for instance_id in instance_ids:
+                waiter.wait(
+                    CommandId=command_id,
+                    InstanceId=instance_id,
+                    WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+                )
+
+        return command_id