From dc7d0fe441505646d90745108da74173ad967ec3 Mon Sep 17 00:00:00 2001 From: eladkal <45845474+eladkal@users.noreply.github.com> Date: Mon, 2 May 2022 16:34:50 +0300 Subject: [PATCH 1/4] Add `OpsgenieDeleteAlertOperator` --- .../example_dags/example_opsgenie_alert.py | 7 +++ airflow/providers/opsgenie/hooks/opsgenie.py | 30 +++++++++++ .../providers/opsgenie/operators/opsgenie.py | 52 +++++++++++++++++++ .../operators/opsgenie_alert.rst | 17 ++++++ .../providers/opsgenie/hooks/test_opsgenie.py | 16 ++++++ .../opsgenie/operators/test_opsgenie.py | 30 +++++++++++ 6 files changed, 152 insertions(+) diff --git a/airflow/providers/opsgenie/example_dags/example_opsgenie_alert.py b/airflow/providers/opsgenie/example_dags/example_opsgenie_alert.py index 96e6293d3b106..ad294e5e6e1e0 100644 --- a/airflow/providers/opsgenie/example_dags/example_opsgenie_alert.py +++ b/airflow/providers/opsgenie/example_dags/example_opsgenie_alert.py @@ -20,6 +20,7 @@ from airflow.providers.opsgenie.operators.opsgenie import ( OpsgenieCloseAlertOperator, OpsgenieCreateAlertOperator, + OpsgenieDeleteAlertOperator, ) with DAG( @@ -38,3 +39,9 @@ task_id="opsgenie_close_task", identifier="identifier_example" ) # [END howto_opsgenie_close_alert_operator] + + # [START howto_opsgenie_delete_alert_operator] + opsgenie_delete_alert_operator = OpsgenieDeleteAlertOperator( + task_id="opsgenie_delete_task", identifier="identifier_example" + ) + # [END howto_opsgenie_delete_alert_operator] diff --git a/airflow/providers/opsgenie/hooks/opsgenie.py b/airflow/providers/opsgenie/hooks/opsgenie.py index 28e7c7452ab15..78d023ef5d273 100644 --- a/airflow/providers/opsgenie/hooks/opsgenie.py +++ b/airflow/providers/opsgenie/hooks/opsgenie.py @@ -132,3 +132,33 @@ def close_alert( except OpenApiException as e: self.log.exception('Exception when closing alert in opsgenie with payload: %s', payload) raise e + + def delete_alert( + self, + identifier: str, + identifier_type: Optional[str], + user: Optional[str], + source: Optional[str], + ) -> SuccessResponse: + """ + Delete an alert in Opsgenie + + :param identifier: Identifier of alert which could be alert id, tiny id or alert alias. + :param identifier_type: Type of the identifier that is provided as an in-line parameter. + Possible values are 'id', 'alias' or 'tiny' + :param user: Display name of the request owner. + :param source: Display name of the request source + :return: SuccessResponse + :rtype: opsgenie_sdk.SuccessResponse + """ + try: + api_response = self.alert_api_instance.delete_alert( + identifier=identifier, + identifier_type=identifier_type, + user=user, + source=source, + ) + return api_response + except OpenApiException as e: + self.log.exception('Exception when calling AlertApi->delete_alert: %s\n', e) + raise e diff --git a/airflow/providers/opsgenie/operators/opsgenie.py b/airflow/providers/opsgenie/operators/opsgenie.py index d072d82f41cca..4dee8d4292e59 100644 --- a/airflow/providers/opsgenie/operators/opsgenie.py +++ b/airflow/providers/opsgenie/operators/opsgenie.py @@ -208,3 +208,55 @@ def execute(self, context: 'Context') -> None: payload=self._build_opsgenie_close_alert_payload(), kwargs=self.close_alert_kwargs, ) + + +class OpsgenieDeleteAlertOperator(BaseOperator): + """ + This operator allows you to Delete alerts to Opsgenie. + Accepts a connection that has an Opsgenie API key as the connection's password. + This operator 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 operator. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:OpsgenieDeleteAlertOperator` + + :param opsgenie_conn_id: The name of the Opsgenie connection to use + :param identifier: Identifier of alert which could be alert id, tiny id or alert alias + :param identifier_type: Type of the identifier that is provided as an in-line parameter. + Possible values are 'id', 'alias' or 'tiny' + :param user: Display name of the request owner + :param source: Display name of the request source + """ + + def __init__( + self, + *, + identifier: str, + opsgenie_conn_id: str = 'opsgenie_default', + identifier_type: Optional[str] = None, + user: Optional[str] = None, + source: Optional[str] = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + + self.opsgenie_conn_id = opsgenie_conn_id + self.identifier = identifier + self.identifier_type = identifier_type + self.user = user + self.source = source + self.hook: Optional[OpsgenieAlertHook] = None + + def execute(self, context: 'Context') -> None: + """Call the OpsgenieAlertHook to delete alert""" + self.hook = OpsgenieAlertHook(self.opsgenie_conn_id) + self.hook.delete_alert( + identifier=self.identifier, + identifier_type=self.identifier_type, + user=self.user, + source=self.source, + ) diff --git a/docs/apache-airflow-providers-opsgenie/operators/opsgenie_alert.rst b/docs/apache-airflow-providers-opsgenie/operators/opsgenie_alert.rst index 8fad51f52f380..60f7e1c6bb1b4 100644 --- a/docs/apache-airflow-providers-opsgenie/operators/opsgenie_alert.rst +++ b/docs/apache-airflow-providers-opsgenie/operators/opsgenie_alert.rst @@ -48,3 +48,20 @@ Close alert in Opsgenie. :language: python :start-after: [START howto_opsgenie_close_alert_operator] :end-before: [END howto_opsgenie_close_alert_operator] + +.. _howto/operator:OpsgenieDeleteAlertOperator: + +OpsgenieDeleteAlertOperator +=========================== + +Use the :class:`~airflow.providers.opsgenie.operators.opsgenie.OpsgenieDeleteAlertOperator` to delete alert in opsgenie. + + +Using the Operator +^^^^^^^^^^^^^^^^^^ +Delete alert in Opsgenie. + +.. exampleinclude:: /../../airflow/providers/opsgenie/example_dags/example_opsgenie_alert.py + :language: python + :start-after: [START howto_opsgenie_delete_alert_operator] + :end-before: [END howto_opsgenie_delete_alert_operator] diff --git a/tests/providers/opsgenie/hooks/test_opsgenie.py b/tests/providers/opsgenie/hooks/test_opsgenie.py index 0c6c573044584..8963c1b5613c9 100644 --- a/tests/providers/opsgenie/hooks/test_opsgenie.py +++ b/tests/providers/opsgenie/hooks/test_opsgenie.py @@ -138,3 +138,19 @@ def test_close_alert(self, close_alert_mock): close_alert_payload=CloseAlertPayload(**pay_load), kwargs=kwargs, ) + + @mock.patch.object(AlertApi, 'delete_alert') + def test_delete_alert(self, delete_alert_mock): + hook = OpsgenieAlertHook(opsgenie_conn_id=self.conn_id) + + # When + identifier = 'identifier_example' + identifier_type = 'id' + user = "some_user" + source = "airflow" + + # Then + hook.delete_alert(identifier=identifier, identifier_type=identifier_type, user=user, source=source) + delete_alert_mock.assert_called_once_with( + identifier=identifier, identifier_type=identifier_type, user=user, source=source + ) diff --git a/tests/providers/opsgenie/operators/test_opsgenie.py b/tests/providers/opsgenie/operators/test_opsgenie.py index 444b19b1657bc..cbded2f04d4f2 100644 --- a/tests/providers/opsgenie/operators/test_opsgenie.py +++ b/tests/providers/opsgenie/operators/test_opsgenie.py @@ -18,11 +18,13 @@ # import unittest +from unittest import mock from airflow.models.dag import DAG from airflow.providers.opsgenie.operators.opsgenie import ( OpsgenieCloseAlertOperator, OpsgenieCreateAlertOperator, + OpsgenieDeleteAlertOperator, ) from airflow.utils import timezone @@ -141,3 +143,31 @@ def test_properties(self): assert self._config['user'] == operator.user assert self._config['note'] == operator.note assert self._config['source'] == operator.source + + +class TestOpsgenieDeleteAlertOperator(unittest.TestCase): + def setUp(self): + args = {'owner': 'airflow', 'start_date': DEFAULT_DATE} + self.dag = DAG('test_dag_id', default_args=args) + + @mock.patch('airflow.providers.opsgenie.operators.opsgenie.OpsgenieAlertHook') + def test_operator(self, mock_opsgenie_hook): + mock_opsgenie_hook.return_value = mock.Mock() + mock_opsgenie_hook.return_value.delete_alert.return_value = True + + operator = OpsgenieDeleteAlertOperator( + task_id='opsgenie_test_delete_job', + dag=self.dag, + identifier="id", + identifier_type='id', + user="name", + source="source", + ) + operator.execute(None) + mock_opsgenie_hook.assert_called_once_with('opsgenie_default') + mock_opsgenie_hook.return_value.delete_alert.assert_called_once_with( + identifier='id', + identifier_type='id', + source='source', + user='name', + ) From 725fa13ae6fe734bdf959eabe1c281ef8c13e197 Mon Sep 17 00:00:00 2001 From: eladkal <45845474+eladkal@users.noreply.github.com> Date: Wed, 4 May 2022 07:37:52 +0300 Subject: [PATCH 2/4] fixes --- airflow/providers/opsgenie/operators/opsgenie.py | 9 +++++---- .../operators/opsgenie_alert.rst | 1 + 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/airflow/providers/opsgenie/operators/opsgenie.py b/airflow/providers/opsgenie/operators/opsgenie.py index 4dee8d4292e59..af2d176fbe17a 100644 --- a/airflow/providers/opsgenie/operators/opsgenie.py +++ b/airflow/providers/opsgenie/operators/opsgenie.py @@ -212,7 +212,7 @@ def execute(self, context: 'Context') -> None: class OpsgenieDeleteAlertOperator(BaseOperator): """ - This operator allows you to Delete alerts to Opsgenie. + This operator allows you to delete alerts in Opsgenie. Accepts a connection that has an Opsgenie API key as the connection's password. This operator sets the domain to conn_id.host, and if not set will default to ``https://api.opsgenie.com``. @@ -232,6 +232,8 @@ class OpsgenieDeleteAlertOperator(BaseOperator): :param source: Display name of the request source """ + template_fields: Sequence[str] = 'identifier' + def __init__( self, *, @@ -249,12 +251,11 @@ def __init__( self.identifier_type = identifier_type self.user = user self.source = source - self.hook: Optional[OpsgenieAlertHook] = None def execute(self, context: 'Context') -> None: """Call the OpsgenieAlertHook to delete alert""" - self.hook = OpsgenieAlertHook(self.opsgenie_conn_id) - self.hook.delete_alert( + hook = OpsgenieAlertHook(self.opsgenie_conn_id) + hook.delete_alert( identifier=self.identifier, identifier_type=self.identifier_type, user=self.user, diff --git a/docs/apache-airflow-providers-opsgenie/operators/opsgenie_alert.rst b/docs/apache-airflow-providers-opsgenie/operators/opsgenie_alert.rst index 60f7e1c6bb1b4..7da63f55e58a0 100644 --- a/docs/apache-airflow-providers-opsgenie/operators/opsgenie_alert.rst +++ b/docs/apache-airflow-providers-opsgenie/operators/opsgenie_alert.rst @@ -63,5 +63,6 @@ Delete alert in Opsgenie. .. exampleinclude:: /../../airflow/providers/opsgenie/example_dags/example_opsgenie_alert.py :language: python + :dedent: 4 :start-after: [START howto_opsgenie_delete_alert_operator] :end-before: [END howto_opsgenie_delete_alert_operator] From 3b63ee97724ae6b7dc15baadcef8072245e0a52e Mon Sep 17 00:00:00 2001 From: eladkal <45845474+eladkal@users.noreply.github.com> Date: Wed, 4 May 2022 08:14:00 +0300 Subject: [PATCH 3/4] fix templated field --- airflow/providers/opsgenie/operators/opsgenie.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/opsgenie/operators/opsgenie.py b/airflow/providers/opsgenie/operators/opsgenie.py index af2d176fbe17a..7521316313a72 100644 --- a/airflow/providers/opsgenie/operators/opsgenie.py +++ b/airflow/providers/opsgenie/operators/opsgenie.py @@ -232,7 +232,7 @@ class OpsgenieDeleteAlertOperator(BaseOperator): :param source: Display name of the request source """ - template_fields: Sequence[str] = 'identifier' + template_fields: Sequence[str] = ('identifier',) def __init__( self, From 5a981c2b8c2d78bc943b173c4f7d03d55660217f Mon Sep 17 00:00:00 2001 From: eladkal <45845474+eladkal@users.noreply.github.com> Date: Fri, 6 May 2022 09:36:46 +0300 Subject: [PATCH 4/4] set defaults --- airflow/providers/opsgenie/hooks/opsgenie.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/providers/opsgenie/hooks/opsgenie.py b/airflow/providers/opsgenie/hooks/opsgenie.py index 78d023ef5d273..567e90c17eae1 100644 --- a/airflow/providers/opsgenie/hooks/opsgenie.py +++ b/airflow/providers/opsgenie/hooks/opsgenie.py @@ -136,9 +136,9 @@ def close_alert( def delete_alert( self, identifier: str, - identifier_type: Optional[str], - user: Optional[str], - source: Optional[str], + identifier_type: Optional[str] = None, + user: Optional[str] = None, + source: Optional[str] = None, ) -> SuccessResponse: """ Delete an alert in Opsgenie