diff --git a/airflow/providers/opsgenie/notifications/__init__.py b/airflow/providers/opsgenie/notifications/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /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 0000000000000..950d92939e3b7 --- /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 d60f55f6d7d98..c60ca588f9924 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 0000000000000..13a83393a9124 --- /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 0000000000000..4e6621785d63d --- /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 f1409d7eb5c80..ec32005261ce7 100644 --- a/docs/apache-airflow-providers-opsgenie/index.rst +++ b/docs/apache-airflow-providers-opsgenie/index.rst @@ -35,6 +35,7 @@ :caption: Guides Operators + Notifications .. 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 0000000000000..15f9b53bcb47a --- /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 0000000000000..54cff970d4a68 --- /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 0000000000000..13a83393a9124 --- /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 0000000000000..d5915c35579b5 --- /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": "trinity@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": "trinity@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 0000000000000..13a83393a9124 --- /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 0000000000000..b782305187fb9 --- /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 0000000000000..8d0847817bbd3 --- /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)