apache-airflow-providers-amazon 9.15.0__py3-none-any.whl → 9.18.1rc1__py3-none-any.whl

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/quicksight.py

@@ -21,8 +21,8 @@ import time
 
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowException
 
 
 class QuickSightHook(AwsBaseHook):

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

@@ -23,9 +23,9 @@ import time
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
 
 if TYPE_CHECKING:
     from mypy_boto3_rds import RDSClient  # noqa: F401

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

@@ -20,12 +20,13 @@ from functools import cached_property
 from typing import TYPE_CHECKING
 
 import redshift_connector
-from redshift_connector import Connection as RedshiftConnection
+import tenacity
+from redshift_connector import Connection as RedshiftConnection, InterfaceError, OperationalError
 from sqlalchemy import create_engine
 from sqlalchemy.engine.url import URL
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 
 if TYPE_CHECKING:
@@ -51,7 +52,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 +156,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."""
@@ -195,6 +207,14 @@ class RedshiftSQLHook(DbApiHook):
         pk_columns = [row[0] for row in self.get_records(sql, (schema, table))]
         return pk_columns or None
 
+    @tenacity.retry(
+        stop=tenacity.stop_after_attempt(5),
+        wait=tenacity.wait_exponential(max=20),
+        # OperationalError is thrown when the connection times out
+        # InterfaceError is thrown when the connection is refused
+        retry=tenacity.retry_if_exception_type((OperationalError, InterfaceError)),
+        reraise=True,
+    )
     def get_conn(self) -> RedshiftConnection:
         """Get a ``redshift_connector.Connection`` object."""
         conn_params = self._get_conn_params()
@@ -237,7 +257,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,26 +43,25 @@ 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
 from boto3.s3.transfer import S3Transfer, TransferConfig
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException, AirflowNotFoundException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.exceptions import S3HookUriParseFailure
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
 from airflow.utils.helpers import chunks
 
 logger = logging.getLogger(__name__)
@@ -945,7 +944,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 +1780,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.py

@@ -31,11 +31,11 @@ from typing import Any, cast
 from asgiref.sync import sync_to_async
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.tags import format_tags
+from airflow.providers.common.compat.sdk import AirflowException
 
 try:
     from airflow.sdk import timezone

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

@@ -24,9 +24,8 @@ import time
 from sagemaker_studio import ClientConfig
 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 AirflowException, 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/links/emr.py

@@ -19,10 +19,10 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any
 from urllib.parse import ParseResult, quote_plus, urlparse
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.emr import EmrServerlessHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.utils.helpers import exactly_one
 
 if TYPE_CHECKING:

airflow/providers/amazon/aws/log/cloudwatch_task_handler.py

@@ -22,6 +22,7 @@ import copy
 import json
 import logging
 import os
+from collections.abc import Generator
 from datetime import date, datetime, timedelta, timezone
 from functools import cached_property
 from pathlib import Path
@@ -40,8 +41,15 @@ if TYPE_CHECKING:
     import structlog.typing
 
     from airflow.models.taskinstance import TaskInstance
+    from airflow.providers.amazon.aws.hooks.logs import CloudWatchLogEvent
     from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
-    from airflow.utils.log.file_task_handler import LogMessages, LogSourceInfo
+    from airflow.utils.log.file_task_handler import (
+        LogMessages,
+        LogResponse,
+        LogSourceInfo,
+        RawLogStream,
+        StreamingLogResponse,
+    )
 
 
 def json_serialize_legacy(value: Any) -> str | None:
@@ -163,20 +171,31 @@ class CloudWatchRemoteLogIO(LoggingMixin):  # noqa: D101
         self.close()
         return
 
-    def read(self, relative_path, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages | None]:
-        logs: LogMessages | None = []
+    def read(self, relative_path: str, ti: RuntimeTI) -> LogResponse:
+        messages, logs = self.stream(relative_path, ti)
+        str_logs: list[str] = [f"{msg}\n" for group in logs for msg in group]
+
+        return messages, str_logs
+
+    def stream(self, relative_path: str, ti: RuntimeTI) -> StreamingLogResponse:
+        logs: list[RawLogStream] = []
         messages = [
             f"Reading remote log from Cloudwatch log_group: {self.log_group} log_stream: {relative_path}"
         ]
         try:
-            logs = [self.get_cloudwatch_logs(relative_path, ti)]
+            gen: RawLogStream = (
+                self._parse_log_event_as_dumped_json(event)
+                for event in self.get_cloudwatch_logs(relative_path, ti)
+            )
+            logs = [gen]
         except Exception as e:
-            logs = None
             messages.append(str(e))
 
         return messages, logs
 
-    def get_cloudwatch_logs(self, stream_name: str, task_instance: RuntimeTI):
+    def get_cloudwatch_logs(
+        self, stream_name: str, task_instance: RuntimeTI
+    ) -> Generator[CloudWatchLogEvent, None, None]:
         """
         Return all logs from the given log stream.
 
