apache-airflow-providers-amazon 9.19.0rc1__py3-none-any.whl → 9.21.0rc1__py3-none-any.whl

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

@@ -19,12 +19,14 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
+from botocore.exceptions import WaiterError
+
 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
+from airflow.providers.common.compat.sdk import conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context
@@ -50,6 +52,12 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
         (default: 120)
     :param waiter_max_attempts: Maximum number of attempts to check for job
         completion. (default: 75)
+    :param fail_on_nonzero_exit: If True (default), the operator will fail when
+        the command returns a non-zero exit code. If False, the operator will
+        complete successfully regardless of the command exit code, allowing
+        downstream tasks to handle exit codes for workflow routing. Note that
+        AWS-level failures (Cancelled, TimedOut) will still raise exceptions
+        even when this is False. (default: True)
     :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)
@@ -81,6 +89,7 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
         wait_for_completion: bool = True,
         waiter_delay: int = 120,
         waiter_max_attempts: int = 75,
+        fail_on_nonzero_exit: bool = True,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ):
@@ -88,6 +97,7 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
         self.wait_for_completion = wait_for_completion
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
+        self.fail_on_nonzero_exit = fail_on_nonzero_exit
         self.deferrable = deferrable
 
         self.document_name = document_name
@@ -96,6 +106,19 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
         event = validate_execute_complete_event(event)
 
+        if event["status"] == "failed":
+            # Command failed - raise an exception with detailed information
+            command_status = event.get("command_status", "Unknown")
+            exit_code = event.get("exit_code", -1)
+            instance_id = event.get("instance_id", "Unknown")
+            message = event.get("message", "Command failed")
+
+            error_msg = (
+                f"SSM run command {event['command_id']} failed on instance {instance_id}. "
+                f"Status: {command_status}, Exit code: {exit_code}. {message}"
+            )
+            raise RuntimeError(error_msg)
+
         if event["status"] != "success":
             raise RuntimeError(f"Error while running run command: {event}")
 
@@ -118,6 +141,7 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
                     command_id=command_id,
                     waiter_delay=self.waiter_delay,
                     waiter_max_attempts=self.waiter_max_attempts,
+                    fail_on_nonzero_exit=self.fail_on_nonzero_exit,
                     aws_conn_id=self.aws_conn_id,
                     region_name=self.region_name,
                     verify=self.verify,
@@ -132,14 +156,35 @@ class SsmRunCommandOperator(AwsBaseOperator[SsmHook]):
 
             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,
-                    },
-                )
+                try:
+                    waiter.wait(
+                        CommandId=command_id,
+                        InstanceId=instance_id,
+                        WaiterConfig={
+                            "Delay": self.waiter_delay,
+                            "MaxAttempts": self.waiter_max_attempts,
+                        },
+                    )
+                except WaiterError:
+                    if not self.fail_on_nonzero_exit:
+                        # Enhanced mode: distinguish between AWS-level and command-level failures
+                        invocation = self.hook.get_command_invocation(command_id, instance_id)
+                        status = invocation.get("Status", "")
+
+                        # AWS-level failures should always raise
+                        if SsmHook.is_aws_level_failure(status):
+                            raise
+
+                        # Command-level failure - tolerate it in enhanced mode
+                        self.log.info(
+                            "Command completed with status %s (exit code: %s). "
+                            "Continuing due to fail_on_nonzero_exit=False",
+                            status,
+                            invocation.get("ResponseCode", "unknown"),
+                        )
+                    else:
+                        # Traditional mode: all failures raise
+                        raise
 
         return command_id
 
@@ -148,14 +193,6 @@ class SsmGetCommandInvocationOperator(AwsBaseOperator[SsmHook]):
     """
     Retrieves the output and execution details of an SSM command invocation.
 
-    This operator allows you to fetch the standard output, standard error,
-    execution status, and other details from SSM commands. It can be used to
-    retrieve output from commands executed by SsmRunCommandOperator in previous
-    tasks, or from commands executed outside of Airflow entirely.
-
-    The operator returns structured data including stdout, stderr, execution
-    times, and status information for each instance that executed the command.
-
     .. seealso::
         For more information on how to use this operator, take a look at the
         guide:

airflow/providers/amazon/aws/operators/step_function.py

@@ -21,7 +21,6 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook
 from airflow.providers.amazon.aws.links.step_function import (
     StateMachineDetailsLink,
@@ -31,7 +30,7 @@ from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.step_function import StepFunctionsExecutionCompleteTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/queues/sqs.py

@@ -19,8 +19,8 @@ from __future__ import annotations
 import re
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
 try:
     from airflow.providers.common.messaging.providers.base_provider import BaseMessageQueueProvider

airflow/providers/amazon/aws/secrets/secrets_manager.py

@@ -197,11 +197,12 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
 
         return conn_d
 
-    def get_conn_value(self, conn_id: str) -> str | None:
+    def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | None:
         """
         Get serialized representation of Connection.
 
         :param conn_id: connection id
