crawler_task_management/main.py

import os
import yaml
import uuid
import asyncio
import uvicorn
from pathlib import Path
from typing import Dict, Any, Optional, List
from fastapi import FastAPI, HTTPException, Depends

from public_function.public_func import read_config
from model.model import CrawlerTask, DataReceive
from model.model import AccountCreate, AccountUpdate, DeviceResetData

app = FastAPI()


def get_config():
    """获取配置文件"""
    config_path = os.path.join(Path(__file__).resolve().parent, 'public_function/config.yaml')
    try:
        # 这里假设read_config函数存在
        from public_function.public_func import read_config
        return read_config(config_path)
    except ImportError:
        logger.warning("未找到read_config函数，使用默认配置")
        return {'host': 'localhost', 'port': 3306, 'user': 'root', 'password': 'password', 'db': 'test_db', 'max_overflow': 10}


def get_account_manager():
    """获取账号管理器实例"""
    config = get_config()
    try:
        from account_management.deal_account import DealAccount
        return DealAccount(config)
    except ImportError:
        logger.warning("未找到DealAccount类，返回模拟实例")
    return None


def get_task_manager():
    """获任务管理器实例"""
    config = get_config()
    try:
        from task_management.all_task_management import AllTask
        return AllTask(config)
    except ImportError:
        logger.warning("未找到AllTask类，返回模拟实例")
    return None


# 账号处理相关
@app.get("/obtain_account", summary="获取可用账号")
async def obtain_account(app_name: str, country: str, account_manager: Any = Depends(get_account_manager)):
    """
    获取指定应用的可用账号
    - **app_name**: 应用名称
    """
    if not app_name or not app_name.strip():
        raise HTTPException(status_code=400, detail="应用名称不能为空")
    if not country or not country.strip():
        raise HTTPException(status_code=400, detail="国家不能为空会")
    try:
        result = await account_manager.obtain_account_info(app_name, country)
        if result:
            return {"code": 200, "message": "获取账号成功", "data": result[0]}
        else:
            raise HTTPException(status_code=404, detail="没有可用的账号")
    except Exception as e:
        print(f"获取账号失败: {e}")
        raise HTTPException(status_code=404, detail="{}".format(e))


@app.post("/add_account", summary="新增账号")
async def add_account(account_data: AccountCreate, account_manager: Any = Depends(get_account_manager)):
    """
    新增爬虫账号
    - **account_id**: 账号ID
    - **password**: 密码
    - **app_name**: 应用名称
    """
    try:
        print(account_data.dict())
        # 这里应该调用实际的添加账号方法
        result = await account_manager.add_account([account_data.dict()])
        if result:
            return {"code": 200, "message": "新增账号成功", "data": result}
        raise HTTPException(status_code=404, detail="新增账号失败")
    except Exception as e:
        print(f"新增账号失败: {e}")
        raise HTTPException(status_code=500, detail="新增账号失败,失败原因：{}".format(e))


@app.post("/update_account", summary="删除账号")
async def update_account(account_data: AccountUpdate, account_manager: Any = Depends(get_account_manager)):
    """
    删除爬虫账号
    - **account_id**: 账号ID
    - **password**: 密码
    - **app_name**: 应用名称
    """
    try:
        data = account_data.dict()
        set_param = {"status": data["status"]}
        # params = {"account_id": data["account_id"], "app_name": data["app_name"]}
        params = (data["account_id"], data["app_name"])
        where_conditions = "account_id = %s and app_name = %s "
        # # 这里应该调用实际的添加账号方法
        result = await account_manager.update_account_info(set_param=set_param, where_conditions=where_conditions, params=params)
        if result:
            return {"code": 200, "message": "删除账号状态修改成功", "data": result}
        raise HTTPException(status_code=404, detail="删除账号状态修改失败")
    except Exception as e:
        print(f"新增账号失败: {e}")
        raise HTTPException(status_code=500, detail="删除账号失败，请重试,失败原因：{}".format(e))


@app.post("/receive_data", summary="接收抓取得数据")
async def receive_data(task_data: DataReceive, task_manager: Any = Depends(get_task_manager)):
    """数据接收接口"""
    try:
        params = task_data.dict()
        print(params)
        result = await task_manager.deal_receive_data(params)
        if result:
            return {"code": 200, "message": "数据保存成功", "data": result}
        raise HTTPException(status_code=404, detail="抓取商品数据失败，请重新尝试")
    except Exception as e:
        print(f"<UNK>: {e}")
        raise HTTPException(status_code=500, detail="获取数据失败；失败原因{}".format(e))


@app.get("/crawler_task", summary="爬虫任务调用接口")
async def crawler_task(task_data: CrawlerTask, task_manager: Any = Depends(get_task_manager)):
    """爬虫任务接口"""
    try:
        params = task_data.dict()
        params['task_id'] = uuid.uuid4().hex
        result = await task_manager.task_distribution(params)
        if result:
            return {"code": 200, "message": "<UNK>", "data": result}
        raise HTTPException(status_code=404, detail="抓取商品数据失败，请重新尝试")
    except Exception as e:
        print(f"<UNK>: {e}")
        raise HTTPException(status_code=500, detail="获取数据失败；失败原因{}".format(e))


# @app.get("/device_reset", summary="爬虫任务调用接口")
# async def device_reset(task_data: CrawlerTask, reset_manager: Any = Depends(get_reset_manager)):
#     """设备重置接口"""
#     try:
#         params = task_data.dict()
#         params['task_id'] = uuid.uuid4().hex
#         # 将任务记录到mysql
#         # result = await task_manager.task_distribution(params)
#         # if result:
#         #     return {"code": 200, "message": "<UNK>", "data": result}
#         raise HTTPException(status_code=404, detail="抓取商品数据失败，请重新尝试")
#     except Exception as e:
#         print(f"<UNK>: {e}")
#         raise HTTPException(status_code=500, detail="获取数据失败；失败原因{}".format(e))


if __name__ == '__main__':
    uvicorn.run(app)