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 是否存在changsha_user_accounts中
        """
        sql = f"""select user_name from changsha_user_accounts where piaoquan_account_id = {uid};"""
        result = MysqlHelper.get_values(
            log_type=self.mode, crawler=self.platform, env="prod", sql=sql
        )
        return result

    def author_limitation(self, user_id):
        """
        限制账号， 服务长沙四名同学
        """
        if self.mode == "author":
            result = self.find_tag(user_id)
            if result:
                user_name = result[0]['user_name']
                AliyunLogger.logging(
                    code="8807",
                    platform=self.platform,
                    mode=self.mode,
                    env="prod",
                    message="找到个人账号，{}".format(user_name)
                )
                R = RedisClient()
                if R.connect():
                    tag = self.limit_tag_dict[user_name]
                    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