apache-airflow-providers-amazon 9.15.0__py3-none-any.whl → 9.18.0rc2__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -29,11 +29,11 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.15.0"
+__version__ = "9.18.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.10.0"
+    "2.11.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-amazon:{__version__}` needs Apache Airflow 2.10.0+"
+        f"The package `apache-airflow-providers-amazon:{__version__}` needs Apache Airflow 2.11.0+"
     )

airflow/providers/amazon/aws/auth_manager/routes/login.py

@@ -35,6 +35,7 @@ from airflow.configuration import conf
 from airflow.providers.amazon.aws.auth_manager.constants import CONF_SAML_METADATA_URL_KEY, CONF_SECTION_NAME
 from airflow.providers.amazon.aws.auth_manager.datamodels.login import LoginResponse
 from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_1_1_PLUS
 
 try:
     from onelogin.saml2.auth import OneLogin_Saml2_Auth
@@ -101,7 +102,12 @@ def login_callback(request: Request):
     if relay_state == "login-redirect":
         response = RedirectResponse(url=url, status_code=303)
         secure = bool(conf.get("api", "ssl_cert", fallback=""))
-        response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
+        # In Airflow 3.1.1 authentication changes, front-end no longer handle the token
+        # See https://github.com/apache/airflow/pull/55506
+        if AIRFLOW_V_3_1_1_PLUS:
+            response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure, httponly=True)
+        else:
+            response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
         return response
     if relay_state == "login-token":
         return LoginResponse(access_token=token)

airflow/providers/amazon/aws/executors/aws_lambda/docker/app.py

@@ -66,7 +66,11 @@ def run_and_report(command, task_key):
     try:
         log.info("Starting execution for task: %s", task_key)
         result = subprocess.run(
-            command, shell=isinstance(command, str), stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+            command,
+            check=False,
+            shell=isinstance(command, str),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
         )
         return_code = result.returncode
         log.info("Execution completed for task %s with return code %s", task_key, return_code)

airflow/providers/amazon/aws/hooks/athena.py

@@ -40,11 +40,15 @@ MULTI_LINE_QUERY_LOG_PREFIX = "\n\t\t"
 
 def query_params_to_string(params: dict[str, str | Collection[str]]) -> str:
     result = ""
-    for key, value in params.items():
+    for key, original_value in params.items():
+        value: str | Collection[str]
         if key == "QueryString":
             value = (
-                MULTI_LINE_QUERY_LOG_PREFIX + str(value).replace("\n", MULTI_LINE_QUERY_LOG_PREFIX).rstrip()
+                MULTI_LINE_QUERY_LOG_PREFIX
+                + str(original_value).replace("\n", MULTI_LINE_QUERY_LOG_PREFIX).rstrip()
             )
+        else:
+            value = original_value
         result += f"\t{key}: {value}\n"
     return result.rstrip()
 

airflow/providers/amazon/aws/hooks/athena_sql.py

@@ -56,7 +56,7 @@ class AthenaSQLHook(AwsBaseHook, DbApiHook):
         :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
 
     .. note::
-        get_uri() depends on SQLAlchemy and PyAthena.
+        get_uri() depends on SQLAlchemy and PyAthena
     """
 
     conn_name_attr = "athena_conn_id"
@@ -163,7 +163,7 @@ class AthenaSQLHook(AwsBaseHook, DbApiHook):
             port=443,
             database=conn_params["schema_name"],
             query={"aws_session_token": creds.token, **self.conn.extra_dejson},
-        )
+        ).render_as_string(hide_password=False)
 
     def get_conn(self) -> AthenaConnection:
         """Get a ``pyathena.Connection`` object."""

airflow/providers/amazon/aws/hooks/base_aws.py

@@ -60,7 +60,7 @@ from airflow.exceptions import (
 from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper
 from airflow.providers.amazon.aws.utils.identifiers import generate_uuid
 from airflow.providers.amazon.aws.utils.suppress import return_on_error
-from airflow.providers.amazon.version_compat import BaseHook
+from airflow.providers.common.compat.sdk import BaseHook
 from airflow.providers_manager import ProvidersManager
 from airflow.utils.helpers import exactly_one
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -790,7 +790,7 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
     async def get_async_conn(self):
         """Get an aiobotocore client to use for async operations."""
         # We have to wrap the call `self.get_client_type` in another call `_get_async_conn`,
-        # because one of it's arguments `self.region_name` is a `@property` decorated function
+        # because one of its arguments `self.region_name` is a `@property` decorated function
         # calling the cached property `self.conn_config` at the end.
         return await sync_to_async(self._get_async_conn)()
 

airflow/providers/amazon/aws/hooks/batch_client.py

@@ -386,8 +386,7 @@ class BatchClientHook(AwsBaseHook):
             )
             if job_status in match_status:
                 return True
-        else:
-            raise AirflowException(f"AWS Batch job ({job_id}) status checks exceed max_retries")
+        raise AirflowException(f"AWS Batch job ({job_id}) status checks exceed max_retries")
 
     def get_job_description(self, job_id: str) -> dict:
         """