@@ -192,29 +211,22 @@ class CloudWatchRemoteLogIO(LoggingMixin):  # noqa: D101
             if (end_date := getattr(task_instance, "end_date", None)) is None
             else datetime_to_epoch_utc_ms(end_date + timedelta(seconds=30))
         )
-        events = self.hook.get_log_events(
+        return self.hook.get_log_events(
             log_group=self.log_group,
             log_stream_name=stream_name,
             end_time=end_time,
         )
-        return "\n".join(self._event_to_str(event) for event in events)
 
-    def _event_to_dict(self, event: dict) -> dict:
+    def _parse_log_event_as_dumped_json(self, event: CloudWatchLogEvent) -> str:
         event_dt = datetime.fromtimestamp(event["timestamp"] / 1000.0, tz=timezone.utc).isoformat()
-        message = event["message"]
+        event_msg = event["message"]
         try:
-            message = json.loads(message)
+            message = json.loads(event_msg)
             message["timestamp"] = event_dt
-            return message
         except Exception:
-            return {"timestamp": event_dt, "event": message}
+            message = {"timestamp": event_dt, "event": event_msg}
 
-    def _event_to_str(self, event: dict) -> str:
-        event_dt = datetime.fromtimestamp(event["timestamp"] / 1000.0, tz=timezone.utc)
-        # Format a datetime object to a string in Zulu time without milliseconds.
-        formatted_event_dt = event_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
-        message = event["message"]
-        return f"[{formatted_event_dt}] {message}"
+        return json.dumps(message)
 
 
 class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin):
@@ -291,4 +303,22 @@ class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin):
     ) -> tuple[LogSourceInfo, LogMessages]:
         stream_name = self._render_filename(task_instance, try_number)
         messages, logs = self.io.read(stream_name, task_instance)
-        return messages, logs or []
+
+        messages = [
+            f"Reading remote log from Cloudwatch log_group: {self.io.log_group} log_stream: {stream_name}"
+        ]
+        try:
+            events = self.io.get_cloudwatch_logs(stream_name, task_instance)
+            logs = ["\n".join(self._event_to_str(event) for event in events)]
+        except Exception as e:
+            logs = []
+            messages.append(str(e))
+
+        return messages, logs
+
+    def _event_to_str(self, event: CloudWatchLogEvent) -> str:
+        event_dt = datetime.fromtimestamp(event["timestamp"] / 1000.0, tz=timezone.utc)
+        # Format a datetime object to a string in Zulu time without milliseconds.
+        formatted_event_dt = event_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
+        message = event["message"]
+        return f"[{formatted_event_dt}] {message}"

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

@@ -20,11 +20,11 @@ import time
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, cast
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.appflow import AppflowHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.utils import datetime_to_epoch_ms
 from airflow.providers.amazon.aws.utils.mixins import AwsBaseHookMixin, AwsHookParams, aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.common.compat.standard.operators import ShortCircuitOperator
 
 if TYPE_CHECKING:

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

@@ -22,13 +22,13 @@ from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.athena import AthenaHook
 from airflow.providers.amazon.aws.links.athena import AthenaQueryResultsLink
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.athena import AthenaTrigger
 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
 
 if TYPE_CHECKING:
     from airflow.providers.common.compat.openlineage.facet import BaseFacet, Dataset, DatasetFacet

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/batch.py

@@ -31,7 +31,6 @@ from datetime import timedelta
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
 from airflow.providers.amazon.aws.links.batch import (
     BatchJobDefinitionLink,
@@ -47,6 +46,7 @@ from airflow.providers.amazon.aws.triggers.batch import (
 from airflow.providers.amazon.aws.utils import trim_none_values, validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Any
 from botocore.exceptions import ClientError
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.bedrock import (
     BedrockAgentHook,
     BedrockAgentRuntimeHook,
@@ -41,6 +40,7 @@ from airflow.providers.amazon.aws.triggers.bedrock 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.utils.helpers import prune_dict
 from airflow.utils.timezone import utcnow
 
@@ -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/comprehend.py

@@ -21,7 +21,6 @@ from functools import cached_property
 from typing import TYPE_CHECKING, Any, ClassVar
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.comprehend import ComprehendHook
 from airflow.providers.amazon.aws.links.comprehend import (
     ComprehendDocumentClassifierLink,
@@ -34,6 +33,7 @@ from airflow.providers.amazon.aws.triggers.comprehend 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.utils.timezone import utcnow
 
 if TYPE_CHECKING:

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

@@ -23,11 +23,11 @@ import random
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException, AirflowTaskTimeout
 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 AirflowException, AirflowTaskTimeout
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -22,7 +22,6 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, ClassVar
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.dms import DmsHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.aws.triggers.dms import (
@@ -33,6 +32,7 @@ from airflow.providers.amazon.aws.triggers.dms import (
     DmsReplicationTerminalStatusTrigger,
 )
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.utils.context import Context
 
 if TYPE_CHECKING: