2.0/start_run_yaoshibang_v2.py

"""药师帮统一入口 — 根据任务 snapshot_collect_status 字段自动切换普通/快照模式。
支持账号池轮换：一个账号被锁自动换下一个。
"""
import random
import time
import json

from spiders.yaoshibang.ysbang_crawl import YsbSpider
from spiders.yaoshibang.ysb_snapshot_crawl import YaoShiBangSnapshot
from commons.scheduler import CrawlerScheduler
from commons.conn_mysql import MySQLPoolOn2
from commons.Logger import logger
from commons.feishu_webhook import send_text

PLATFORM_NAME = "药师帮"
PLATFORM_ID = 5

# 账号池
_account_pool = []
_account_index = 0


def _load_account_pool():
    """从 DB 加载所有可用药师帮账号"""
    global _account_pool, _account_index
    db = MySQLPoolOn2()
    sql = """SELECT `name`, `phone`, `password` FROM `accounts_platform`
             WHERE `platform`=%s AND `status`=1 AND `equipment_id`=3
             ORDER BY `cookie_timestamp` ASC"""
    rows = db.select_data(sql, (PLATFORM_ID,))
    if not rows:
        logger.error("无可用药师帮账号（platform=%s, status=1）", PLATFORM_ID)
        return False
    _account_pool = [(r["phone"], r["password"], r["name"]) for r in rows]
    _account_index = 0
    logger.info("从 DB 加载 %s 个药师帮账号", len(_account_pool))
    return True


def _get_current_account():
    """获取当前账号，池空则重新加载"""
    global _account_pool, _account_index
    if not _account_pool:
        if not _load_account_pool():
            return None
    return _account_pool[_account_index]


def _rotate_account():
    """换下一个账号，返回 (phone, password, name)"""
    global _account_index
    if not _account_pool:
        if not _load_account_pool():
            return None
    _account_index = (_account_index + 1) % len(_account_pool)
    new_acc = _account_pool[_account_index]
    logger.warning("账号被锁，切换: %s → %s (第%s/%s个)",
                   _account_pool[(_account_index - 1) % len(_account_pool)][2],
                   new_acc[2], _account_index + 1, len(_account_pool))
    return new_acc


def _build_task_dict(api_task, device_id, account_password, account_name, scheduler):
    """将调度 API 返回的原始任务字段组装为 spider 需要的 task_dict。"""
    current_page = api_task.get("current_page", 0)
    start_page = current_page if current_page > 0 else 1
    task_dict = {
        # 新 API 字段
        "id": api_task.get("id"),
        "company_id": api_task.get("company_id"),
        "product_name": api_task.get("product_name", ""),
        "product_specs": api_task.get("product_specs", ""),
        "product_brand": api_task.get("product_brand", ""),
        "product_keyword": api_task.get("product_keyword", ""),
        "current_page": current_page,
        "start_offset": api_task.get("start_offset"),
        # 兼容旧字段
        "start_page": start_page,
        "end_page": 0,
        # spider 逐页回告用
        "scheduler": scheduler,
        "_device_id": device_id,
        "_platform": str(PLATFORM_ID),
        "collect_task_id": str(api_task.get("id", "")),
        "snapshot_collect_status": api_task.get("snapshot_collect_status"),
        "sampling_cycle": "",
        "sampling_start_time": "",
        "sampling_end_time": "",
        "collect_equipment_account_id": "",
        "collect_region_id": "",
        "collect_equipment_id": "",
        "collect_round": 1,
        # 入口传入的账号信息，spider 可复用避免重复查 DB
        "account_phone": device_id,
        "account_password": account_password,
        "account_name": account_name,
    }
    return task_dict


def _report_progress(scheduler, task_id, is_finished, crawled_count=0,
                     current_page=0, total_pages=0,
                     exception_type=None, remark=None):
    """通过调度器 API 上报任务进度。"""
    data = {
        "task_id": task_id,
        "platform": str(PLATFORM_ID),
        "username": scheduler.username,
        "current_page": current_page,
        "is_finished": is_finished,
        "total_pages": total_pages,
        "crawled_count": crawled_count,
    }
    if exception_type is not None:
        data["exception_type"] = exception_type
    if remark:
        data["remark"] = remark
    try:
        scheduler.post_report(data)
        logger.info("任务状态上报成功: is_finished=%s, crawled_count=%s",
                    is_finished, crawled_count)
    except Exception as e:
        logger.warning("任务状态上报失败: %s", e)



if __name__ == "__main__":
    # 1. 加载账号池
    if not _load_account_pool():
        logger.error("无可用药师帮账号，退出")
        exit(1)

    idle_seconds = random.randint(180, 300)
    empty_count = 0          # 连续拉不到任务的次数
    max_empty_before_rotate = 2  # 连续 2 次拉不到就换号

    # 2. 初始化调度器
    phone, password, name = _get_current_account()
    scheduler = CrawlerScheduler(DEVICE_ID=phone, platform=str(PLATFORM_ID))
    scheduler.start()
    time.sleep(2)
    logger.info("调度器已启动，当前账号: %s (phone=%s)", name, phone)

    logger.info("循环任务已启动，平台=%s，每轮间隔 %s 秒", PLATFORM_NAME, idle_seconds)

    while True:
        try:
            logger.info("开始执行%s爬虫任务", PLATFORM_NAME)
            api_task = scheduler.get_task()

            if not api_task:
                empty_count += 1
                logger.info("%s暂无任务 (连续%s次)，等待 %s 秒后重试",
                            PLATFORM_NAME, empty_count, idle_seconds)
                if empty_count >= max_empty_before_rotate:
                    # 换个号试试
                    new_acc = _rotate_account()
                    if new_acc:
                        scheduler.stop()
                        phone, password, name = new_acc
                        scheduler = CrawlerScheduler(DEVICE_ID=phone, platform=str(PLATFORM_ID))
                        scheduler.start()
                        empty_count = 0
                        logger.info("已切换到账号: %s", name)
                time.sleep(idle_seconds)
                continue

            empty_count = 0  # 拉到任务了，重置

            task_id = api_task.get("id", "")
            task_dict = _build_task_dict(api_task, phone, password, name, scheduler)

            log_dict = {k: v for k, v in task_dict.items() if k != "scheduler"}
            logger.info("拉取到任务: %s", json.dumps(log_dict, ensure_ascii=False))

            # 根据 snapshot_collect_status 选爬虫
            snapshot_status = task_dict.get("snapshot_collect_status", 0)
            if not snapshot_status:
                spider_cls = YaoShiBangSnapshot
                logger.info("任务 %s 启用快照模式", task_id)
            else:
                spider_cls = YsbSpider
                logger.info("任务 %s 使用普通模式", task_id)

            result = spider_cls(task_dict).run()
            if len(result) == 3:
                crawl_count, is_success, total_pages = result
            else:
                crawl_count, is_success = result
                total_pages = 0

            send_text(
                f"{time.strftime('%Y-%m-%d %H:%M:%S')} 通知:\n"
                f"平台: {PLATFORM_NAME}, 药品: {task_dict.get('product_name')}, "
                f"爬取数据: {crawl_count}条"
            )

            if is_success:
                _report_progress(scheduler, task_id, is_finished=1,
                                 crawled_count=crawl_count,
                                 current_page=total_pages, total_pages=total_pages)
            else:
                _report_progress(scheduler, task_id, is_finished=0,
                                 crawled_count=crawl_count,
                                 current_page=total_pages, total_pages=total_pages)

        except Exception as e:
            logger.error("%s爬虫任务执行失败: %s", PLATFORM_NAME, e, exc_info=True)
        time.sleep(idle_seconds)