pytools/services/account_pool.py

# -*- coding: utf-8 -*-
"""
Created on 2023-04-24
---------
@summary:  竞品账号管理
---------
@author: Dzr
"""
from common.databases import mongo_table, object_id
from common.log import logger
from common.utils import now_date


class AccountManagePool:

    def __init__(self, db: str, table: str):
        self.account_tab = mongo_table(db, table)
        self.sleep_interval = 300
        self.output_log = True

        list_attr = dict(
            name='crawl_list',
            lock=dict(crawl_list=True),
            release=dict(crawl_list=False),
        )
        detail_attr = dict(
            name='crawl_detail',
            lock=dict(crawl_detail=True),
            release=dict(crawl_detail=False),
        )
        self._schedule = {'list': list_attr, 'detail': detail_attr}

    def _update_account(self, uid, **kwargs):
        kwargs['update_time'] = now_date()
        self.account_tab.update_one({'_id': uid}, {'$set': kwargs})

    def _find_account(self, query, *args, **kwargs):
        return self.account_tab.find_one(query, *args, **kwargs)

    def _find_and_update_account(self, query, update, *args, **kwargs):
        update['$set']['update_time'] = now_date()
        return self.account_tab.find_one_and_update(query, update, *args, **kwargs)

    def release_account(self, uid, crawl_type, ip):
        if isinstance(uid, str):
            uid = object_id(uid)

        query = dict(_id=uid)
        account = self._find_account(query)
        application = self._schedule[crawl_type]['name']
        response = {'status': 'duplicate'}
        if account is not None and account[application]:
            release = self._schedule[crawl_type]['release']
            self._update_account(uid, **release)
            logger.info(f'[账号池]释放账号:{uid},用途:{crawl_type}, {ip}')
            response = {'status': 'ok'}
        return response

    def lock_account(self, site, crawl_type, ip):
        release = self._schedule[crawl_type]['release']
        query = dict(site=site, **release)
        lock = self._schedule[crawl_type]['lock']
        update = {'$set': lock}
        sort = [('update_time', 1)]
        account = self._find_and_update_account(query, update, sort=sort)
        if account is not None:
            _account = {}
            for key, val in account.items():
                if key == '_id':
                    _account[key] = str(val)
                else:
                    _account[key] = val
            account = _account
            logger.info(f'[账号池]锁定账号:{account["_id"]},用途:{crawl_type},{ip}')
        return account