@@ -426,10 +425,9 @@ class BatchClientHook(AwsBaseHook):
                     "check Amazon Provider AWS Connection documentation for more details.",
                     str(err),
                 )
-        else:
-            raise AirflowException(
-                f"AWS Batch job ({job_id}) description error: exceeded status_retries ({self.status_retries})"
-            )
+        raise AirflowException(
+            f"AWS Batch job ({job_id}) description error: exceeded status_retries ({self.status_retries})"
+        )
 
     @staticmethod
     def parse_job_description(job_id: str, response: dict) -> dict:

airflow/providers/amazon/aws/hooks/batch_waiters.py

@@ -33,7 +33,6 @@ from copy import deepcopy
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
-import botocore.client
 import botocore.exceptions
 import botocore.waiter
 

airflow/providers/amazon/aws/hooks/chime.py

@@ -33,7 +33,7 @@ class ChimeWebhookHook(HttpHook):
     """
     Interact with Amazon Chime Webhooks to create notifications.
 
-    .. warning:: This hook is only designed to work with web hooks and not chat bots.
+    .. warning:: This hook is only designed to work with web hooks and not chatbots.
 
     :param chime_conn_id: :ref:`Amazon Chime Connection ID <howto/connection:chime>`
         with Endpoint as `https://hooks.chime.aws` and the webhook token

airflow/providers/amazon/aws/hooks/datasync.py

@@ -21,8 +21,9 @@ from __future__ import annotations
 import time
 from urllib.parse import urlsplit
 
-from airflow.exceptions import AirflowBadRequest, AirflowException, AirflowTaskTimeout
+from airflow.exceptions import AirflowBadRequest, AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowTaskTimeout
 
 
 class DataSyncHook(AwsBaseHook):
@@ -319,5 +320,4 @@ class DataSyncHook(AwsBaseHook):
             else:
                 raise AirflowException(f"Unknown status: {status}")  # Should never happen
             time.sleep(self.wait_interval_seconds)
-        else:
-            raise AirflowTaskTimeout("Max iterations exceeded!")
+        raise AirflowTaskTimeout("Max iterations exceeded!")

airflow/providers/amazon/aws/hooks/firehose.py

@@ -0,0 +1,56 @@
+#
+# 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.
+"""This module contains AWS Firehose hook."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class FirehoseHook(AwsBaseHook):
+    """
+    Interact with Amazon Kinesis Firehose.
+
+    Provide thick wrapper around :external+boto3:py:class:`boto3.client("firehose") <Firehose.Client>`.
+
+    :param delivery_stream: Name of the delivery stream
+
+    Additional arguments (such as ``aws_conn_id``) may be specified and
+    are passed down to the underlying AwsBaseHook.
+
+    .. seealso::
+        - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+    """
+
+    def __init__(self, delivery_stream: str, *args, **kwargs) -> None:
+        self.delivery_stream = delivery_stream
+        kwargs["client_type"] = "firehose"
+        super().__init__(*args, **kwargs)
+
+    def put_records(self, records: Iterable) -> dict:
+        """
+        Write batch records to Kinesis Firehose.
+
+        .. seealso::
+            - :external+boto3:py:meth:`Firehose.Client.put_record_batch`
+
+        :param records: list of records
+        """
+        return self.get_conn().put_record_batch(DeliveryStreamName=self.delivery_stream, Records=records)

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

