diff --git a/airflow/providers/discord/notifications/__init__.py b/airflow/providers/discord/notifications/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/airflow/providers/discord/notifications/__init__.py @@ -0,0 +1,17 @@ +# +# 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/discord/notifications/discord.py b/airflow/providers/discord/notifications/discord.py new file mode 100644 index 0000000000000..629b8c2b5962a --- /dev/null +++ b/airflow/providers/discord/notifications/discord.py @@ -0,0 +1,82 @@ +# +# 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 + +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" + ) + +from airflow.providers.discord.hooks.discord_webhook import DiscordWebhookHook + +ICON_URL: str = "https://raw.githubusercontent.com/apache/airflow/main/airflow/www/static/pin_100.png" + + +class DiscordNotifier(BaseNotifier): + """ + Discord BaseNotifier. + + :param discord_conn_id: Http connection ID with host as "https://discord.com/api/" and + default webhook endpoint in the extra field in the form of + {"webhook_endpoint": "webhooks/{webhook.id}/{webhook.token}"} + :param text: The content of the message + :param username: The username to send the message as. Optional + :param avatar_url: The URL of the avatar to use for the message. Optional + :param tts: Text to speech. + """ + + # A property that specifies the attributes that can be templated. + template_fields = ("discord_conn_id", "text", "username", "avatar_url", "tts") + + def __init__( + self, + discord_conn_id: str = "discord_webhook_default", + text: str = "This is a default message", + username: str = "Airflow", + avatar_url: str = ICON_URL, + tts: bool = False, + ): + super().__init__() + self.discord_conn_id = discord_conn_id + self.text = text + self.username = username + self.avatar_url = avatar_url + + # If you're having problems with tts not being recognized in __init__(), + # you can define that after instantiating the class + self.tts = tts + + @cached_property + def hook(self) -> DiscordWebhookHook: + """Discord Webhook Hook.""" + return DiscordWebhookHook(http_conn_id=self.discord_conn_id) + + def notify(self, context): + """Send a message to a Discord channel.""" + self.hook.username = self.username + self.hook.message = self.text + self.hook.avatar_url = self.avatar_url + self.hook.tts = self.tts + + self.hook.execute() diff --git a/airflow/providers/discord/provider.yaml b/airflow/providers/discord/provider.yaml index ee751779bc459..6dedfc1fe0e7e 100644 --- a/airflow/providers/discord/provider.yaml +++ b/airflow/providers/discord/provider.yaml @@ -58,3 +58,6 @@ hooks: connection-types: - hook-class-name: airflow.providers.discord.hooks.discord_webhook.DiscordWebhookHook connection-type: discord + +notifications: + - airflow.providers.discord.notifications.discord.DiscordNotifier diff --git a/tests/providers/discord/notifications/__init__.py b/tests/providers/discord/notifications/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/providers/discord/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/discord/notifications/test_discord.py b/tests/providers/discord/notifications/test_discord.py new file mode 100644 index 0000000000000..279dafa515523 --- /dev/null +++ b/tests/providers/discord/notifications/test_discord.py @@ -0,0 +1,58 @@ +# +# 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.mock import MagicMock, patch + +import pytest + +from airflow.models import Connection +from airflow.providers.discord.notifications.discord import DiscordNotifier +from airflow.utils import db + + +@pytest.fixture(autouse=True) +def setup(): + db.merge_conn( + Connection( + conn_id="my_discord_conn_id", + conn_type="discord", + host="https://discordapp.com/api/", + extra='{"webhook_endpoint": "webhooks/00000/some-discord-token_000"}', + ) + ) + + +@patch("airflow.providers.discord.notifications.discord.DiscordWebhookHook.execute") +def test_discord_notifier_notify(mock_execute): + notifier = DiscordNotifier( + discord_conn_id="my_discord_conn_id", + text="This is a test message", + username="test_user", + avatar_url="https://example.com/avatar.png", + tts=False, + ) + context = MagicMock() + + notifier.notify(context) + + mock_execute.assert_called_once() + assert notifier.hook.username == "test_user" + assert notifier.hook.message == "This is a test message" + assert notifier.hook.avatar_url == "https://example.com/avatar.png" + assert notifier.hook.tts is False