Add Amazon SNS Notifier

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

@@ -68,7 +68,6 @@ def publish_to_target(
 
         :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/notifications/sns.py

@@ -0,0 +1,94 @@
+# 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 functools import cached_property
+from typing import Sequence
+
+from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.amazon.aws.hooks.sns import SnsHook
+
+try:
+    from airflow.notifications.basenotifier import BaseNotifier
+except ImportError:
+    raise AirflowOptionalProviderFeatureException(
+        "Failed to import BaseNotifier. This feature is only available in Airflow versions >= 2.6.0"
+    )
+
+
+class SnsNotifier(BaseNotifier):
+    """
+    Amazon SNS (Simple Notification Service) Notifier.
+
+    .. seealso::
+        For more information on how to use this notifier, take a look at the guide:
+        :ref:`howto/notifier:SnsNotifier`
+
+    :param aws_conn_id: The :ref:`Amazon Web Services Connection id <howto/connection:aws>`
+        used for AWS credentials. If this is None or empty then the default boto3 behaviour is used.
+    :param target_arn: Either a TopicArn or an EndpointArn.
+    :param message: The message you want to send.
+    :param subject: The message subject you want to send.
+    :param message_attributes: The message attributes you want to send as a flat dict (data type will be
+        determined automatically).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    """
+
+    template_fields: Sequence[str] = (
+        "target_arn",
+        "message",
+        "subject",
+        "message_attributes",
+        "aws_conn_id",
+        "region_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        aws_conn_id: str | None = SnsHook.default_conn_name,
+        target_arn: str,
+        message: str,
+        subject: str | None = None,
+        message_attributes: dict | None = None,
+        region_name: str | None = None,
+    ):
+        super().__init__()
+        self.aws_conn_id = aws_conn_id
+        self.region_name = region_name
+        self.target_arn = target_arn
+        self.message = message
+        self.subject = subject
+        self.message_attributes = message_attributes
+
+    @cached_property
+    def hook(self) -> SnsHook:
+        """Amazon SNS Hook (cached)."""
+        return SnsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+    def notify(self, context):
+        """Publish the notification message to Amazon SNS."""
+        self.hook.publish_to_target(
+            target_arn=self.target_arn,
+            message=self.message,
+            subject=self.subject,
+            message_attributes=self.message_attributes,
+        )
+
+
+send_sns_notification = SnsNotifier

airflow/providers/amazon/provider.yaml

@@ -689,7 +689,7 @@ connection-types:
 
 notifications:
   - airflow.providers.amazon.aws.notifications.chime.ChimeNotifier
-
+  - airflow.providers.amazon.aws.notifications.sns.SnsNotifier
 
 secrets-backends:
   - airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend

docs/apache-airflow-providers-amazon/notifications/index.rst

@@ -20,6 +20,8 @@
 Amazon AWS Notifications
 ========================
 
+.. important:: This feature is only available in Airflow versions >= 2.6.0
+
 
 .. toctree::
     :maxdepth: 1

docs/apache-airflow-providers-amazon/notifications/sns.rst

@@ -0,0 +1,62 @@
+ .. 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.
+
+.. _howto/notifier:SnsNotifier:
+
+How-to Guide for SNS notifications
+==================================
+
+Introduction
+------------
+`AWS SNS <https://aws.amazon.com/sns/>`__ notifier :class:`~airflow.providers.amazon.aws.notifications.sns.SnsNotifier`
+allows users to push messages to a SNS Topic using the various ``on_*_callbacks`` at both the DAG level and Task level.
+
+You can also use a notifier with ``sla_miss_callback``.
+
+.. note::
+    When notifiers are used with ``sla_miss_callback`` the context will contain only values passed to the callback,
+    refer :ref:`sla_miss_callback<concepts:sla_miss_callback>`.
+
+Example Code:
+-------------
+
+.. code-block:: python
+
+    from datetime import datetime
+    from airflow import DAG
+    from airflow.operators.python import PythonOperator
+    from airflow.providers.amazon.aws.notifications.sns import send_sns_notification
+
+    dag_failure_sns_notification = send_sns_notification(
+        aws_conn_id="aws_default",
+        region_name="eu-west-2",
+        message="The DAG {{ dag.dag_id }} failed",
+        target_arn="arn:aws:sns:us-west-2:123456789098:TopicName",
+    )
+    task_failure_sns_notification = send_sns_notification(
+        aws_conn_id="aws_default",
+        region_name="eu-west-2",
+        message="The task {{ ti.task_id }} failed",
+        target_arn="arn:aws:sns:us-west-2:123456789098:AnotherTopicName",
+    )
+
+    with DAG(start_date=datetime(2023, 1, 1), on_failure_callback=[dag_failure_sns_notification]):
+        BashOperator(
+            task_id="mytask",
+            on_failure_callback=[task_failure_sns_notification],
+            bash_command="fail",
+        )

tests/providers/amazon/aws/notifications/test_sns.py

@@ -0,0 +1,53 @@
+# 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 unittest import mock
+
+import pytest
+
+from airflow.providers.amazon.aws.notifications.sns import SnsNotifier, send_sns_notification
+from airflow.utils.types import NOTSET
+
+PARAM_DEFAULT_VALUE = pytest.param(NOTSET, id="default-value")
+
+
+class TestSnsNotifier:
+    def test_class_and_notifier_are_same(self):
+        assert send_sns_notification is SnsNotifier
+
+    @pytest.mark.parametrize("aws_conn_id", ["aws_test_conn_id", None, PARAM_DEFAULT_VALUE])
+    @pytest.mark.parametrize("region_name", ["eu-west-2", None, PARAM_DEFAULT_VALUE])
+    def test_parameters_propagate_to_hook(self, aws_conn_id, region_name):
+        """Test notifier attributes propagate to SnsHook."""
+        notifier_kwargs = {
+            "target_arn": "arn:aws:sns:us-west-2:123456789098:TopicName",
+            "message": "foo-bar",
+        }
+        if aws_conn_id is not NOTSET:
+            notifier_kwargs["aws_conn_id"] = aws_conn_id
+        if region_name is not NOTSET:
+            notifier_kwargs["region_name"] = region_name
+
+        notifier = SnsNotifier(**notifier_kwargs)
+        with mock.patch("airflow.providers.amazon.aws.notifications.sns.SnsHook") as mock_hook:
+            hook = notifier.hook
+            assert hook is notifier.hook, "Hook property not cached"
+            mock_hook.assert_called_once_with(
+                aws_conn_id=(aws_conn_id if aws_conn_id is not NOTSET else "aws_default"),
+                region_name=(region_name if region_name is not NOTSET else None),
+            )