+        :param team_name: Team name associated to the task trying to access the connection (if any)
         """
         if self.connections_prefix is None:
             return None

airflow/providers/amazon/aws/secrets/systems_manager.py

@@ -132,11 +132,12 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
         session = SessionFactory(conn=conn_config).create_session()
         return session.client(service_name="ssm", **client_kwargs)
 
-    def get_conn_value(self, conn_id: str) -> str | None:
+    def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | None:
         """
         Get param value.
 
         :param conn_id: connection id
+        :param team_name: Team name associated to the task trying to access the connection (if any)
         """
         if self.connections_prefix is None:
             return None

airflow/providers/amazon/aws/sensors/batch.py

@@ -20,12 +20,11 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.batch import BatchJobTrigger
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/bedrock.py

@@ -21,7 +21,6 @@ import abc
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, TypeVar
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.bedrock import (
@@ -33,7 +32,7 @@ from airflow.providers.amazon.aws.triggers.bedrock import (
     BedrockProvisionModelThroughputCompletedTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.triggers.bedrock import BedrockBaseBatchInferenceTrigger

airflow/providers/amazon/aws/sensors/comprehend.py

@@ -20,7 +20,6 @@ import abc
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.comprehend import ComprehendHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.comprehend import (
@@ -28,7 +27,7 @@ from airflow.providers.amazon.aws.triggers.comprehend import (
     ComprehendPiiEntitiesDetectionJobCompletedTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/ec2.py

@@ -20,13 +20,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.ec2 import EC2Hook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.ec2 import EC2StateSensorTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/emr.py

@@ -21,7 +21,6 @@ from collections.abc import Iterable, Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
 from airflow.providers.amazon.aws.links.emr import EmrClusterLink, EmrLogsLink, get_log_uri
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
@@ -32,7 +31,7 @@ from airflow.providers.amazon.aws.triggers.emr import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/glue.py

@@ -20,7 +20,6 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.glue import GlueDataQualityHook, GlueJobHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.glue import (
@@ -30,7 +29,7 @@ from airflow.providers.amazon.aws.triggers.glue import (
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/glue_catalog_partition.py

@@ -21,13 +21,12 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.glue_catalog import GlueCatalogHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.glue import GlueCatalogPartitionTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/kinesis_analytics.py

@@ -19,14 +19,13 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.kinesis_analytics import KinesisAnalyticsV2Hook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.kinesis_analytics import (
     KinesisAnalyticsV2ApplicationOperationCompleteTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/mwaa.py

@@ -20,12 +20,11 @@ from __future__ import annotations
 from collections.abc import Collection, Sequence
 from typing import TYPE_CHECKING, Any, Literal
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.mwaa import MwaaDagRunCompletedTrigger, MwaaTaskCompletedTrigger
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 from airflow.utils.state import DagRunState, TaskInstanceState
 
 if TYPE_CHECKING:

airflow/providers/amazon/aws/sensors/opensearch_serverless.py

@@ -19,14 +19,13 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.opensearch_serverless import OpenSearchServerlessHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.opensearch_serverless import (
     OpenSearchServerlessCollectionActiveTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 from airflow.utils.helpers import exactly_one
 
 if TYPE_CHECKING:

airflow/providers/amazon/aws/sensors/redshift_cluster.py

@@ -20,13 +20,12 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.redshift_cluster import RedshiftClusterTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/s3.py

@@ -25,13 +25,12 @@ from collections.abc import Callable, Sequence
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, cast
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.s3 import S3KeysUnchangedTrigger, S3KeyTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.providers.common.compat.sdk import AirflowException, poke_mode_only
+from airflow.providers.common.compat.sdk import AirflowException, conf, poke_mode_only
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/amazon/aws/sensors/sqs.py

@@ -23,14 +23,13 @@ from collections.abc import Collection, Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.sqs import process_response
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, conf
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection

airflow/providers/amazon/aws/sensors/ssm.py

@@ -20,12 +20,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.ssm import SsmHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 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
+from airflow.providers.common.compat.sdk import conf
 
 if TYPE_CHECKING:
     from airflow.sdk import Context
@@ -44,6 +44,11 @@ class SsmRunCommandCompletedSensor(AwsBaseSensor[SsmHook]):
         :ref:`howto/sensor:SsmRunCommandCompletedSensor`
 
     :param command_id: The ID of the AWS SSM Run Command.
+    :param fail_on_nonzero_exit: If True (default), the sensor will fail when the command
+        returns a non-zero exit code. If False, the sensor will complete successfully
+        for both Success and Failed command statuses, allowing downstream tasks to handle
+        exit codes. AWS-level failures (Cancelled, TimedOut) will still raise exceptions.
+        (default: True)
     :param deferrable: If True, the sensor will operate in deferrable mode.
         This mode requires aiobotocore module to be installed.
         (default: False, but can be overridden in config file by setting
@@ -85,6 +90,7 @@ class SsmRunCommandCompletedSensor(AwsBaseSensor[SsmHook]):
         self,
         *,
         command_id,
+        fail_on_nonzero_exit: bool = True,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         poke_interval: int = 120,
         max_retries: int = 75,
@@ -92,6 +98,7 @@ class SsmRunCommandCompletedSensor(AwsBaseSensor[SsmHook]):
     ):
         super().__init__(**kwargs)
         self.command_id = command_id
+        self.fail_on_nonzero_exit = fail_on_nonzero_exit
         self.deferrable = deferrable
         self.poke_interval = poke_interval
         self.max_retries = max_retries
@@ -112,7 +119,19 @@ class SsmRunCommandCompletedSensor(AwsBaseSensor[SsmHook]):
             state = invocation["Status"]
 
             if state in self.FAILURE_STATES:
-                raise RuntimeError(self.FAILURE_MESSAGE)
+                # Check if we should tolerate this failure
+                if self.fail_on_nonzero_exit:
+                    raise RuntimeError(self.FAILURE_MESSAGE)  # Traditional behavior
+
+                # Only fail on AWS-level issues, tolerate command failures
+                if SsmHook.is_aws_level_failure(state):
+                    raise RuntimeError(f"SSM command {self.command_id} {state}")
+
+                # Command failed but we're tolerating it
+                self.log.info(
+                    "Command invocation has status %s. Continuing due to fail_on_nonzero_exit=False",
+                    state,
+                )
 
             if state in self.INTERMEDIATE_STATES:
                 return False
@@ -127,6 +146,7 @@ class SsmRunCommandCompletedSensor(AwsBaseSensor[SsmHook]):
                     waiter_delay=int(self.poke_interval),
                     waiter_max_attempts=self.max_retries,
                     aws_conn_id=self.aws_conn_id,
+                    fail_on_nonzero_exit=self.fail_on_nonzero_exit,
                 ),
                 method_name="execute_complete",
             )

airflow/providers/amazon/aws/transfers/azure_blob_to_s3.py

@@ -28,7 +28,7 @@ from airflow.providers.common.compat.sdk import BaseOperator
 try:
     from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
 except ModuleNotFoundError as e:
-    from airflow.exceptions import AirflowOptionalProviderFeatureException
+    from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
     raise AirflowOptionalProviderFeatureException(e)
 

airflow/providers/amazon/aws/transfers/sql_to_s3.py

@@ -183,7 +183,7 @@ class SqlToS3Operator(BaseOperator):
             import numpy as np
             import pandas as pd
         except ImportError as e:
-            from airflow.exceptions import AirflowOptionalProviderFeatureException
+            from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
             raise AirflowOptionalProviderFeatureException(e)
 

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

@@ -351,8 +351,8 @@ class EmrServerlessStartJobTrigger(AwsBaseWaiterTrigger):
             failure_message="Serverless Job failed",
             status_message="Serverless Job status is",
             status_queries=["jobRun.state", "jobRun.stateDetails"],
-            return_key="job_id",
-            return_value=job_id,
+            return_key="job_details",
+            return_value={"application_id": application_id, "job_id": job_id},
             waiter_delay=waiter_delay,
             waiter_max_attempts=waiter_max_attempts,
             aws_conn_id=aws_conn_id,

airflow/providers/amazon/aws/triggers/ssm.py

@@ -35,6 +35,9 @@ class SsmRunCommandTrigger(AwsBaseWaiterTrigger):
     :param command_id: The ID of the AWS SSM Run Command.
     :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
     :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
+    :param fail_on_nonzero_exit: If True (default), the trigger will fail when the command returns
+        a non-zero exit code. If False, the trigger will complete successfully regardless of the
+        command exit code. (default: True)
     :param aws_conn_id: The Airflow connection used for AWS credentials.
     :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:
@@ -49,18 +52,19 @@ class SsmRunCommandTrigger(AwsBaseWaiterTrigger):
         command_id: str,
         waiter_delay: int = 120,
         waiter_max_attempts: int = 75,
+        fail_on_nonzero_exit: bool = True,
         aws_conn_id: str | None = None,
         region_name: str | None = None,
         verify: bool | str | None = None,
         botocore_config: dict | None = None,
     ) -> None:
         super().__init__(
-            serialized_fields={"command_id": command_id},
+            serialized_fields={"command_id": command_id, "fail_on_nonzero_exit": fail_on_nonzero_exit},
             waiter_name="command_executed",
             waiter_args={"CommandId": command_id},
             failure_message="SSM run command failed.",
             status_message="Status of SSM run command is",
-            status_queries=["status"],
+            status_queries=["Status"],
             return_key="command_id",
             return_value=command_id,
             waiter_delay=waiter_delay,
@@ -71,6 +75,7 @@ class SsmRunCommandTrigger(AwsBaseWaiterTrigger):
             botocore_config=botocore_config,
         )
         self.command_id = command_id
+        self.fail_on_nonzero_exit = fail_on_nonzero_exit
 
     def hook(self) -> AwsGenericHook:
         return SsmHook(
@@ -89,14 +94,66 @@ class SsmRunCommandTrigger(AwsBaseWaiterTrigger):
 
             for instance_id in instance_ids:
                 self.waiter_args["InstanceId"] = instance_id
-                await async_wait(
-                    waiter,
-                    self.waiter_delay,
-                    self.attempts,
-                    self.waiter_args,
-                    self.failure_message,
-                    self.status_message,
-                    self.status_queries,
-                )
+                try:
+                    await async_wait(
+                        waiter,
+                        self.waiter_delay,
+                        self.attempts,
+                        self.waiter_args,
+                        self.failure_message,
+                        self.status_message,
+                        self.status_queries,
+                    )
+                except Exception:
+                    # Get detailed invocation information to determine failure type
+                    invocation = await client.get_command_invocation(
+                        CommandId=self.command_id, InstanceId=instance_id
+                    )
+                    status = invocation.get("Status", "")
+                    response_code = invocation.get("ResponseCode", -1)
+
+                    # AWS-level failures should always raise
+                    if SsmHook.is_aws_level_failure(status):
+                        self.log.error(
+                            "AWS-level failure for command %s on instance %s: status=%s",
+                            self.command_id,
+                            instance_id,
+                            status,
+                        )
+                        raise
+
+                    # Command-level failure (non-zero exit code)
+                    if not self.fail_on_nonzero_exit:
+                        # Enhanced mode: tolerate command-level failures
+                        self.log.info(
+                            "Command %s completed with status %s (exit code: %s) for instance %s. "
+                            "Continuing due to fail_on_nonzero_exit=False",
+                            self.command_id,
+                            status,
+                            response_code,
+                            instance_id,
+                        )
+                        continue
+                    else:
+                        # Traditional mode: yield failure event instead of raising
+                        # This allows the operator to handle the failure gracefully
+                        self.log.warning(
+                            "Command %s failed with status %s (exit code: %s) for instance %s",
+                            self.command_id,
+                            status,
+                            response_code,
+                            instance_id,
+                        )
+                        yield TriggerEvent(
+                            {
+                                "status": "failed",
+                                "message": f"Command failed with status {status} (exit code: {response_code})",
+                                "command_status": status,
+                                "exit_code": response_code,
+                                "instance_id": instance_id,
+                                self.return_key: self.return_value,
+                            }
+                        )
+                        return
 
         yield TriggerEvent({"status": "success", self.return_key: self.return_value})

airflow/providers/amazon/get_provider_info.py

@@ -1207,7 +1207,7 @@ def get_provider_info():
                         "example": "True",
                         "default": "True",
                     },
-                    "max_run_task_attempts": {
+                    "max_invoke_attempts": {
                         "description": "The maximum number of times the Lambda Executor should attempt to start an Airflow task.\n",
                         "version_added": "9.9.0",
                         "type": "integer",
@@ -1390,5 +1390,6 @@ def get_provider_info():
         },
         "executors": ["airflow.providers.amazon.aws.executors.ecs.ecs_executor.AwsEcsExecutor"],
         "auth-managers": ["airflow.providers.amazon.aws.auth_manager.aws_auth_manager.AwsAuthManager"],
+        "cli": ["airflow.providers.amazon.aws.cli.definition.get_aws_cli_commands"],
         "queues": ["airflow.providers.amazon.aws.queues.sqs.SqsMessageQueueProvider"],
     }