131 lines
4.7 KiB
Python
131 lines
4.7 KiB
Python
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 AccountCreate, AccountUpdate, CrawlerTask, AccountDelete
|
||
|
||
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()])
|
||
return {"code": 200, "message": "新增账号成功", "data": result}
|
||
except Exception as e:
|
||
print(f"新增账号失败: {e}")
|
||
raise HTTPException(status_code=500, detail="新增账号失败,失败原因：{}".format(e))
|
||
|
||
|
||
@app.post("/delete_account", summary="删除账号")
|
||
async def delete_account(account_data: AccountDelete, account_manager: Any = Depends(get_account_manager)):
|
||
"""
|
||
删除爬虫账号
|
||
- **account_id**: 账号ID
|
||
- **password**: 密码
|
||
- **app_name**: 应用名称
|
||
"""
|
||
try:
|
||
print(account_data.dict())
|
||
# 这里应该调用实际的添加账号方法
|
||
result = await account_manager.delete_account(account_data.dict())
|
||
return {"code": 200, "message": "删除账号成功", "data": result}
|
||
except Exception as e:
|
||
print(f"新增账号失败: {e}")
|
||
raise HTTPException(status_code=500, detail="删除账号失败，请重试,失败原因：{}".format(e))
|
||
|
||
|
||
@app.post("/receive_data", summary="接收抓取得数据")
|
||
async def receive_data(params: Dict[str, Any]):
|
||
"""数据接收接口"""
|
||
# 接收数据表写入redis后，存入mysql
|
||
|
||
|
||
@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))
|
||
|
||
|
||
if __name__ == '__main__':
|
||
uvicorn.run(app)
|