@@ -565,7 +565,13 @@ class GlueDataQualityHook(AwsBaseHook):
         Rule_3    ColumnLength "marketplace" between 1 and 2     FAIL        {'Column.marketplace.MaximumLength': 9.0, 'Column.marketplace.MinimumLength': 3.0}     Value: 9.0 does not meet the constraint requirement!
 
         """
-        import pandas as pd
+        try:
+            import pandas as pd
+        except ImportError:
+            self.log.warning(
+                "Pandas is not installed. Please install pandas to see the detailed Data Quality results."
+            )
+            return
 
         pd.set_option("display.max_rows", None)
         pd.set_option("display.max_columns", None)

airflow/providers/amazon/aws/hooks/kinesis.py

@@ -19,12 +19,14 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+import warnings
 
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.hooks.firehose import FirehoseHook as _FirehoseHook
 
 
-class FirehoseHook(AwsBaseHook):
+class FirehoseHook(_FirehoseHook):
     """
     Interact with Amazon Kinesis Firehose.
 
@@ -37,20 +39,36 @@ class FirehoseHook(AwsBaseHook):
 
     .. seealso::
         - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+    .. deprecated::
+        This hook was moved. Import from
+       :class:`airflow.providers.amazon.aws.hooks.firehose.FirehoseHook`
+       instead of kinesis.py
     """
 
-    def __init__(self, delivery_stream: str, *args, **kwargs) -> None:
-        self.delivery_stream = delivery_stream
-        kwargs["client_type"] = "firehose"
+    def __init__(self, *args, **kwargs) -> None:
+        warnings.warn(
+            "Importing FirehoseHook from kinesis.py is deprecated "
+            "and will be removed in a future release. "
+            "Please import it from firehose.py instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
         super().__init__(*args, **kwargs)
 
-    def put_records(self, records: Iterable):
-        """
-        Write batch records to Kinesis Firehose.
 
-        .. seealso::
-            - :external+boto3:py:meth:`Firehose.Client.put_record_batch`
+class KinesisHook(AwsBaseHook):
+    """
+    Interact with Amazon Kinesis.
+
+    Provide thin wrapper around :external+boto3:py:class:`boto3.client("kinesis") <Kinesis.Client>`.
+
+    Additional arguments (such as ``aws_conn_id``) may be specified and
+    are passed down to the underlying AwsBaseHook.
 
-        :param records: list of records
-        """
-        return self.get_conn().put_record_batch(DeliveryStreamName=self.delivery_stream, Records=records)
+    .. seealso::
+        - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        kwargs["client_type"] = "kinesis"
+        super().__init__(*args, **kwargs)

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

@@ -18,6 +18,9 @@
 
 from __future__ import annotations
 
+import warnings
+from typing import Literal
+
 import requests
 from botocore.exceptions import ClientError
 
@@ -55,6 +58,7 @@ class MwaaHook(AwsBaseHook):
         body: dict | None = None,
         query_params: dict | None = None,
         generate_local_token: bool = False,
+        airflow_version: Literal[2, 3] | None = None,
     ) -> dict:
         """
         Invoke the REST API on the Airflow webserver with the specified inputs.
@@ -70,6 +74,8 @@ class MwaaHook(AwsBaseHook):
         :param generate_local_token: If True, only the local web token method is used without trying boto's
             `invoke_rest_api` first. If False, the local web token method is used as a fallback after trying
             boto's `invoke_rest_api`
+        :param airflow_version: The Airflow major version the MWAA environment runs.
+            This parameter is only used if the local web token method is used to call Airflow API.
         """
         # Filter out keys with None values because Airflow REST API doesn't accept requests otherwise
         body = {k: v for k, v in body.items() if v is not None} if body else {}
@@ -83,7 +89,7 @@ class MwaaHook(AwsBaseHook):
         }
 
         if generate_local_token:
