[airflow] branch main updated: Add `DiscordNotifier` (#31273)

[eladkal]

This is an automated email from the ASF dual-hosted git repository.

eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new bdfebad5c9 Add `DiscordNotifier` (#31273)
bdfebad5c9 is described below

commit bdfebad5c9491234a78453856bd8c3baac98f75e
Author: Guilherme Cardoso de Vargas 
<77084039+vargacyp...@users.noreply.github.com>
AuthorDate: Fri Jun 16 02:49:03 2023 -0300

    Add `DiscordNotifier` (#31273)
    
    * Add DiscordNotifier
    
    ---------
    
    Co-authored-by: guilherme.vargas <guilherme.var...@otocrm.com.br>
    Co-authored-by: Pankaj Koti <pankajkoti...@gmail.com>
    Co-authored-by: Shivam Shrey <31056227+shivamsh...@users.noreply.github.com>
    Co-authored-by: eladkal <45845474+elad...@users.noreply.github.com>
---
 .../providers/discord/notifications/__init__.py    | 17 +++++
 airflow/providers/discord/notifications/discord.py | 82 ++++++++++++++++++++++
 airflow/providers/discord/provider.yaml            |  3 +
 tests/providers/discord/notifications/__init__.py  | 16 +++++
 .../discord/notifications/test_discord.py          | 58 +++++++++++++++
 5 files changed, 176 insertions(+)

diff --git a/airflow/providers/discord/notifications/__init__.py 
b/airflow/providers/discord/notifications/__init__.py
new file mode 100644
index 0000000000..217e5db960
--- /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 0000000000..629b8c2b59
--- /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 ee751779bc..6dedfc1fe0 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 0000000000..13a83393a9
--- /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 0000000000..279dafa515
--- /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