diff --git a/airflow/providers/amazon/aws/notifications/__init__.py b/airflow/providers/amazon/aws/notifications/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow/providers/amazon/aws/notifications/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow/providers/amazon/aws/notifications/chime.py b/airflow/providers/amazon/aws/notifications/chime.py new file mode 100644 index 0000000000000..c505b3e227eac --- /dev/null +++ b/airflow/providers/amazon/aws/notifications/chime.py @@ -0,0 +1,61 @@ +# 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 airflow.exceptions import AirflowOptionalProviderFeatureException +from airflow.providers.amazon.aws.hooks.chime import ChimeWebhookHook +from airflow.utils.context import Context + +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 ChimeNotifier(BaseNotifier): + """ + Chime notifier to send messages to a chime room via callbacks. + + :param: chime_conn_id: The chime connection to use with Endpoint as "https://hooks.chime.aws" and + the webhook token in the form of ```{webhook.id}?token{webhook.token}``` + :param: message: The message to send to the chime room associated with the webhook. + + """ + + template_fields = ("message",) + + def __init__(self, *, chime_conn_id: str, message: str = "This is the default chime notifier message"): + super().__init__() + self.chime_conn_id = chime_conn_id + self.message = message + + @cached_property + def hook(self): + """To reduce overhead cache the hook for the notifier.""" + return ChimeWebhookHook(chime_conn_id=self.chime_conn_id) + + def notify(self, context: Context) -> None: + """Send a message to a Chime Chat Room.""" + self.hook.send_message(message=self.message) + + +send_chime_notification = ChimeNotifier diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml index e4f16ce398460..bf2391d9a8388 100644 --- a/airflow/providers/amazon/provider.yaml +++ b/airflow/providers/amazon/provider.yaml @@ -642,6 +642,9 @@ connection-types: - hook-class-name: airflow.providers.amazon.aws.hooks.redshift_sql.RedshiftSQLHook connection-type: redshift +notifications: + - airflow.providers.amazon.aws.notifications.chime.ChimeNotifier + secrets-backends: - airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend diff --git a/docs/apache-airflow-providers-amazon/index.rst b/docs/apache-airflow-providers-amazon/index.rst index e4fd28da4a369..85d5542f8c00c 100644 --- a/docs/apache-airflow-providers-amazon/index.rst +++ b/docs/apache-airflow-providers-amazon/index.rst @@ -34,6 +34,7 @@ :caption: Guides Connection types + Notifications Operators Transfers Deferrable Operators diff --git a/docs/apache-airflow-providers-amazon/notifications/chime_notifier_howto_guide.rst b/docs/apache-airflow-providers-amazon/notifications/chime_notifier_howto_guide.rst new file mode 100644 index 0000000000000..bf346aad5b0b8 --- /dev/null +++ b/docs/apache-airflow-providers-amazon/notifications/chime_notifier_howto_guide.rst @@ -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. + +How-to Guide for Chime notifications +==================================== + +Introduction +------------ +Chime notifier (:class:`airflow.providers.amazon.aws.notifications.chime.ChimeNotifier`) allows users to send +messages to a Chime chat room setup via a webhook 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`. + +Example Code: +------------- + +.. code-block:: python + + from datetime import datetime + from airflow import DAG + from airflow.operators.bash import BashOperator + from airflow.providers.amazon.aws.notifications.chime import send_chime_notification + + with DAG( + start_date=datetime(2023, 6, 27), + on_success_callback=[ + send_chime_notification(chime_conn_id="my_chime_conn", message="The DAG {{ dag.dag_id }} succeeded") + ], + ): + BashOperator( + task_id="mytask", + on_failure_callback=[ + send_chime_notification(chime_conn_id="my_chime_conn", message="The task {{ ti.task_id }} failed") + ], + bash_command="fail", + ) diff --git a/docs/apache-airflow-providers-amazon/notifications/index.rst b/docs/apache-airflow-providers-amazon/notifications/index.rst new file mode 100644 index 0000000000000..6cd2c853c22ec --- /dev/null +++ b/docs/apache-airflow-providers-amazon/notifications/index.rst @@ -0,0 +1,28 @@ + .. 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. + + + +Amazon AWS Notifications +======================== + + +.. toctree:: + :maxdepth: 1 + :glob: + + * diff --git a/tests/providers/amazon/aws/notifications/__init__.py b/tests/providers/amazon/aws/notifications/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/providers/amazon/aws/notifications/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/providers/amazon/aws/notifications/test_chime.py b/tests/providers/amazon/aws/notifications/test_chime.py new file mode 100644 index 0000000000000..5dc774d0d5e8a --- /dev/null +++ b/tests/providers/amazon/aws/notifications/test_chime.py @@ -0,0 +1,73 @@ +# 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 + +from airflow.models import Connection +from airflow.operators.empty import EmptyOperator +from airflow.providers.amazon.aws.hooks.chime import ChimeWebhookHook +from airflow.providers.amazon.aws.notifications.chime import ChimeNotifier, send_chime_notification +from airflow.utils import db + + +class TestChimeNotifier: + # Chime webhooks can't really have a default connection, so we need to create one for tests. + def setup_method(self): + db.merge_conn( + Connection( + conn_id="default-chime-webhook", + conn_type="chime", + host="hooks.chime.aws/incomingwebhooks/", + password="abcd-1134-ZeDA?token=somechimetoken111", + schema="https", + ) + ) + + @mock.patch.object(ChimeWebhookHook, "send_message") + def test_chime_notifier(self, mock_chime_hook, dag_maker): + with dag_maker("test_chime_notifier") as dag: + EmptyOperator(task_id="task1") + + notifier = send_chime_notification( + chime_conn_id="default-chime-webhook", message="Chime Test Message" + ) + notifier({"dag": dag}) + mock_chime_hook.assert_called_once_with(message="Chime Test Message") + + @mock.patch.object(ChimeWebhookHook, "send_message") + def test_chime_notifier_with_notifier_class(self, mock_chime_hook, dag_maker): + with dag_maker("test_chime_notifier") as dag: + EmptyOperator(task_id="task1") + + notifier = ChimeNotifier( + chime_conn_id="default-chime-webhook", message="Test Chime Message for Class" + ) + notifier({"dag": dag}) + mock_chime_hook.assert_called_once_with(message="Test Chime Message for Class") + + @mock.patch.object(ChimeWebhookHook, "send_message") + def test_chime_notifier_templated(self, mock_chime_hook, dag_maker): + with dag_maker("test_chime_notifier") as dag: + EmptyOperator(task_id="task1") + + notifier = send_chime_notification( + chime_conn_id="default-chime-webhook", message="Test Chime Message. Dag is {{ dag.dag_id }}." + ) + notifier({"dag": dag}) + mock_chime_hook.assert_called_once_with(message="Test Chime Message. Dag is test_chime_notifier.")