apache-airflow-providers-amazon 8.6.0__py3-none-any.whl → 8.7.0__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -28,7 +28,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.6.0"
+__version__ = "8.7.0"
 
 try:
     from airflow import __version__ as airflow_version

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

@@ -24,7 +24,6 @@ from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
 
 if TYPE_CHECKING:
     from mypy_boto3_appflow.client import AppflowClient
-    from mypy_boto3_appflow.type_defs import TaskOutputTypeDef, TaskTypeDef
 
 
 class AppflowHook(AwsBaseHook):
@@ -93,9 +92,7 @@ class AppflowHook(AwsBaseHook):
         exec_details = last_execs[execution_id]
         self.log.info("Run complete, execution details: %s", exec_details)
 
-    def update_flow_filter(
-        self, flow_name: str, filter_tasks: list[TaskTypeDef], set_trigger_ondemand: bool = False
-    ) -> None:
+    def update_flow_filter(self, flow_name: str, filter_tasks, set_trigger_ondemand: bool = False) -> None:
         """
         Update the flow task filter; all filters will be removed if an empty array is passed to filter_tasks.
 
@@ -106,7 +103,7 @@ class AppflowHook(AwsBaseHook):
         """
         response = self.conn.describe_flow(flowName=flow_name)
         connector_type = response["sourceFlowConfig"]["connectorType"]
-        tasks: list[TaskTypeDef | TaskOutputTypeDef] = []
+        tasks = []
 
         # cleanup old filter tasks
         for task in response["tasks"]:

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

@@ -25,14 +25,15 @@ This module contains AWS Athena hook.
 from __future__ import annotations
 
 import warnings
-from typing import Any
-
-from botocore.paginate import PageIterator
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
 
+if TYPE_CHECKING:
+    from botocore.paginate import PageIterator
+
 
 class AthenaHook(AwsBaseHook):
     """Interact with Amazon Athena.

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

@@ -29,11 +29,9 @@ import inspect
 import json
 import logging
 import os
-import uuid
 import warnings
 from copy import deepcopy
 from functools import cached_property, wraps
-from os import PathLike
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union
 
@@ -43,9 +41,7 @@ import botocore.session
 import jinja2
 import requests
 import tenacity
-from botocore.client import ClientMeta
 from botocore.config import Config
-from botocore.credentials import ReadOnlyCredentials
 from botocore.waiter import Waiter, WaiterModel
 from dateutil.tz import tzlocal
 from slugify import slugify
@@ -58,6 +54,8 @@ from airflow.exceptions import (
 )
 from airflow.hooks.base import BaseHook
 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_manager import ProvidersManager
 from airflow.utils.helpers import exactly_one
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -66,6 +64,9 @@ from airflow.utils.log.secrets_masker import mask_secret
 BaseAwsConnection = TypeVar("BaseAwsConnection", bound=Union[boto3.client, boto3.resource])
 
 if TYPE_CHECKING:
+    from botocore.client import ClientMeta
+    from botocore.credentials import ReadOnlyCredentials
+
     from airflow.models.connection import Connection  # Avoid circular imports.
 
 
@@ -470,21 +471,17 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         self._verify = verify
 
     @classmethod
+    @return_on_error("Unknown")
     def _get_provider_version(cls) -> str:
         """Check the Providers Manager for the package version."""
-        try:
-            manager = ProvidersManager()
-            hook = manager.hooks[cls.conn_type]
-            if not hook:
-                # This gets caught immediately, but without it MyPy complains
-                # Item "None" of "Optional[HookInfo]" has no attribute "package_name"
-                # on the following line and static checks fail.
-                raise ValueError(f"Hook info for {cls.conn_type} not found in the Provider Manager.")
-            provider = manager.providers[hook.package_name]
-            return provider.version
-        except Exception:
-            # Under no condition should an error here ever cause an issue for the user.
-            return "Unknown"
+        manager = ProvidersManager()
+        hook = manager.hooks[cls.conn_type]
+        if not hook:
+            # This gets caught immediately, but without it MyPy complains
+            # Item "None" of "Optional[HookInfo]" has no attribute "package_name"
+            # on the following line and static checks fail.
+            raise ValueError(f"Hook info for {cls.conn_type} not found in the Provider Manager.")
+        return manager.providers[hook.package_name].version
 
     @staticmethod
     def _find_class_name(target_function_name: str) -> str:
@@ -504,19 +501,17 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         # Return the name of the class object.
         return frame_class_object.__name__
 
+    @return_on_error("Unknown")
     def _get_caller(self, target_function_name: str = "execute") -> str:
         """Given a function name, walk the stack and return the name of the class which called it last."""
-        try:
-            caller = self._find_class_name(target_function_name)
-            if caller == "BaseSensorOperator":
-                # If the result is a BaseSensorOperator, then look for whatever last called "poke".
-                return self._get_caller("poke")
-            return caller
-        except Exception:
-            # Under no condition should an error here ever cause an issue for the user.
-            return "Unknown"
+        caller = self._find_class_name(target_function_name)
+        if caller == "BaseSensorOperator":
+            # If the result is a BaseSensorOperator, then look for whatever last called "poke".
+            return self._get_caller("poke")
+        return caller
 
     @staticmethod
+    @return_on_error("00000000-0000-0000-0000-000000000000")
     def _generate_dag_key() -> str:
         """Generate a DAG key.
 
