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

airflow/providers/amazon/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.18.0"
+__version__ = "9.19.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.11.0"

airflow/providers/amazon/aws/auth_manager/avp/facade.py

@@ -23,7 +23,6 @@ from pathlib import Path
 from typing import TYPE_CHECKING, TypedDict
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities, get_action_id, get_entity_type
 from airflow.providers.amazon.aws.auth_manager.constants import (
     CONF_AVP_POLICY_STORE_ID_KEY,
@@ -32,6 +31,7 @@ from airflow.providers.amazon.aws.auth_manager.constants import (
     CONF_SECTION_NAME,
 )
 from airflow.providers.amazon.aws.hooks.verified_permissions import VerifiedPermissionsHook
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.utils.helpers import prune_dict
 from airflow.utils.log.logging_mixin import LoggingMixin
 

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

@@ -22,9 +22,9 @@ from pathlib import Path
 import structlog
 
 from airflow.dag_processing.bundles.base import BaseDagBundle
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.common.compat.sdk import AirflowException
 
 
 class S3DagBundle(BaseDagBundle):

airflow/providers/amazon/aws/exceptions.py

@@ -17,7 +17,7 @@
 # under the License.
 from __future__ import annotations
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 
 # Note: Any AirflowException raised is expected to cause the TaskInstance
 #       to be marked in an ERROR state

airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py

@@ -26,7 +26,6 @@ from boto3.session import NoCredentialsError
 from botocore.utils import ClientError
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.executors.base_executor import BaseExecutor
 from airflow.models.taskinstancekey import TaskInstanceKey
 from airflow.providers.amazon.aws.executors.aws_lambda.utils import (
@@ -42,7 +41,7 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
 )
 from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
-from airflow.stats import Stats
+from airflow.providers.common.compat.sdk import AirflowException, Stats
 
 try:
     from airflow.sdk import timezone

airflow/providers/amazon/aws/executors/batch/batch_executor.py

@@ -29,7 +29,6 @@ from typing import TYPE_CHECKING, Any
 from botocore.exceptions import ClientError, NoCredentialsError
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
 from airflow.executors.base_executor import BaseExecutor
 from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry import (
     calculate_next_attempt_delay,
@@ -37,7 +36,7 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
 )
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
 from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
-from airflow.stats import Stats
+from airflow.providers.common.compat.sdk import AirflowException, Stats
 
 try:
     from airflow.sdk import timezone

airflow/providers/amazon/aws/executors/ecs/ecs_executor.py

@@ -32,7 +32,6 @@ from typing import TYPE_CHECKING
 
 from botocore.exceptions import ClientError, NoCredentialsError
 
-from airflow.exceptions import AirflowException
 from airflow.executors.base_executor import BaseExecutor
 from airflow.providers.amazon.aws.executors.ecs.boto_schema import BotoDescribeTasksSchema, BotoRunTaskSchema
 from airflow.providers.amazon.aws.executors.ecs.utils import (
@@ -49,7 +48,7 @@ from airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
 )
 from airflow.providers.amazon.aws.hooks.ecs import EcsHook
 from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
-from airflow.stats import Stats
+from airflow.providers.common.compat.sdk import AirflowException, Stats
 
 try:
     from airflow.sdk import timezone

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

@@ -28,9 +28,9 @@ from __future__ import annotations
 from collections.abc import Collection
 from typing import TYPE_CHECKING, Any
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from botocore.paginate import PageIterator

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

@@ -23,9 +23,9 @@ from typing import TYPE_CHECKING, Any
 import pyathena
 from sqlalchemy.engine.url import URL
 
-from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 
 if TYPE_CHECKING:

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

@@ -52,15 +52,11 @@ from dateutil.tz import tzlocal
 from slugify import slugify
 
 from airflow.configuration import conf
-from airflow.exceptions import (
-    AirflowException,
-    AirflowNotFoundException,
-    AirflowProviderDeprecationWarning,
-)
+from airflow.exceptions import AirflowProviderDeprecationWarning
 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.common.compat.sdk import BaseHook
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException, BaseHook
 from airflow.providers_manager import ProvidersManager
 from airflow.utils.helpers import exactly_one
 from airflow.utils.log.logging_mixin import LoggingMixin

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

@@ -37,8 +37,8 @@ import botocore.client
 import botocore.exceptions
 import botocore.waiter
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher

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

@@ -36,8 +36,8 @@ from typing import TYPE_CHECKING, Any
 import botocore.exceptions
 import botocore.waiter
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
+from airflow.providers.common.compat.sdk import AirflowException
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher

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

@@ -25,7 +25,7 @@ import re
 from functools import cached_property
 from typing import Any
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.http.hooks.http import HttpHook
 
 

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

@@ -16,8 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowException
 
 
 class ComprehendHook(AwsBaseHook):

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

@@ -21,9 +21,9 @@ from __future__ import annotations
 import time
 from urllib.parse import urlsplit
 
-from airflow.exceptions import AirflowBadRequest, AirflowException
+from airflow.exceptions import AirflowBadRequest
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.common.compat.sdk import AirflowTaskTimeout
+from airflow.providers.common.compat.sdk import AirflowException, AirflowTaskTimeout
 
 
 class DataSyncHook(AwsBaseHook):

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

@@ -25,8 +25,8 @@ from typing import TYPE_CHECKING
 
 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
 
 if TYPE_CHECKING:
     from botocore.client import BaseClient

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

@@ -22,8 +22,8 @@ import time
 from collections.abc import Callable
 from typing import ParamSpec, TypeVar
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowException
 
 PS = ParamSpec("PS")
 RT = TypeVar("RT")

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

@@ -19,8 +19,8 @@ from __future__ import annotations
 
 import time
 
-from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowException
 
 
 class ElastiCacheReplicationGroupHook(AwsBaseHook):

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

@@ -26,9 +26,9 @@ import tenacity
 from botocore.exceptions import ClientError
 from tenacity import retry_if_exception, stop_after_attempt, wait_fixed
 
-from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
 
 
 class EmrHook(AwsBaseHook):

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

@@ -33,9 +33,10 @@ from tenacity import (
     wait_exponential,
 )
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.common.compat.sdk import AirflowException
 
 DEFAULT_LOG_SUFFIX = "output"
 ERROR_LOG_SUFFIX = "error"

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

@@ -23,8 +23,8 @@ from typing import Any
 
 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 GlueCatalogHook(AwsBaseHook):

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

@@ -19,7 +19,7 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import AsyncGenerator, Generator
-from typing import Any
+from typing import Any, TypedDict
 
 from botocore.exceptions import ClientError
 
@@ -35,6 +35,14 @@ from airflow.utils.helpers import prune_dict
 NUM_CONSECUTIVE_EMPTY_RESPONSE_EXIT_THRESHOLD = 3
 
 
+class CloudWatchLogEvent(TypedDict):
+    """TypedDict for CloudWatch Log Event."""
+
+    timestamp: int
+    message: str
+    ingestionTime: int
+
+
 class AwsLogsHook(AwsBaseHook):
     """
     Interact with Amazon CloudWatch Logs.
@@ -67,7 +75,7 @@ class AwsLogsHook(AwsBaseHook):
         start_from_head: bool | None = None,
         continuation_token: ContinuationToken | None = None,
         end_time: int | None = None,
-    ) -> Generator:
+    ) -> Generator[CloudWatchLogEvent, None, None]:
         """
         Return a generator for log items in a single stream; yields all items available at the current moment.
 

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:
@@ -206,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()

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

@@ -56,11 +56,12 @@ 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__)

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.common.compat.sdk import BaseHook
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
 
 
 class SageMakerNotebookHook(BaseHook):
