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

potiuk 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 2d01cfd4a8 feature(providers): added `OpsgenieNotifier` (#35530)
2d01cfd4a8 is described below

commit 2d01cfd4a8b3875f3b511730a62fc4b897c40ac4
Author: Eric Mendes <42689328+eric-men...@users.noreply.github.com>
AuthorDate: Mon Nov 20 17:49:27 2023 -0300

    feature(providers): added `OpsgenieNotifier` (#35530)
    
    
    
    ---------
    
    Co-authored-by: Utkarsh Sharma <utkarshar...@gmail.com>
    Co-authored-by: Andrey Anshin <andrey.ans...@taragol.is>
---
 .../providers/opsgenie/notifications/__init__.py   |  16 +++
 .../providers/opsgenie/notifications/opsgenie.py   |  81 ++++++++++++++
 airflow/providers/opsgenie/provider.yaml           |   3 +
 airflow/providers/opsgenie/typing/__init__.py      |  16 +++
 airflow/providers/opsgenie/typing/opsgenie.py      |  59 ++++++++++
 docs/apache-airflow-providers-opsgenie/index.rst   |   1 +
 .../notifications/index.rst                        |  30 ++++++
 .../notifications/opsgenie_notifier.rst            |  33 ++++++
 tests/providers/opsgenie/notifications/__init__.py |  16 +++
 .../opsgenie/notifications/test_opsgenie.py        | 119 +++++++++++++++++++++
 tests/providers/opsgenie/typing/__init__.py        |  16 +++
 tests/providers/opsgenie/typing/test_opsgenie.py   |  29 +++++
 .../opsgenie/example_opsgenie_notifier.py          |  42 ++++++++
 13 files changed, 461 insertions(+)

diff --git a/airflow/providers/opsgenie/notifications/__init__.py 
b/airflow/providers/opsgenie/notifications/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++ b/airflow/providers/opsgenie/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/opsgenie/notifications/opsgenie.py 
b/airflow/providers/opsgenie/notifications/opsgenie.py
new file mode 100644
index 0000000000..950d92939e
--- /dev/null
+++ b/airflow/providers/opsgenie/notifications/opsgenie.py
@@ -0,0 +1,81 @@
+#
+# 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 TYPE_CHECKING, Sequence
+
+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.opsgenie.hooks.opsgenie import OpsgenieAlertHook
+
+if TYPE_CHECKING:
+    from airflow.providers.opsgenie.typing.opsgenie import CreateAlertPayload
+    from airflow.utils.context import Context
+
+
+class OpsgenieNotifier(BaseNotifier):
+    """
+    This notifier allows you to post alerts to Opsgenie.
+
+    Accepts a connection that has an Opsgenie API key as the connection's 
password.
+    This notifier sets the domain to conn_id.host, and if not set will default
+    to ``https://api.opsgenie.com``.
+
+    Each Opsgenie API key can be pre-configured to a team integration.
+    You can override these defaults in this notifier.
+
+    .. seealso::
+        For more information on how to use this notifier, take a look at the 
guide:
+        :ref:`howto/notifier:OpsgenieNotifier`
+
+    :param payload: The payload necessary for creating an alert.
+    :param opsgenie_conn_id: Optional. The name of the Opsgenie connection to 
use. Default conn_id is opsgenie_default
+    """
+
+    template_fields: Sequence[str] = ("payload",)
+
+    def __init__(
+        self,
+        *,
+        payload: CreateAlertPayload,
+        opsgenie_conn_id: str = "opsgenie_default",
+    ) -> None:
+        super().__init__()
+
+        self.payload = payload
+        self.opsgenie_conn_id = opsgenie_conn_id
+
+    @cached_property
+    def hook(self) -> OpsgenieAlertHook:
+        """Opsgenie alert Hook."""
+        return OpsgenieAlertHook(self.opsgenie_conn_id)
+
+    def notify(self, context: Context) -> None:
+        """Call the OpsgenieAlertHook to post message."""
+        self.hook.get_conn().create_alert(self.payload)
+
+
+send_opsgenie_notification = OpsgenieNotifier
diff --git a/airflow/providers/opsgenie/provider.yaml 
b/airflow/providers/opsgenie/provider.yaml
index d60f55f6d7..c60ca588f9 100644
--- a/airflow/providers/opsgenie/provider.yaml
+++ b/airflow/providers/opsgenie/provider.yaml
@@ -65,3 +65,6 @@ hooks:
 connection-types:
   - hook-class-name: 
airflow.providers.opsgenie.hooks.opsgenie.OpsgenieAlertHook
     connection-type: opsgenie
+
+notifications:
+  - airflow.providers.opsgenie.notifications.opsgenie.OpsgenieNotifier
diff --git a/airflow/providers/opsgenie/typing/__init__.py 
b/airflow/providers/opsgenie/typing/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++ b/airflow/providers/opsgenie/typing/__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/opsgenie/typing/opsgenie.py 
b/airflow/providers/opsgenie/typing/opsgenie.py
new file mode 100644
index 0000000000..4e6621785d
--- /dev/null
+++ b/airflow/providers/opsgenie/typing/opsgenie.py
@@ -0,0 +1,59 @@
+# 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 typing import TypedDict
+
+from typing_extensions import NotRequired, Required  # For compat with Python 
< 3.11
+
+
+class CreateAlertPayload(TypedDict):
+    """
+    Payload schema for creating an Opsgenie alert.
+
+    :param message: The Message of the Opsgenie alert.
+    :param alias: Client-defined identifier of the alert.
+    :param description: Description field of the alert.
+    :param responders: Teams, users, escalations and schedules that
+        the alert will be routed to send notifications.
+    :param visible_to: Teams and users that the alert will become visible
+        to without sending any notification.
+    :param actions: Custom actions that will be available for the alert.
+    :param tags: Tags of the alert.
+    :param details: Map of key-value pairs to use as custom properties of the 
alert.
+    :param entity: Entity field of the alert that is
+        generally used to specify which domain alert is related to.
+    :param source: Source field of the alert. Default value is
+        IP address of the incoming request.
+    :param priority: Priority level of the alert. Default value is P3.
+    :param user: Display name of the request owner.
+    :param note: Additional note that will be added while creating the alert.
+    """
+
+    message: Required[str]
+    alias: NotRequired[str | None]
+    description: NotRequired[str | None]
+    responders: NotRequired[list[dict] | None]
+    visible_to: NotRequired[list[dict] | None]
+    actions: NotRequired[list[str] | None]
+    tags: NotRequired[list[str] | None]
+    details: NotRequired[dict | None]
+    entity: NotRequired[str | None]
+    source: NotRequired[str | None]
+    priority: NotRequired[str | None]
+    user: NotRequired[str | None]
+    note: NotRequired[str | None]
diff --git a/docs/apache-airflow-providers-opsgenie/index.rst 
b/docs/apache-airflow-providers-opsgenie/index.rst
index f1409d7eb5..ec32005261 100644
--- a/docs/apache-airflow-providers-opsgenie/index.rst
+++ b/docs/apache-airflow-providers-opsgenie/index.rst
@@ -35,6 +35,7 @@
     :caption: Guides
 
     Operators <operators/index>
+    Notifications <notifications/index>
 
 .. toctree::
     :hidden:
diff --git a/docs/apache-airflow-providers-opsgenie/notifications/index.rst 
b/docs/apache-airflow-providers-opsgenie/notifications/index.rst
new file mode 100644
index 0000000000..15f9b53bcb
--- /dev/null
+++ b/docs/apache-airflow-providers-opsgenie/notifications/index.rst
@@ -0,0 +1,30 @@
+ .. 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.
+
+
+
+Opsgenie Notifiers
+====================
+
+.. important:: This feature is only available in Airflow versions >= 2.6.0
+
+
+.. toctree::
+    :maxdepth: 1
+    :glob:
+
+    *
diff --git 
a/docs/apache-airflow-providers-opsgenie/notifications/opsgenie_notifier.rst 
b/docs/apache-airflow-providers-opsgenie/notifications/opsgenie_notifier.rst
new file mode 100644
index 0000000000..54cff970d4
--- /dev/null
+++ b/docs/apache-airflow-providers-opsgenie/notifications/opsgenie_notifier.rst
@@ -0,0 +1,33 @@
+ .. 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:OpsgenieNotifier:
+
+Opsgenie Alert Notifier
+=======================
+
+Use the 
:class:`~airflow.providers.opsgenie.notifications.opsgenie.OpsgenieNotifier` to 
send an alert to opsgenie.
+
+
+Using the Notifier
+^^^^^^^^^^^^^^^^^^
+Send an alert to Opsgenie with a specific message.
+
+.. exampleinclude:: 
/../../tests/system/providers/opsgenie/example_opsgenie_notifier.py
+    :language: python
+    :start-after: [START howto_notifier_opsgenie]
+    :end-before: [END howto_notifier_opsgenie]
diff --git a/tests/providers/opsgenie/notifications/__init__.py 
b/tests/providers/opsgenie/notifications/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++ b/tests/providers/opsgenie/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/opsgenie/notifications/test_opsgenie.py 
b/tests/providers/opsgenie/notifications/test_opsgenie.py
new file mode 100644
index 0000000000..d5915c3557
--- /dev/null
+++ b/tests/providers/opsgenie/notifications/test_opsgenie.py
@@ -0,0 +1,119 @@
+#
+# 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.operators.empty import EmptyOperator
+from airflow.providers.opsgenie.hooks.opsgenie import OpsgenieAlertHook
+from airflow.providers.opsgenie.notifications.opsgenie import 
OpsgenieNotifier, send_opsgenie_notification
+
+pytestmark = pytest.mark.db_test
+
+
+class TestOpsgenieNotifier:
+    _config = {
+        "message": "An example alert message",
+        "alias": "Life is too short for no alias",
+        "description": "Every alert needs a description",
+        "responders": [
+            {"id": "4513b7ea-3b91-438f-b7e4-e3e54af9147c", "type": "team"},
+            {"name": "NOC", "type": "team"},
+            {"id": "bb4d9938-c3c2-455d-aaab-727aa701c0d8", "type": "user"},
+            {"username": "trin...@opsgenie.com", "type": "user"},
+            {"id": "aee8a0de-c80f-4515-a232-501c0bc9d715", "type": 
"escalation"},
+            {"name": "Nightwatch Escalation", "type": "escalation"},
+            {"id": "80564037-1984-4f38-b98e-8a1f662df552", "type": "schedule"},
+            {"name": "First Responders Schedule", "type": "schedule"},
+        ],
+        "visible_to": [
+            {"id": "4513b7ea-3b91-438f-b7e4-e3e54af9147c", "type": "team"},
+            {"name": "rocket_team", "type": "team"},
+            {"id": "bb4d9938-c3c2-455d-aaab-727aa701c0d8", "type": "user"},
+            {"username": "trin...@opsgenie.com", "type": "user"},
+        ],
+        "actions": ["Restart", "AnExampleAction"],
+        "tags": ["OverwriteQuietHours", "Critical"],
+        "details": {"key1": "value1", "key2": "value2"},
+        "entity": "An example entity",
+        "source": "Airflow",
+        "priority": "P1",
+        "user": "Jesse",
+        "note": "Write this down",
+    }
+
+    expected_payload_dict = {
+        "message": _config["message"],
+        "alias": _config["alias"],
+        "description": _config["description"],
+        "responders": _config["responders"],
+        "visible_to": _config["visible_to"],
+        "actions": _config["actions"],
+        "tags": _config["tags"],
+        "details": _config["details"],
+        "entity": _config["entity"],
+        "source": _config["source"],
+        "priority": _config["priority"],
+        "user": _config["user"],
+        "note": _config["note"],
+    }
+
+    @mock.patch.object(OpsgenieAlertHook, "get_conn")
+    def test_notifier(self, mock_opsgenie_alert_hook, dag_maker):
+        with dag_maker("test_notifier") as dag:
+            EmptyOperator(task_id="task1")
+        notifier = send_opsgenie_notification(payload=self._config)
+        notifier({"dag": dag})
+        
mock_opsgenie_alert_hook.return_value.create_alert.assert_called_once_with(self.expected_payload_dict)
+
+    @mock.patch.object(OpsgenieAlertHook, "get_conn")
+    def test_notifier_with_notifier_class(self, mock_opsgenie_alert_hook, 
dag_maker):
+        with dag_maker("test_notifier") as dag:
+            EmptyOperator(task_id="task1")
+        notifier = OpsgenieNotifier(payload=self._config)
+        notifier({"dag": dag})
+        
mock_opsgenie_alert_hook.return_value.create_alert.assert_called_once_with(self.expected_payload_dict)
+
+    @mock.patch.object(OpsgenieAlertHook, "get_conn")
+    def test_notifier_templated(self, mock_opsgenie_alert_hook, dag_maker):
+        dag_id = "test_notifier"
+        with dag_maker(dag_id) as dag:
+            EmptyOperator(task_id="task1")
+
+        template_fields = ("message", "alias", "description", "entity", 
"priority", "note")
+        templated_config = {}
+        for key, value in self._config.items():
+            if key in template_fields:
+                templated_config[key] = value + " {{dag.dag_id}}"
+            else:
+                templated_config[key] = value
+
+        templated_expected_payload_dict = {}
+        for key, value in self.expected_payload_dict.items():
+            if key in template_fields:
+                templated_expected_payload_dict[key] = value + f" {dag_id}"
+            else:
+                templated_expected_payload_dict[key] = value
+
+        notifier = OpsgenieNotifier(payload=templated_config)
+        notifier({"dag": dag})
+        
mock_opsgenie_alert_hook.return_value.create_alert.assert_called_once_with(
+            templated_expected_payload_dict
+        )
diff --git a/tests/providers/opsgenie/typing/__init__.py 
b/tests/providers/opsgenie/typing/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++ b/tests/providers/opsgenie/typing/__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/opsgenie/typing/test_opsgenie.py 
b/tests/providers/opsgenie/typing/test_opsgenie.py
new file mode 100644
index 0000000000..b782305187
--- /dev/null
+++ b/tests/providers/opsgenie/typing/test_opsgenie.py
@@ -0,0 +1,29 @@
+#
+# 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
+
+import pytest
+
+from airflow.providers.opsgenie.typing.opsgenie import CreateAlertPayload
+
+pytestmark = pytest.mark.db_test
+
+
+class TestCreateAlertPayload:
+    def test_create_alert_payload_required_keys(self):
+        CreateAlertPayload.__required_keys__ = frozenset({"message"})
diff --git a/tests/system/providers/opsgenie/example_opsgenie_notifier.py 
b/tests/system/providers/opsgenie/example_opsgenie_notifier.py
new file mode 100644
index 0000000000..8d0847817b
--- /dev/null
+++ b/tests/system/providers/opsgenie/example_opsgenie_notifier.py
@@ -0,0 +1,42 @@
+# 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.
+
+# [START howto_notifier_opsgenie]
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow import DAG
+from airflow.operators.bash import BashOperator
+from airflow.providers.opsgenie.notifications.opsgenie import 
send_opsgenie_notification
+
+with DAG(
+    "opsgenie_notifier",
+    start_date=datetime(2023, 1, 1),
+    on_failure_callback=[send_opsgenie_notification(payload={"message": 
"Something went wrong!"})],
+) as dag:
+    BashOperator(
+        task_id="mytask",
+        bash_command="fail",
+        on_failure_callback=[send_opsgenie_notification(payload={"message": 
"Something went wrong!"})],
+    )
+# [END howto_notifier_opsgenie]
+
+from tests.system.utils import get_test_run
+
+# Needed to run the example DAG with pytest (see: 
tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)

Reply via email to