@@ -525,25 +520,17 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         can not (reasonably) be reversed.  No personal data can be inferred or
         extracted from the resulting UUID.
         """
-        try:
-            dag_id = os.environ["AIRFLOW_CTX_DAG_ID"]
-            return str(uuid.uuid5(uuid.NAMESPACE_OID, dag_id))
-        except Exception:
-            # Under no condition should an error here ever cause an issue for the user.
-            return "00000000-0000-0000-0000-000000000000"
+        return generate_uuid(os.environ.get("AIRFLOW_CTX_DAG_ID"))
 
     @staticmethod
+    @return_on_error("Unknown")
     def _get_airflow_version() -> str:
         """Fetch and return the current Airflow version."""
-        try:
-            # This can be a circular import under specific configurations.
-            # Importing locally to either avoid or catch it if it does happen.
-            from airflow import __version__ as airflow_version
+        # This can be a circular import under specific configurations.
+        # Importing locally to either avoid or catch it if it does happen.
+        from airflow import __version__ as airflow_version
 
-            return airflow_version
-        except Exception:
-            # Under no condition should an error here ever cause an issue for the user.
-            return "Unknown"
+        return airflow_version
 
     def _generate_user_agent_extra_field(self, existing_user_agent_extra: str) -> str:
         user_agent_extra_values = [
@@ -831,7 +818,7 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
             return False, str(f"{type(e).__name__!r} error occurred while testing connection: {e}")
 
     @cached_property
-    def waiter_path(self) -> PathLike[str] | None:
+    def waiter_path(self) -> os.PathLike[str] | None:
         filename = self.client_type if self.client_type else self.resource_type
         path = Path(__file__).parents[1].joinpath(f"waiters/{filename}.json").resolve()
         return path if path.exists() else None

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

@@ -27,9 +27,9 @@ A client for AWS Batch services.
 from __future__ import annotations
 
 import itertools
-from random import uniform
+import random
 from time import sleep
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import botocore.client
 import botocore.exceptions
@@ -37,9 +37,11 @@ import botocore.waiter
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
 from airflow.typing_compat import Protocol, runtime_checkable
 
+if TYPE_CHECKING:
+    from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
+
 
 @runtime_checkable
 class BatchProtocol(Protocol):
@@ -527,7 +529,7 @@ class BatchClientHook(AwsBaseHook):
         minima = abs(minima)
         lower = max(minima, delay - width)
         upper = delay + width
-        return uniform(lower, upper)
+        return random.uniform(lower, upper)
 
     @staticmethod
     def delay(delay: int | float | None = None) -> None:
@@ -544,7 +546,7 @@ class BatchClientHook(AwsBaseHook):
             when many concurrent tasks request job-descriptions.
         """
         if delay is None:
-            delay = uniform(BatchClientHook.DEFAULT_DELAY_MIN, BatchClientHook.DEFAULT_DELAY_MAX)
+            delay = random.uniform(BatchClientHook.DEFAULT_DELAY_MIN, BatchClientHook.DEFAULT_DELAY_MAX)
         else:
             delay = BatchClientHook.add_jitter(delay)
         sleep(delay)
@@ -592,4 +594,4 @@ class BatchClientHook(AwsBaseHook):
         max_interval = 600.0  # results in 3 to 10 minute delay
         delay = 1 + pow(tries * 0.6, 2)
         delay = min(max_interval, delay)
-        return uniform(delay / 3, delay)
+        return random.uniform(delay / 3, delay)

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

@@ -29,7 +29,7 @@ import json
 import sys
 from copy import deepcopy
 from pathlib import Path
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import botocore.client
 import botocore.exceptions
@@ -37,7 +37,9 @@ import botocore.waiter
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
-from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
+
+if TYPE_CHECKING:
+    from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
 
 
 class BatchWaitersHook(BatchClientHook):

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

@@ -21,6 +21,7 @@ from __future__ import annotations
 
 import json
 import re
+from functools import cached_property
 from typing import Any
 
 from airflow.exceptions import AirflowException
@@ -28,19 +29,19 @@ from airflow.providers.http.hooks.http import HttpHook
 
 
 class ChimeWebhookHook(HttpHook):
-    """Interact with Chime Web Hooks to create notifications.
+    """Interact with Amazon Chime Webhooks to create notifications.
 
     .. warning:: This hook is only designed to work with web hooks and not chat bots.
 
-    :param chime_conn_id: Chime connection ID with Endpoint as "https://hooks.chime.aws" and
-                         the webhook token in the form of ```{webhook.id}?token{webhook.token}```
-
+    :param chime_conn_id: :ref:`Amazon Chime Connection ID <howto/connection:chime>`
+        with Endpoint as `https://hooks.chime.aws` and the webhook token
+        in the form of ``{webhook.id}?token{webhook.token}``
     """
 
     conn_name_attr = "chime_conn_id"
     default_conn_name = "chime_default"
     conn_type = "chime"
-    hook_name = "Chime Web Hook"
+    hook_name = "Amazon Chime Webhook"
 
     def __init__(
         self,
@@ -49,7 +50,11 @@ class ChimeWebhookHook(HttpHook):
         **kwargs: Any,
     ) -> None:
         super().__init__(*args, **kwargs)
-        self.webhook_endpoint = self._get_webhook_endpoint(chime_conn_id)
+        self._chime_conn_id = chime_conn_id
+
+    @cached_property
+    def webhook_endpoint(self):
+        return self._get_webhook_endpoint(self._chime_conn_id)
 
     def _get_webhook_endpoint(self, conn_id: str) -> str:
         """
@@ -65,7 +70,7 @@ class ChimeWebhookHook(HttpHook):
         url = conn.schema + "://" + conn.host
         endpoint = url + token
         # Check to make sure the endpoint matches what Chime expects
-        if not re.match(r"^[a-zA-Z0-9_-]+\?token=[a-zA-Z0-9_-]+$", token):
+        if not re.fullmatch(r"[a-zA-Z0-9_-]+\?token=[a-zA-Z0-9_-]+", token):
             raise AirflowException(
                 "Expected Chime webhook token in the form of '{webhook.id}?token={webhook.token}'."
             )