-            return self._invoke_rest_api_using_local_session_token(**api_kwargs)
+            return self._invoke_rest_api_using_local_session_token(airflow_version, **api_kwargs)
 
         try:
             response = self.conn.invoke_rest_api(**api_kwargs)
@@ -100,7 +106,7 @@ class MwaaHook(AwsBaseHook):
                 self.log.info(
                     "Access Denied due to missing airflow:InvokeRestApi in IAM policy. Trying again by generating local token..."
                 )
-                return self._invoke_rest_api_using_local_session_token(**api_kwargs)
+                return self._invoke_rest_api_using_local_session_token(airflow_version, **api_kwargs)
             to_log = e.response
             # ResponseMetadata is removed because it contains data that is either very unlikely to be
             # useful in XComs and logs, or redundant given the data already included in the response
@@ -110,14 +116,35 @@ class MwaaHook(AwsBaseHook):
 
     def _invoke_rest_api_using_local_session_token(
         self,
+        airflow_version: Literal[2, 3] | None = None,
         **api_kwargs,
     ) -> dict:
+        if not airflow_version:
+            warnings.warn(
+                "The parameter ``airflow_version`` in ``MwaaHook.invoke_rest_api`` is not "
+                "specified and the local web token method is being used. "
+                "The default Airflow version being used is 2 but this value will change in the future. "
+                "To avoid any unexpected behavior, please explicitly specify the Airflow version.",
+                FutureWarning,
+                stacklevel=3,
+            )
+            airflow_version = 2
+
         try:
-            session, hostname = self._get_session_conn(api_kwargs["Name"])
+            session, hostname, login_response = self._get_session_conn(api_kwargs["Name"], airflow_version)
+
+            headers = {}
+            if airflow_version == 3:
+                headers = {
+                    "Authorization": f"Bearer {login_response.cookies['_token']}",
+                    "Content-Type": "application/json",
+                }
 
+            api_version = "v1" if airflow_version == 2 else "v2"
             response = session.request(
                 method=api_kwargs["Method"],
-                url=f"https://{hostname}/api/v1{api_kwargs['Path']}",
+                url=f"https://{hostname}/api/{api_version}{api_kwargs['Path']}",
+                headers=headers,
                 params=api_kwargs["QueryParameters"],
                 json=api_kwargs["Body"],
                 timeout=10,
@@ -134,15 +161,19 @@ class MwaaHook(AwsBaseHook):
         }
 
     # Based on: https://docs.aws.amazon.com/mwaa/latest/userguide/access-mwaa-apache-airflow-rest-api.html#create-web-server-session-token
-    def _get_session_conn(self, env_name: str) -> tuple:
+    def _get_session_conn(self, env_name: str, airflow_version: Literal[2, 3]) -> tuple:
         create_token_response = self.conn.create_web_login_token(Name=env_name)
         web_server_hostname = create_token_response["WebServerHostname"]
         web_token = create_token_response["WebToken"]
 
-        login_url = f"https://{web_server_hostname}/aws_mwaa/login"
+        login_url = (
+            f"https://{web_server_hostname}/aws_mwaa/login"
+            if airflow_version == 2
+            else f"https://{web_server_hostname}/pluginsv2/aws_mwaa/login"
+        )
         login_payload = {"token": web_token}
         session = requests.Session()
         login_response = session.post(login_url, data=login_payload, timeout=10)
         login_response.raise_for_status()
 
-        return session, web_server_hostname
+        return session, web_server_hostname, login_response

airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -51,7 +51,7 @@ class RedshiftSQLHook(DbApiHook):
         :ref:`Amazon Redshift connection id<howto/connection:redshift>`
 
     .. note::
