钉钉 log handler

dingtalk_handler.py

# -*- coding: utf-8 -*-
# created:     Fri Dec 29 10:56:33 CST 2017
# filename:    dingtalk_handler.py
# author:      juntao liu
# email:       [email protected]
# descritpion:

import json
import logging
import threading
import time
import datetime
import traceback

import requests

local = threading.local()

logger = logging.getLogger(__name__)

class DingTalkMessageBuffer:
    '''如果超过限制则，自动缓存钉钉消息，统一在到下一周期发送
    '''

    def __init__(self, period, rate):
        self.period = period
        self.rate = rate

        self.message_count = 0
        self.message_buffer = []
        self.period_timestamp = self.get_period_timestamp(self.period)

    @staticmethod
    def get_period_timestamp(period):
        '''根据周期，计算当前周期的unix timestamp
        '''
        now = int(time.time())
        return now - now % period

    def push(self, title, content):
        '''压入钉钉消息，如果超过限制，放入缓存，否则返回钉钉机器人data
        '''
        data = None
        message = "# {} at {}\n{}\n".format(
            title,
            datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            content
        )

        current_timestamp = self.get_period_timestamp(self.period)
        if self.period_timestamp != current_timestamp:
            # 非当前周期
            if self.message_buffer:
                # 有缓存的消息
                title = "{} with archived {}".format(title, len(self.message_buffer))
                for _message in self.message_buffer:
                    message += _message
            self.message_count = 0
            self.message_buffer = []
            self.period_timestamp = current_timestamp

        if self.message_count < self.rate:
            # 没有超过限额
            data = {
                'msgtype': 'markdown',
                'markdown': {
                    'title': title,
                    'text': message
                }
            }
            logger.info("under rate limit, generate dingtalk message:%s", message)
        else:
            logger.warning("over rate limit(rate:%s, current:%s, period:%s), buffer message",
                           self.rate,
                           self.message_count,
                           self.period)
            self.message_buffer.append(message)
        self.message_count += 1
        return data

class DingTalkHandler(logging.Handler):
    """钉钉机器人日志处理
    """

    def __init__(self, access_token, title, timeout=0.2, period=60, rate=20):
        '''
        钉钉机器人日志handler

        默认限制60秒(period)发送20条(rate)消息。
        当前周期超过rate-1条消息之后，消息会被合并，并且在下一个period合并发送。

        已知问题：
        1在一些边缘情况下会丢失合并的消息
        1.1 在当前周期结束之后没有新的消息进入。
        1.2 当前周期还未结束，程序就退出了。

        2 多进程依然会超过钉钉消息限制。因为现在使用得失 threading.local

        Args:
          access_token (str): 钉钉机器人access_token
          title (str): 钉钉消息标题
          timeout (float): 发送钉钉机器人消息超时时间
          period (int):  消息限制条数
          rate (int): 消息限制时间
        '''

        self.access_token = access_token
        self.title = title
        self.timeout = timeout

        self.dingtalk_message_buffer = getattr(local, 'dingtalk_message_buffer', None)
        if not self.dingtalk_message_buffer:
            self.dingtalk_message_buffer = DingTalkMessageBuffer(period, rate)
            local.dingtalk_message_buffer = self.dingtalk_message_buffer

        super(DingTalkHandler, self).__init__()


    def send(self, data):
        print("send", data)
        try:
            resp = requests.post(
                'https://oapi.dingtalk.com/robot/send?access_token={}'.format(self.access_token),
                data=json.dumps(data),
                headers={
                    'Content-Type': 'application/json'
                },
                timeout=self.timeout,
            )
        except requests.exceptions.Timeout:
            logger.error("send dingtalk robot message timeout")

    def emit(self, record):
        content = "{}\n{}".format(
            self.format(record),
            traceback.format_exc()
        )
        data = self.dingtalk_message_buffer.push(self.title, content)
        if data:
            return self.send(data)
        return True


if __name__ == "__main__":
    log = logging.getLogger()
    dingtalk = DingTalkHandler(
        "abc079eea1168990885b12e1ca272ac2ba5becd528617fd4a9a5e8527d65b402",
        'debug dingtalk handler',
        period=1,
        rate=1)
    dingtalk.setLevel(logging.ERROR)

    log.addHandler(dingtalk)
    try:
        hah = asdfa
    except Exception:
        for i in range(10):
            log.error('test')
        for i in range(10):
            time.sleep(0.5)
            log.error('test')