@@ -104,7 +109,7 @@ class ChimeWebhookHook(HttpHook):
             "hidden_fields": ["login", "port", "extra"],
             "relabeling": {
                 "host": "Chime Webhook Endpoint",
-                "password": "Webhook Token",
+                "password": "Chime Webhook token",
             },
             "placeholders": {
                 "schema": "https",

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

@@ -18,11 +18,15 @@
 """This module contains AWS CloudFormation Hook."""
 from __future__ import annotations
 
-from boto3 import client, resource
+from typing import TYPE_CHECKING
+
 from botocore.exceptions import ClientError
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
+if TYPE_CHECKING:
+    from boto3 import client, resource
+
 
 class CloudFormationHook(AwsBaseHook):
     """

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

@@ -301,25 +301,18 @@ class DataSyncHook(AwsBaseHook):
         if not task_execution_arn:
             raise AirflowBadRequest("task_execution_arn not specified")
 
-        status = None
-        iterations = max_iterations
-        while status is None or status in self.TASK_EXECUTION_INTERMEDIATE_STATES:
+        for _ in range(max_iterations):
             task_execution = self.get_conn().describe_task_execution(TaskExecutionArn=task_execution_arn)
             status = task_execution["Status"]
             self.log.info("status=%s", status)
-            iterations -= 1
-            if status in self.TASK_EXECUTION_FAILURE_STATES:
-                break
             if status in self.TASK_EXECUTION_SUCCESS_STATES:
-                break
-            if iterations <= 0:
-                break
+                return True
+            elif status in self.TASK_EXECUTION_FAILURE_STATES:
+                return False
+            elif status is None or status in self.TASK_EXECUTION_INTERMEDIATE_STATES:
+                time.sleep(self.wait_interval_seconds)
+            else:
+                raise AirflowException(f"Unknown status: {status}")  # Should never happen
             time.sleep(self.wait_interval_seconds)
-
-        if status in self.TASK_EXECUTION_SUCCESS_STATES:
-            return True
-        if status in self.TASK_EXECUTION_FAILURE_STATES:
-            return False
-        if iterations <= 0:
+        else:
             raise AirflowTaskTimeout("Max iterations exceeded!")
-        raise AirflowException(f"Unknown status: {status}")  # Should never happen

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

@@ -20,11 +20,14 @@ from __future__ import annotations
 import base64
 import logging
 from dataclasses import dataclass
-from datetime import datetime
+from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.utils.log.secrets_masker import mask_secret
 
+if TYPE_CHECKING:
+    from datetime import datetime
+
 logger = logging.getLogger(__name__)
 
 

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

@@ -17,13 +17,16 @@
 # under the License.
 from __future__ import annotations
 
-from botocore.waiter import Waiter
+from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.exceptions import EcsOperatorError, EcsTaskFailToStart
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
 from airflow.providers.amazon.aws.utils import _StringCompareEnum
 from airflow.typing_compat import Protocol, runtime_checkable
 
+if TYPE_CHECKING:
+    from botocore.waiter import Waiter
+
 
 def should_retry(exception: Exception):
     """Check if exception is related to ECS resource quota (CPU, MEM)."""

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

@@ -21,7 +21,6 @@ import warnings
 from typing import Any, Sequence
 
 import botocore.exceptions
-from botocore.exceptions import ClientError
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseAsyncHook, AwsBaseHook
@@ -70,17 +69,14 @@ class RedshiftHook(AwsBaseHook):
             for the cluster that is being created.
         :param params: Remaining AWS Create cluster API params.
         """
-        try:
-            response = self.get_conn().create_cluster(
-                ClusterIdentifier=cluster_identifier,
-                NodeType=node_type,
-                MasterUsername=master_username,
-                MasterUserPassword=master_user_password,
-                **params,
-            )
-            return response
-        except ClientError as e:
-            raise e
+        response = self.get_conn().create_cluster(
+            ClusterIdentifier=cluster_identifier,
+            NodeType=node_type,
+            MasterUsername=master_username,
+            MasterUserPassword=master_user_password,
+            **params,
+        )
+        return response
 
     # TODO: Wrap create_cluster_snapshot
     def cluster_status(self, cluster_identifier: str) -> str:

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

@@ -188,7 +188,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
         pk_columns = []
         token = ""
         while True:
-            kwargs = dict(Id=stmt_id)
+            kwargs = {"Id": stmt_id}
             if token:
                 kwargs["NextToken"] = token
             response = self.conn.get_statement_result(**kwargs)

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

@@ -41,10 +41,8 @@ from urllib.parse import urlsplit
 from uuid import uuid4
 
 if TYPE_CHECKING:
-    try:
+    with suppress(ImportError):
         from aiobotocore.client import AioBaseClient
-    except ImportError:
-        pass
 
 from asgiref.sync import sync_to_async
 from boto3.s3.transfer import TransferConfig
@@ -470,7 +468,7 @@ class S3Hook(AwsBaseHook):
         :param bucket_name: the name of the bucket
         :param key: the path to the key
         """
-        prefix = re.split(r"[\[\*\?]", key, 1)[0]
+        prefix = re.split(r"[\[*?]", key, 1)[0]
         delimiter = ""
         paginator = client.get_paginator("list_objects_v2")
         response = paginator.paginate(Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter)
@@ -572,7 +570,7 @@ class S3Hook(AwsBaseHook):
         for key in bucket_keys:
             prefix = key
             if wildcard_match:
-                prefix = re.split(r"[\[\*\?]", key, 1)[0]
+                prefix = re.split(r"[\[*?]", key, 1)[0]
 
             paginator = client.get_paginator("list_objects_v2")
             response = paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter=delimiter)
@@ -1017,7 +1015,7 @@ class S3Hook(AwsBaseHook):
         :param delimiter: the delimiter marks key hierarchy
         :return: the key object from the bucket or None if none has been found.
         """
-        prefix = re.split(r"[\[\*\?]", wildcard_key, 1)[0]
+        prefix = re.split(r"[\[*?]", wildcard_key, 1)[0]
         key_list = self.list_keys(bucket_name, prefix=prefix, delimiter=delimiter)
         key_matches = [k for k in key_list if fnmatch.fnmatch(k, wildcard_key)]
         if key_matches:

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

@@ -252,12 +252,12 @@ class SageMakerHook(AwsBaseHook):
         ]
         events: list[Any | None] = []
         for event_stream in event_iters:
-            if not event_stream:
-                events.append(None)
-                continue
-            try:
-                events.append(next(event_stream))
-            except StopIteration:
+            if event_stream:
+                try:
+                    events.append(next(event_stream))
+                except StopIteration:
+                    events.append(None)
+            else:
                 events.append(None)
 
         while any(events):
@@ -979,8 +979,7 @@ class SageMakerHook(AwsBaseHook):
         found_name: str,
         job_name_suffix: str | None = None,
     ) -> bool:
-        pattern = re.compile(f"^{processing_job_name}({job_name_suffix})?$")
-        return pattern.fullmatch(found_name) is not None
+        return re.fullmatch(f"{processing_job_name}({job_name_suffix})?", found_name) is not None
 
     def count_processing_jobs_by_name(
         self,

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

@@ -68,7 +68,6 @@ class SnsHook(AwsBaseHook):
 
         :param target_arn: either a TopicArn or an EndpointArn
         :param message: the default message you want to send
-        :param message: str
         :param subject: subject of message
         :param message_attributes: additional attributes to publish for message filtering. This should be
             a flat dict; the DataType to be sent depends on the type of the value:

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

@@ -16,15 +16,16 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Any
-
-import boto3
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 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.utils.helpers import exactly_one
 
+if TYPE_CHECKING:
+    import boto3
+
 
 class EmrClusterLink(BaseAwsLink):
     """Helper class for constructing AWS EMR Cluster Link."""

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

@@ -19,16 +19,19 @@ from __future__ import annotations
 
 from datetime import datetime, timedelta
 from functools import cached_property
+from typing import TYPE_CHECKING
 
 import watchtower
 
 from airflow.configuration import conf
-from airflow.models import TaskInstance
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.utils import datetime_to_epoch_utc_ms
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 
+if TYPE_CHECKING:
+    from airflow.models import TaskInstance
+
 
 class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin):
     """

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

@@ -201,7 +201,7 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
         try:
             if append and self.s3_log_exists(remote_log_location):
                 old_log = self.s3_read(remote_log_location)
-                log = "\n".join([old_log, log]) if old_log else log
+                log = f"{old_log}\n{log}" if old_log else log
         except Exception:
             self.log.exception("Could not verify previous log to append")
             return False

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

@@ -18,10 +18,13 @@
 from __future__ import annotations
 
 from functools import cached_property
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.providers.amazon.aws.hooks.chime import ChimeWebhookHook
-from airflow.utils.context import Context
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
 
 try:
     from airflow.notifications.basenotifier import BaseNotifier