-        get_sqlalchemy_engine() and get_uri() depend on sqlalchemy-amazon-redshift
+        get_sqlalchemy_engine() and get_uri() depend on sqlalchemy-amazon-redshift.
     """
 
     conn_name_attr = "redshift_conn_id"
@@ -155,10 +155,21 @@ class RedshiftSQLHook(DbApiHook):
         if "user" in conn_params:
             conn_params["username"] = conn_params.pop("user")
 
-        # Compatibility: The 'create' factory method was added in SQLAlchemy 1.4
-        # to replace calling the default URL constructor directly.
-        create_url = getattr(URL, "create", URL)
-        return str(create_url(drivername="postgresql", **conn_params))
+        # Use URL.create for SQLAlchemy 2 compatibility
+        username = conn_params.get("username")
+        password = conn_params.get("password")
+        host = conn_params.get("host")
+        port = conn_params.get("port")
+        database = conn_params.get("database")
+
+        return URL.create(
+            drivername="postgresql",
+            username=str(username) if username is not None else None,
+            password=str(password) if password is not None else None,
+            host=str(host) if host is not None else None,
+            port=int(port) if port is not None else None,
+            database=str(database) if database is not None else None,
+        ).render_as_string(hide_password=False)
 
     def get_sqlalchemy_engine(self, engine_kwargs=None):
         """Overridden to pass Redshift-specific arguments."""
@@ -237,7 +248,10 @@ class RedshiftSQLHook(DbApiHook):
             region_name = AwsBaseHook(aws_conn_id=self.aws_conn_id).region_name
             identifier = f"{cluster_identifier}.{region_name}"
         if not cluster_identifier:
-            identifier = self._get_identifier_from_hostname(connection.host)
+            if connection.host:
+                identifier = self._get_identifier_from_hostname(connection.host)
+            else:
+                raise AirflowException("Host is required when cluster_identifier is not provided.")
         return f"{identifier}:{port}"
 
     def _get_identifier_from_hostname(self, hostname: str) -> str:

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

@@ -43,15 +43,13 @@ from urllib.parse import urlsplit
 from uuid import uuid4
 
 if TYPE_CHECKING:
+    from aiobotocore.client import AioBaseClient
     from mypy_boto3_s3.service_resource import (
         Bucket as S3Bucket,
         Object as S3ResourceObject,
     )
 
-    from airflow.utils.types import ArgNotSet
-
-    with suppress(ImportError):
-        from aiobotocore.client import AioBaseClient
+    from airflow.providers.amazon.version_compat import ArgNotSet
 
 
 from asgiref.sync import sync_to_async
@@ -945,7 +943,14 @@ class S3Hook(AwsBaseHook):
             stacklevel=2,
         )
 
-        return list(self.iter_file_metadata(prefix=prefix, page_size=page_size, max_items=max_items))
+        return list(
+            self.iter_file_metadata(
+                prefix=prefix,
+                bucket_name=bucket_name,
+                page_size=page_size,
+                max_items=max_items,
+            )
+        )
 
     @provide_bucket_name
     def iter_file_metadata(
@@ -1774,6 +1779,8 @@ class S3Hook(AwsBaseHook):
         local_s3_objects = []
         s3_bucket = self.get_bucket(bucket_name)
         for obj in s3_bucket.objects.filter(Prefix=s3_prefix):
+            if obj.key.endswith("/"):
+                continue
             obj_path = Path(obj.key)
             local_target_path = local_dir.joinpath(obj_path.relative_to(s3_prefix))
             if not local_target_path.parent.exists():

airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py

@@ -26,7 +26,7 @@ from sagemaker_studio.sagemaker_studio_api import SageMakerStudioAPI
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.utils.sagemaker_unified_studio import is_local_runner
-from airflow.providers.amazon.version_compat import BaseHook
+from airflow.providers.common.compat.sdk import BaseHook
 
 
 class SageMakerNotebookHook(BaseHook):

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

@@ -20,7 +20,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.utils.types import NOTSET, ArgNotSet
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet, is_arg_set
 
 if TYPE_CHECKING:
     from airflow.sdk.execution_time.secrets_masker import mask_secret
@@ -38,7 +38,8 @@ class SsmHook(AwsBaseHook):
     """
     Interact with Amazon Systems Manager (SSM).
 
-    Provide thin wrapper around :external+boto3:py:class:`boto3.client("ssm") <SSM.Client>`.
+    Provide thin wrapper around
+    :external+boto3:py:class:`boto3.client("ssm") <SSM.Client>`.
 
     Additional arguments (such as ``aws_conn_id``) may be specified and
     are passed down to the underlying AwsBaseHook.
