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

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

@@ -28,7 +28,8 @@ import os
 import re
 import shutil
 import time
-from collections.abc import AsyncIterator, Callable
+import warnings
+from collections.abc import AsyncIterator, Callable, Iterator
 from contextlib import suppress
 from copy import deepcopy
 from datetime import datetime
@@ -42,22 +43,20 @@ 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
+from airflow.exceptions import AirflowException, AirflowNotFoundException, 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
@@ -931,7 +930,38 @@ class S3Hook(AwsBaseHook):
         max_items: int | None = None,
     ) -> list:
         """
-        List metadata objects in a bucket under prefix.
+        .. deprecated:: <9.13.0> Use `iter_file_metadata` instead.
+
+        This method `get_file_metadata` is deprecated. Calling this method will result in all matching keys
+        being loaded into a single list, and can often result in out-of-memory exceptions.
+        """
+        warnings.warn(
+            "This method `get_file_metadata` is deprecated. Calling this method will result in all matching "
+            "keys being loaded into a single list, and can often result in out-of-memory exceptions. "
+            "Instead, use `iter_file_metadata`.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+
+        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(
+        self,
+        prefix: str,
+        bucket_name: str | None = None,
+        page_size: int | None = None,
+        max_items: int | None = None,
+    ) -> Iterator:
+        """
+        Yield metadata objects from a bucket under a prefix.
 
         .. seealso::
             - :external+boto3:py:class:`S3.Paginator.ListObjectsV2`
@@ -940,7 +970,7 @@ class S3Hook(AwsBaseHook):
         :param bucket_name: the name of the bucket
         :param page_size: pagination size
         :param max_items: maximum items to return
-        :return: a list of metadata of objects
+        :return: an Iterator of metadata of objects
         """
         config = {
             "PageSize": page_size,
@@ -957,11 +987,9 @@ class S3Hook(AwsBaseHook):
             params["RequestPayer"] = "requester"
         response = paginator.paginate(**params)
 
-        files = []
         for page in response:
             if "Contents" in page:
-                files += page["Contents"]
-        return files
+                yield from page["Contents"]
 
     @unify_bucket_name_and_key
     @provide_bucket_name
@@ -1751,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/ses.py

@@ -42,6 +42,20 @@ class SesHook(AwsBaseHook):
         kwargs["client_type"] = "ses"
         super().__init__(*args, **kwargs)
 
+    @staticmethod
+    def _build_headers(
+        custom_headers: dict[str, Any] | None,
+        reply_to: str | None,
+        return_path: str | None,
+    ) -> dict[str, Any]:
+        _custom_headers = custom_headers or {}
+        if reply_to:
+            _custom_headers["Reply-To"] = reply_to
+        if return_path:
+            _custom_headers["Return-Path"] = return_path
+
+        return _custom_headers
+
     def send_email(
         self,
         mail_from: str,
@@ -70,23 +84,17 @@ class SesHook(AwsBaseHook):
         :param files: List of paths of files to be attached
         :param cc: List of email addresses to set as email's CC
         :param bcc: List of email addresses to set as email's BCC
-        :param mime_subtype: Can be used to specify the sub-type of the message. Default = mixed
+        :param mime_subtype: Can be used to specify the subtype of the message. Default = mixed
         :param mime_charset: Email's charset. Default = UTF-8.
         :param return_path: The email address to which replies will be sent. By default, replies
             are sent to the original sender's email address.
         :param reply_to: The email address to which message bounces and complaints should be sent.
             "Return-Path" is sometimes called "envelope from", "envelope sender", or "MAIL FROM".
         :param custom_headers: Additional headers to add to the MIME message.
-            No validations are run on these values and they should be able to be encoded.
+            No validations are run on these values, and they should be able to be encoded.
         :return: Response from Amazon SES service with unique message identifier.
         """
-        ses_client = self.get_conn()
-
-        custom_headers = custom_headers or {}
-        if reply_to:
-            custom_headers["Reply-To"] = reply_to
-        if return_path:
-            custom_headers["Return-Path"] = return_path
+        custom_headers = self._build_headers(custom_headers, reply_to, return_path)
 
         message, recipients = build_mime_message(
             mail_from=mail_from,
@@ -101,6 +109,64 @@ class SesHook(AwsBaseHook):
             custom_headers=custom_headers,
         )
 
-        return ses_client.send_raw_email(
+        return self.conn.send_raw_email(
             Source=mail_from, Destinations=recipients, RawMessage={"Data": message.as_string()}
         )
+
+    async def asend_email(
+        self,
+        mail_from: str,
+        to: str | Iterable[str],
+        subject: str,
+        html_content: str,
+        files: list[str] | None = None,
+        cc: str | Iterable[str] | None = None,
+        bcc: str | Iterable[str] | None = None,
+        mime_subtype: str = "mixed",
+        mime_charset: str = "utf-8",
+        reply_to: str | None = None,
+        return_path: str | None = None,
+        custom_headers: dict[str, Any] | None = None,
+    ) -> dict:
+        """
+        Send email using Amazon Simple Email Service (async).
+
+        .. seealso::
+            - :external+boto3:py:meth:`SES.Client.send_raw_email`
+
+        :param mail_from: Email address to set as email's from
+        :param to: List of email addresses to set as email's to
+        :param subject: Email's subject
+        :param html_content: Content of email in HTML format
+        :param files: List of paths of files to be attached
+        :param cc: List of email addresses to set as email's CC
+        :param bcc: List of email addresses to set as email's BCC
+        :param mime_subtype: Can be used to specify the subtype of the message. Default = mixed
+        :param mime_charset: Email's charset. Default = UTF-8.
+        :param return_path: The email address to which replies will be sent. By default, replies
+            are sent to the original sender's email address.
+        :param reply_to: The email address to which message bounces and complaints should be sent.
+            "Return-Path" is sometimes called "envelope from", "envelope sender", or "MAIL FROM".
+        :param custom_headers: Additional headers to add to the MIME message.
+            No validations are run on these values, and they should be able to be encoded.
+        :return: Response from Amazon SES service with unique message identifier.
+        """
+        custom_headers = self._build_headers(custom_headers, reply_to, return_path)
+
+        message, recipients = build_mime_message(
+            mail_from=mail_from,
+            to=to,
+            subject=subject,
+            html_content=html_content,
+            files=files,
+            cc=cc,
+            bcc=bcc,
+            mime_subtype=mime_subtype,
+            mime_charset=mime_charset,
+            custom_headers=custom_headers,
+        )
+
+        async with await self.get_async_conn() as async_client:
+            return await async_client.send_raw_email(
+                Source=mail_from, Destinations=recipients, RawMessage={"Data": message.as_string()}
+            )

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

@@ -22,6 +22,7 @@ from __future__ import annotations
 import json
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.utils.helpers import prune_dict
 
 
 def _get_message_attribute(o):
@@ -38,6 +39,33 @@ def _get_message_attribute(o):
     )
 
 
+def _build_publish_kwargs(
+    target_arn: str,
+    message: str,
+    subject: str | None = None,
+    message_attributes: dict | None = None,
+    message_deduplication_id: str | None = None,
+    message_group_id: str | None = None,
+) -> dict[str, str | dict]:
+    publish_kwargs: dict[str, str | dict] = prune_dict(
+        {
+            "TargetArn": target_arn,
+            "MessageStructure": "json",
+            "Message": json.dumps({"default": message}),
+            "Subject": subject,
+            "MessageDeduplicationId": message_deduplication_id,
+            "MessageGroupId": message_group_id,
+        }
+    )
+
+    if message_attributes:
+        publish_kwargs["MessageAttributes"] = {
+            key: _get_message_attribute(val) for key, val in message_attributes.items()
+        }
+
+    return publish_kwargs
+
+
 class SnsHook(AwsBaseHook):
     """
     Interact with Amazon Simple Notification Service.
@@ -84,22 +112,50 @@ class SnsHook(AwsBaseHook):
         :param message_group_id: Tag that specifies that a message belongs to a specific message group.
             This parameter applies only to FIFO (first-in-first-out) topics.
         """
-        publish_kwargs: dict[str, str | dict] = {
-            "TargetArn": target_arn,
-            "MessageStructure": "json",
-            "Message": json.dumps({"default": message}),
-        }
+        return self.get_conn().publish(
+            **_build_publish_kwargs(
+                target_arn, message, subject, message_attributes, message_deduplication_id, message_group_id
+            )
+        )
 
-        # Construct args this way because boto3 distinguishes from missing args and those set to None
-        if subject:
-            publish_kwargs["Subject"] = subject
-        if message_deduplication_id:
-            publish_kwargs["MessageDeduplicationId"] = message_deduplication_id
-        if message_group_id:
-            publish_kwargs["MessageGroupId"] = message_group_id
-        if message_attributes:
-            publish_kwargs["MessageAttributes"] = {
-                key: _get_message_attribute(val) for key, val in message_attributes.items()
-            }
-
-        return self.get_conn().publish(**publish_kwargs)
+    async def apublish_to_target(
+        self,
+        target_arn: str,
+        message: str,
+        subject: str | None = None,
+        message_attributes: dict | None = None,
+        message_deduplication_id: str | None = None,
+        message_group_id: str | None = None,
+    ):
+        """
+        Publish a message to a SNS topic or an endpoint.
+
+        .. seealso::
+            - :external+boto3:py:meth:`SNS.Client.publish`
+
+        :param target_arn: either a TopicArn or an EndpointArn
+        :param message: the default message you want to send
+        :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:
+
+            - bytes = Binary
+            - str = String
+            - int, float = Number
+            - iterable = String.Array
+        :param message_deduplication_id: Every message must have a unique message_deduplication_id.
+            This parameter applies only to FIFO (first-in-first-out) topics.
+        :param message_group_id: Tag that specifies that a message belongs to a specific message group.
+            This parameter applies only to FIFO (first-in-first-out) topics.
+        """
+        async with await self.get_async_conn() as async_client:
+            return await async_client.publish(
+                **_build_publish_kwargs(
+                    target_arn,
+                    message,
+                    subject,
+                    message_attributes,
+                    message_deduplication_id,
+                    message_group_id,
+                )
+            )

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

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.utils.helpers import prune_dict
 
 
 class SqsHook(AwsBaseHook):
@@ -52,6 +53,26 @@ class SqsHook(AwsBaseHook):
         """
         return self.get_conn().create_queue(QueueName=queue_name, Attributes=attributes or {})
 
+    @staticmethod
+    def _build_msg_params(
+        queue_url: str,
+        message_body: str,
+        delay_seconds: int = 0,
+        message_attributes: dict | None = None,
+        message_group_id: str | None = None,
+        message_deduplication_id: str | None = None,
+    ) -> dict:
+        return prune_dict(
+            {
+                "QueueUrl": queue_url,
+                "MessageBody": message_body,
+                "DelaySeconds": delay_seconds,
+                "MessageAttributes": message_attributes or {},
+                "MessageGroupId": message_group_id,
+                "MessageDeduplicationId": message_deduplication_id,
+            }
+        )
+
     def send_message(
         self,
         queue_url: str,
@@ -75,15 +96,47 @@ class SqsHook(AwsBaseHook):
         :param message_deduplication_id: This applies only to FIFO (first-in-first-out) queues.
         :return: dict with the information about the message sent
         """
-        params = {
-            "QueueUrl": queue_url,
-            "MessageBody": message_body,
-            "DelaySeconds": delay_seconds,
-            "MessageAttributes": message_attributes or {},
-        }
-        if message_group_id:
-            params["MessageGroupId"] = message_group_id
-        if message_deduplication_id:
-            params["MessageDeduplicationId"] = message_deduplication_id
-
+        params = self._build_msg_params(
+            queue_url=queue_url,
+            message_body=message_body,
+            delay_seconds=delay_seconds,
+            message_attributes=message_attributes,
+            message_group_id=message_group_id,
+            message_deduplication_id=message_deduplication_id,
+        )
         return self.get_conn().send_message(**params)
+
+    async def asend_message(
+        self,
+        queue_url: str,
+        message_body: str,
+        delay_seconds: int = 0,
+        message_attributes: dict | None = None,
+        message_group_id: str | None = None,
+        message_deduplication_id: str | None = None,
+    ) -> dict:
+        """
+        Send message to the queue (async).
+
+        .. seealso::
+            - :external+boto3:py:meth:`SQS.Client.send_message`
+
+        :param queue_url: queue url
+        :param message_body: the contents of the message
+        :param delay_seconds: seconds to delay the message
+        :param message_attributes: additional attributes for the message (default: None)
+        :param message_group_id: This applies only to FIFO (first-in-first-out) queues. (default: None)
+        :param message_deduplication_id: This applies only to FIFO (first-in-first-out) queues.
+        :return: dict with the information about the message sent
+        """
+        params = self._build_msg_params(
+            queue_url=queue_url,
+            message_body=message_body,
+            delay_seconds=delay_seconds,
+            message_attributes=message_attributes,
+            message_group_id=message_group_id,
+            message_deduplication_id=message_deduplication_id,
+        )
+
+        async with await self.get_async_conn() as async_conn:
+            return await async_conn.send_message(**params)

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

@@ -0,0 +1,139 @@
+# 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.
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from functools import cached_property
+from typing import Any
+
+from airflow.providers.amazon.aws.hooks.ses import SesHook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_1_PLUS
+from airflow.providers.common.compat.notifier import BaseNotifier
+from airflow.utils.helpers import prune_dict
+
+
+class SesNotifier(BaseNotifier):
+    """
+    Amazon Simple Email Service (SES) Notifier.
+
+    :param mail_from: Email address to set as email's from
+    :param to: List of email addresses to set as email's to
+    :param subject: Email's subject
+    :param html_content: Content of email in HTML format
+    :param files: List of paths of files to be attached
+    :param cc: List of email addresses to set as email's CC
+    :param bcc: List of email addresses to set as email's BCC
+    :param mime_subtype: Can be used to specify the subtype of the message. Default = mixed
+    :param mime_charset: Email's charset. Default = UTF-8.
+    :param return_path: The email address to which replies will be sent. By default, replies
+        are sent to the original sender's email address.
+    :param reply_to: The email address to which message bounces and complaints should be sent.
+        "Return-Path" is sometimes called "envelope from", "envelope sender", or "MAIL FROM".
+    :param custom_headers: Additional headers to add to the MIME message.
+        No validations are run on these values, and they should be able to be encoded.
+    """
+
+    template_fields: Sequence[str] = (
+        "aws_conn_id",
+        "region_name",
+        "mail_from",
+        "to",
+        "subject",
+        "html_content",
+        "files",
+        "cc",
+        "bcc",
+        "mime_subtype",
+        "mime_charset",
+        "reply_to",
+        "return_path",
+        "custom_headers",
+    )
+
+    def __init__(
+        self,
+        *,
+        aws_conn_id: str | None = SesHook.default_conn_name,
+        region_name: str | None = None,
+        mail_from: str,
+        to: str | Iterable[str],
+        subject: str,
+        html_content: str,
+        files: list[str] | None = None,
+        cc: str | Iterable[str] | None = None,
+        bcc: str | Iterable[str] | None = None,
+        mime_subtype: str = "mixed",
+        mime_charset: str = "utf-8",
+        reply_to: str | None = None,
+        return_path: str | None = None,
+        custom_headers: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        if AIRFLOW_V_3_1_PLUS:
+            #  Support for passing context was added in 3.1.0
+            super().__init__(**kwargs)
+        else:
+            super().__init__()
+        self.aws_conn_id = aws_conn_id
+        self.region_name = region_name
+
+        self.mail_from = mail_from
+        self.to = to
+        self.subject = subject
+        self.html_content = html_content
+        self.files = files
+        self.cc = cc
+        self.bcc = bcc
+        self.mime_subtype = mime_subtype
+        self.mime_charset = mime_charset
+        self.reply_to = reply_to
+        self.return_path = return_path
+        self.custom_headers = custom_headers
+
+    def _build_send_kwargs(self):
+        return prune_dict(
+            {
+                "mail_from": self.mail_from,
+                "to": self.to,
+                "subject": self.subject,
+                "html_content": self.html_content,
+                "files": self.files,
+                "cc": self.cc,
+                "bcc": self.bcc,
+                "mime_subtype": self.mime_subtype,
+                "mime_charset": self.mime_charset,
+                "reply_to": self.reply_to,
+                "return_path": self.return_path,
+                "custom_headers": self.custom_headers,
+            }
+        )
+
+    @cached_property
+    def hook(self) -> SesHook:
+        """Amazon Simple Email Service (SES) Hook (cached)."""
+        return SesHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+    def notify(self, context):
+        """Send email using Amazon Simple Email Service (SES)."""
+        self.hook.send_email(**self._build_send_kwargs())
+
+    async def async_notify(self, context):
+        """Send email using Amazon Simple Email Service (SES) (async)."""
+        await self.hook.asend_email(**self._build_send_kwargs())
+
+
+send_ses_notification = SesNotifier