Skip to content

Commit

Permalink
Organize Opsgenie provider classes (#20454)
Browse files Browse the repository at this point in the history
* organize Opsgenie provider classes
  • Loading branch information
eladkal committed Dec 22, 2021
1 parent 636ae0a commit d57ebb7
Show file tree
Hide file tree
Showing 13 changed files with 278 additions and 218 deletions.
6 changes: 3 additions & 3 deletions airflow/contrib/hooks/opsgenie_alert_hook.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""This module is deprecated. Please use :mod:`airflow.providers.opsgenie.hooks.opsgenie_alert`."""
"""This module is deprecated. Please use :mod:`airflow.providers.opsgenie.hooks.opsgenie`."""

import warnings

from airflow.providers.opsgenie.hooks.opsgenie_alert import OpsgenieAlertHook # noqa
from airflow.providers.opsgenie.hooks.opsgenie import OpsgenieAlertHook # noqa

warnings.warn(
"This module is deprecated. Please use `airflow.providers.opsgenie.hooks.opsgenie_alert`.",
"This module is deprecated. Please use `airflow.providers.opsgenie.hooks.opsgenie`.",
DeprecationWarning,
stacklevel=2,
)
6 changes: 3 additions & 3 deletions airflow/contrib/operators/opsgenie_alert_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""This module is deprecated. Please use :mod:`airflow.providers.opsgenie.operators.opsgenie_alert`."""
"""This module is deprecated. Please use :mod:`airflow.providers.opsgenie.operators.opsgenie`."""

import warnings

from airflow.providers.opsgenie.operators.opsgenie_alert import OpsgenieAlertOperator # noqa
from airflow.providers.opsgenie.operators.opsgenie import OpsgenieAlertOperator # noqa

warnings.warn(
"This module is deprecated. Please use `airflow.providers.opsgenie.operators.opsgenie_alert`.",
"This module is deprecated. Please use `airflow.providers.opsgenie.operators.opsgenie`.",
DeprecationWarning,
stacklevel=2,
)
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from datetime import datetime

from airflow import DAG
from airflow.providers.opsgenie.operators.opsgenie_alert import OpsgenieAlertOperator
from airflow.providers.opsgenie.operators.opsgenie import OpsgenieAlertOperator

with DAG(
dag_id="opsgenie_alert_operator_dag",
Expand Down
100 changes: 100 additions & 0 deletions airflow/providers/opsgenie/hooks/opsgenie.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
#
# 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 typing import Optional

from opsgenie_sdk import (
AlertApi,
ApiClient,
Configuration,
CreateAlertPayload,
OpenApiException,
SuccessResponse,
)

from airflow.hooks.base import BaseHook


class OpsgenieAlertHook(BaseHook):
"""
This hook allows you to post alerts to Opsgenie.
Accepts a connection that has an Opsgenie API key as the connection's password.
This hook 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 hook.
:param opsgenie_conn_id: The name of the Opsgenie connection to use
:type opsgenie_conn_id: str
"""

conn_name_attr = 'opsgenie_conn_id'
default_conn_name = 'opsgenie_default'
conn_type = 'opsgenie'
hook_name = 'Opsgenie'

def __init__(self, opsgenie_conn_id: str = 'opsgenie_default') -> None:
super().__init__() # type: ignore[misc]
self.conn_id = opsgenie_conn_id
configuration = Configuration()
conn = self.get_connection(self.conn_id)
configuration.api_key['Authorization'] = conn.password
configuration.host = conn.host or 'https://api.opsgenie.com'
self.alert_api_instance = AlertApi(ApiClient(configuration))

def _get_api_key(self) -> str:
"""
Get the API key from the connection
:return: API key
:rtype: str
"""
conn = self.get_connection(self.conn_id)
return conn.password

def get_conn(self) -> AlertApi:
"""
Get the underlying AlertApi client
:return: AlertApi client
:rtype: opsgenie_sdk.AlertApi
"""
return self.alert_api_instance

def create_alert(self, payload: Optional[dict] = None) -> SuccessResponse:
"""
Create an alert on Opsgenie
:param payload: Opsgenie API Create Alert payload values
See https://docs.opsgenie.com/docs/alert-api#section-create-alert
:type payload: dict
:return: api response
:rtype: opsgenie_sdk.SuccessResponse
"""
payload = payload or {}

try:
create_alert_payload = CreateAlertPayload(**payload)
api_response = self.alert_api_instance.create_alert(create_alert_payload)
return api_response
except OpenApiException as e:
self.log.exception('Exception when sending alert to opsgenie with payload: %s', payload)
raise e
88 changes: 8 additions & 80 deletions airflow/providers/opsgenie/hooks/opsgenie_alert.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,86 +15,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

from typing import Optional

from opsgenie_sdk import (
AlertApi,
ApiClient,
Configuration,
CreateAlertPayload,
OpenApiException,
SuccessResponse,
)

from airflow.hooks.base import BaseHook


class OpsgenieAlertHook(BaseHook):
"""
This hook allows you to post alerts to Opsgenie.
Accepts a connection that has an Opsgenie API key as the connection's password.
This hook 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 hook.
:param opsgenie_conn_id: The name of the Opsgenie connection to use
:type opsgenie_conn_id: str
"""This module is deprecated. Please use :mod:`airflow.providers.opsgenie.hooks.opsgenie`."""

"""
import warnings

conn_name_attr = 'opsgenie_conn_id'
default_conn_name = 'opsgenie_default'
conn_type = 'opsgenie'
hook_name = 'Opsgenie'
from airflow.providers.opsgenie.hooks.opsgenie import OpsgenieAlertHook # noqa

def __init__(self, opsgenie_conn_id: str = 'opsgenie_default') -> None:
super().__init__() # type: ignore[misc]
self.conn_id = opsgenie_conn_id
configuration = Configuration()
conn = self.get_connection(self.conn_id)
configuration.api_key['Authorization'] = conn.password
configuration.host = conn.host or 'https://api.opsgenie.com'
self.alert_api_instance = AlertApi(ApiClient(configuration))

def _get_api_key(self) -> str:
"""
Get the API key from the connection
:return: API key
:rtype: str
"""
conn = self.get_connection(self.conn_id)
return conn.password

def get_conn(self) -> AlertApi:
"""
Get the underlying AlertApi client
:return: AlertApi client
:rtype: opsgenie_sdk.AlertApi
"""
return self.alert_api_instance

def create_alert(self, payload: Optional[dict] = None) -> SuccessResponse:
"""
Create an alert on Opsgenie
:param payload: Opsgenie API Create Alert payload values
See https://docs.opsgenie.com/docs/alert-api#section-create-alert
:type payload: dict
:return: api response
:rtype: opsgenie_sdk.SuccessResponse
"""
payload = payload or {}

try:
create_alert_payload = CreateAlertPayload(**payload)
api_response = self.alert_api_instance.create_alert(create_alert_payload)
return api_response
except OpenApiException as e:
self.log.exception('Exception when sending alert to opsgenie with payload: %s', payload)
raise e
warnings.warn(
"This module is deprecated. Please use `airflow.providers.opsgenie.hooks.opsgenie`.",
DeprecationWarning,
stacklevel=2,
)
144 changes: 144 additions & 0 deletions airflow/providers/opsgenie/operators/opsgenie.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
#
# 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 typing import Any, Dict, List, Optional

from airflow.models import BaseOperator
from airflow.providers.opsgenie.hooks.opsgenie import OpsgenieAlertHook


class OpsgenieAlertOperator(BaseOperator):
"""
This operator allows you to post 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:OpsgenieAlertOperator`
:param opsgenie_conn_id: The name of the Opsgenie connection to use
:type opsgenie_conn_id: str
:param message: The Message of the Opsgenie alert (templated)
:type message: str
:param alias: Client-defined identifier of the alert (templated)
:type alias: str
:param description: Description field of the alert (templated)
:type description: str
:param responders: Teams, users, escalations and schedules that
the alert will be routed to send notifications.
:type responders: list[dict]
:param visible_to: Teams and users that the alert will become visible
to without sending any notification.
:type visible_to: list[dict]
:param actions: Custom actions that will be available for the alert.
:type actions: list[str]
:param tags: Tags of the alert.
:type tags: list[str]
:param details: Map of key-value pairs to use as custom properties of the alert.
:type details: dict
:param entity: Entity field of the alert that is
generally used to specify which domain alert is related to. (templated)
:type entity: str
:param source: Source field of the alert. Default value is
IP address of the incoming request.
:type source: str
:param priority: Priority level of the alert. Default value is P3. (templated)
:type priority: str
:param user: Display name of the request owner.
:type user: str
:param note: Additional note that will be added while creating the alert. (templated)
:type note: str
"""

template_fields = ('message', 'alias', 'description', 'entity', 'priority', 'note')

def __init__(
self,
*,
message: str,
opsgenie_conn_id: str = 'opsgenie_default',
alias: Optional[str] = None,
description: Optional[str] = None,
responders: Optional[List[dict]] = None,
visible_to: Optional[List[dict]] = None,
actions: Optional[List[str]] = None,
tags: Optional[List[str]] = None,
details: Optional[dict] = None,
entity: Optional[str] = None,
source: Optional[str] = None,
priority: Optional[str] = None,
user: Optional[str] = None,
note: Optional[str] = None,
**kwargs,
) -> None:
super().__init__(**kwargs)

self.message = message
self.opsgenie_conn_id = opsgenie_conn_id
self.alias = alias
self.description = description
self.responders = responders
self.visible_to = visible_to
self.actions = actions
self.tags = tags
self.details = details
self.entity = entity
self.source = source
self.priority = priority
self.user = user
self.note = note
self.hook: Optional[OpsgenieAlertHook] = None

def _build_opsgenie_payload(self) -> Dict[str, Any]:
"""
Construct the Opsgenie JSON payload. All relevant parameters are combined here
to a valid Opsgenie JSON payload.
:return: Opsgenie payload (dict) to send
"""
payload = {}

for key in [
"message",
"alias",
"description",
"responders",
"visible_to",
"actions",
"tags",
"details",
"entity",
"source",
"priority",
"user",
"note",
]:
val = getattr(self, key)
if val:
payload[key] = val
return payload

def execute(self, context) -> None:
"""Call the OpsgenieAlertHook to post message"""
self.hook = OpsgenieAlertHook(self.opsgenie_conn_id)
self.hook.create_alert(self._build_opsgenie_payload())

0 comments on commit d57ebb7

Please sign in to comment.