@@ -53,7 +54,9 @@ class SsmHook(AwsBaseHook):
 
     def get_parameter_value(self, parameter: str, default: str | ArgNotSet = NOTSET) -> str:
         """
-        Return the provided Parameter or an optional default; if it is encrypted, then decrypt and mask.
+        Return the provided Parameter or an optional default.
+
+        If it is encrypted, then decrypt and mask.
 
         .. seealso::
             - :external+boto3:py:meth:`SSM.Client.get_parameter`
@@ -68,6 +71,31 @@ class SsmHook(AwsBaseHook):
                 mask_secret(value)
             return value
         except self.conn.exceptions.ParameterNotFound:
-            if isinstance(default, ArgNotSet):
-                raise
-            return default
+            if is_arg_set(default):
+                return default
+            raise
+
+    def get_command_invocation(self, command_id: str, instance_id: str) -> dict:
+        """
+        Get the output of a command invocation for a specific instance.
+
+        .. seealso::
+            - :external+boto3:py:meth:`SSM.Client.get_command_invocation`
+
+        :param command_id: The ID of the command.
+        :param instance_id: The ID of the instance.
+        :return: The command invocation details including output.
+        """
+        return self.conn.get_command_invocation(CommandId=command_id, InstanceId=instance_id)
+
+    def list_command_invocations(self, command_id: str) -> dict:
+        """
+        List all command invocations for a given command ID.
+
+        .. seealso::
+            - :external+boto3:py:meth:`SSM.Client.list_command_invocations`
+
+        :param command_id: The ID of the command.
+        :return: Response from SSM list_command_invocations API.
+        """
+        return self.conn.list_command_invocations(CommandId=command_id)

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

@@ -18,8 +18,8 @@ from __future__ import annotations
 
 import json
 
-from airflow.exceptions import AirflowFailException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowFailException
 
 
 class StepFunctionHook(AwsBaseHook):

airflow/providers/amazon/aws/links/base_aws.py

@@ -20,7 +20,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, ClassVar
 
 from airflow.providers.amazon.aws.utils.suppress import return_on_error
-from airflow.providers.amazon.version_compat import BaseOperatorLink, XCom
+from airflow.providers.common.compat.sdk import BaseOperatorLink, XCom
 
 if TYPE_CHECKING:
     from airflow.models import BaseOperator

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

@@ -25,8 +25,8 @@ from airflow.providers.amazon.aws.utils.mixins import (
     AwsHookType,
     aws_template_fields,
 )
-from airflow.providers.amazon.version_compat import BaseOperator
-from airflow.utils.types import NOTSET, ArgNotSet
+from airflow.providers.amazon.version_compat import NOTSET, ArgNotSet
+from airflow.providers.common.compat.sdk import BaseOperator
 
 
 class AwsBaseOperator(BaseOperator, AwsBaseHookMixin[AwsHookType]):

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

@@ -482,6 +482,8 @@ class BedrockCreateKnowledgeBaseOperator(AwsBaseOperator[BedrockAgentHook]):
                     # It may also be that permissions haven't even propagated yet to check for the index
                     or "server returned 401" in error_message
                     or "user does not have permissions" in error_message
+                    or "status code: 403" in error_message
+                    or "bad authorization" in error_message
                 )
                 if all(
                     [

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

@@ -32,7 +32,7 @@ if TYPE_CHECKING:
 
 class CloudFormationCreateStackOperator(AwsBaseOperator[CloudFormationHook]):
     """
-    An operator that creates a AWS CloudFormation stack.
+    An operator that creates an AWS CloudFormation stack.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -68,7 +68,7 @@ class CloudFormationCreateStackOperator(AwsBaseOperator[CloudFormationHook]):
 
 class CloudFormationDeleteStackOperator(AwsBaseOperator[CloudFormationHook]):
     """
-    An operator that deletes a AWS CloudFormation stack.
+    An operator that deletes an AWS CloudFormation stack.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:

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

@@ -23,11 +23,12 @@ import random
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException, AirflowTaskTimeout
+from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.datasync import DataSyncHook
 from airflow.providers.amazon.aws.links.datasync import DataSyncTaskExecutionLink, DataSyncTaskLink
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowTaskTimeout
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context