@@ -124,7 +123,7 @@ class SageMakerNotebookHook(BaseHook):
         if self.compute:
             start_execution_params["compute"] = self.compute
         else:
-            start_execution_params["compute"] = {"instance_type": "ml.m4.xlarge"}
+            start_execution_params["compute"] = {"instance_type": "ml.m6i.xlarge"}
 
         print(start_execution_params)
         return self._sagemaker_studio.execution_client.start_execution(**start_execution_params)

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

@@ -25,7 +25,7 @@ from airflow.providers.common.compat.sdk import BaseOperatorLink, XCom
 if TYPE_CHECKING:
     from airflow.models import BaseOperator
     from airflow.models.taskinstancekey import TaskInstanceKey
-    from airflow.utils.context import Context
+    from airflow.sdk import Context
 
 
 BASE_AWS_CONSOLE_LINK = "https://console.{aws_domain}"

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/notifications/chime.py

@@ -24,7 +24,7 @@ from airflow.providers.amazon.aws.hooks.chime import ChimeWebhookHook
 from airflow.providers.common.compat.notifier import BaseNotifier
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.sdk import Context
 
 
 class ChimeNotifier(BaseNotifier):

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:
@@ -34,7 +34,7 @@ if TYPE_CHECKING:
         TaskTypeDef,
     )
 
-    from airflow.utils.context import Context
+    from airflow.sdk import Context
 
 SUPPORTED_SOURCES = {"salesforce", "zendesk"}
 MANDATORY_FILTER_DATE_MSG = "The filter_date argument is mandatory for {entity}!"

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

@@ -22,18 +22,18 @@ 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
     from airflow.providers.openlineage.extractors.base import OperatorLineage
-    from airflow.utils.context import Context
+    from airflow.sdk import Context
 
 
 class AthenaOperator(AwsBaseOperator[AthenaHook]):

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,9 +46,10 @@ 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
+    from airflow.sdk import Context
 
 
 class BatchOperator(AwsBaseOperator[BatchClientHook]):