AIRFLOW 企业微信发送
编写wechat_hook.py
import json
from typing import List, Optional, Union
import requests
from requests import Session
from airflow.exceptions import AirflowException
from airflow.providers.http.hooks.http import HttpHook
class WechatHook(HttpHook):
conn_name_attr = 'wechat_conn_id'
default_conn_name = 'wechat_default'
conn_type = 'http'
hook_name = 'Wechat'
def __init__(
self,
wechat_conn_id='wechat_default',
message_type: str = 'text',
message: Optional[Union[str, dict]] = None,
at_mobiles: Optional[List[str]] = None,
at_all: bool = False,
*args,
**kwargs,
) -> None:
super().__init__(http_conn_id=wechat_conn_id, *args, **kwargs) # type: ignore[misc]
self.message_type = message_type
self.message = message
self.at_mobiles = at_mobiles
self.at_all = at_all
def _get_endpoint(self) -> str:
"""Get WeChat endpoint for sending message."""
conn = self.get_connection(self.http_conn_id)
token = conn.password
if not token:
raise AirflowException(
'WeChat token is requests but get nothing, check you conn_id configuration.'
)
return f'cgi-bin/webhook/send?key={token}'
def _build_message(self) -> str:
"""
Build different type of WeChat message
As most commonly used type, text message just need post message content
rather than a dict like ``{'content': 'message'}``
"""
if self.message_type in ['text', 'markdown']:
data = {
'msgtype': self.message_type,
self.message_type: {'content': self.message} if self.message_type == 'text' else self.message,
'at': {'atMobiles': self.at_mobiles, 'isAtAll': self.at_all},
}
else:
data = {'msgtype': self.message_type, self.message_type: self.message}
return json.dumps(data)
def get_conn(self, headers: Optional[dict] = None) -> Session:
"""
Overwrite HttpHook get_conn because just need base_url and headers and
not don't need generic params
:param headers: additional headers to be passed through as a dictionary
"""
conn = self.get_connection(self.http_conn_id)
self.base_url = conn.host if conn.host else 'https://qyapi.weixin.qq.com'
session = requests.Session()
if headers:
session.headers.update(headers)
return session
def send(self) -> None:
"""Send WeChat message"""
support_type = ['text', 'link', 'markdown', 'actionCard', 'feedCard']
if self.message_type not in support_type:
raise ValueError(
f'WeChatWebhookHook only support {support_type} so far, but receive {self.message_type}'
)
data = self._build_message()
self.log.info('Sending WeChat type %s message %s', self.message_type, data)
resp = self.run(
endpoint=self._get_endpoint(), data=data, headers={'Content-Type': 'application/json'}
)
# WeChat success send message will with errcode equal to 0
if int(resp.json().get('errcode')) != 0:
raise AirflowException(f'Send WeChat message failed, receive error message {resp.text}')
self.log.info('Success Send WeChat message')
编写wechat_operator.py
from typing import TYPE_CHECKING, List, Optional, Sequence, Union
from airflow.models import BaseOperator
from wechat_hook import WechatHook
if TYPE_CHECKING:
from airflow.utils.context import Context
class WechatOperator(BaseOperator):
template_fields: Sequence[str] = ('message',)
ui_color = '#4ea4d4' # Wechat icon color
def __init__(
self,
*,
wechat_conn_id: str = 'wechat_default',
message_type: str = 'text',
message: Union[str, dict, None] = None,
at_mobiles: Optional[List[str]] = None,
at_all: bool = False,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.wechat_conn_id = wechat_conn_id
self.message_type = message_type
self.message = message
self.at_mobiles = at_mobiles
self.at_all = at_all
def execute(self, context: 'Context') -> None:
self.log.info('Sending WeChat message.')
hook = WechatHook(
self.wechat_conn_id,
self.message_type,
self.message,
self.at_mobiles,
self.at_all
)
hook.send()
测试文件
from wechat_operator import WechatOperator
context='22222'
message = '1111111'
WechatOperator(
task_id='wechat_task',
wechat_conn_id='wechat_default',
message_type='text',
message=message,
at_all=True,
).execute(context)
参考1:https://blog.csdn.net/weixin_40791632/article/details/125317428
参考2:https://blog.csdn.net/qq_31405633/article/details/100521402