piaoquan_crawler/common/limit.py

import os
import sys
from datetime import datetime, timedelta

sys.path.append(os.getcwd())

from common.aliyun_log import AliyunLogger
from common.scheduling_db import MysqlHelper
from common.db import RedisClient


def generate_expire_time():
    """
    计算出过期时间
    """
    now = datetime.now()

    # 当天晚上12点的时间
    midnight = datetime(now.year, now.month, now.day) + timedelta(days=1)

    # 计算当前时间到当天晚上12点的秒数
    seconds_until_midnight = int((midnight - now).total_seconds())
    return seconds_until_midnight


class AuthorLimit(object):
    """
    账号爬虫限量
    """

    def __init__(self, mode, platform):
        self.mode = mode
        self.platform = platform
        self.limit_tag_dict = {"352": "余海涛", "353": "罗情", "53": "范军", "51": "鲁涛"}

    def find_tag(self, uid):
        """
        通过 uid 去找符合标准的 tag
        """
        sql = f"""select tag from crawler_user_v3 where uid={uid};"""
        result = MysqlHelper.get_values(
            log_type=self.mode, crawler=self.platform, env="prod", sql=sql
        )
        tags = result[0]["tag"]
        if tags:
            tags = tags.split(",")
            for tag in tags:
                if self.limit_tag_dict.get(tag):
                    return tag
        return None

    def author_limitation(self, user_id):
        """
        限制账号， 服务长沙四名同学
        """
        if self.mode == "author":
            tag = self.find_tag(user_id)
            if tag:
                AliyunLogger.logging(
                    code="8807",
                    platform=self.platform,
                    mode=self.mode,
                    env="prod",
                    message="找到个人账号，{}".format(tag)
                )
                R = RedisClient()
                if R.connect():
                    tag_count = R.select(tag)
                    if tag_count:
                        tag_count = int(tag_count.decode("utf-8"))
                        if tag_count <= 300:
                            tag_count += 1
                            expire_seconds = generate_expire_time()
                            R.insert(
                                key=tag, value=tag_count, expire_time=expire_seconds
                            )
                            return True
                        else:
                            # 报警
                            return False
                    else:
                        tag_count = 1
                        expire_seconds = generate_expire_time()
                        R.insert(
                            key=tag, value=tag_count, expire_time=expire_seconds
                        )
                        return True
                else:
                    return True
            else:
                return True
        else:
            return True