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

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值