""" Web Backend - FastAPI 应用 提供 API 和 SSE 事件推送 """ import asyncio import json import os import subprocess from datetime import datetime from typing import Dict, List, Optional from fastapi import FastAPI, HTTPException, BackgroundTasks from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse from fastapi.staticfiles import StaticFiles from fastapi.templating import Jinja2Templates from fastapi.requests import Request from pydantic import BaseModel from typing import List as ListType from app import config from app.log_manager import get_log_manager from app.scraper import XiaoHongShuScraper # ========== Pydantic 模型 ========== class StartRequest(BaseModel): mode: str = "single" # single / multi max_threads: int = 3 class StatusResponse(BaseModel): is_running: bool mode: str max_threads: int total_users: int completed_users: int failed_users: int active_threads: int users_status: Dict[str, Dict] # ========== 全局状态 ========== class AppState: def __init__(self): self.is_running = False self.mode = "single" self.max_threads = 1 self.total_users = 0 self.completed_users = 0 self.failed_users = 0 self.active_threads = 0 self.users_status: Dict[str, Dict] = {} self.thread_status: Dict[str, Dict] = {} self.event_queue: asyncio.Queue = None self.scraper: Optional[XiaoHongShuScraper] = None self.scraper_task: Optional[asyncio.Task] = None app_state = AppState() log_manager = get_log_manager() # ========== FastAPI 应用 ========== app = FastAPI(title="小红书数据采集器", version="1.0.0") # 挂载静态文件和模板 app.mount("/static", StaticFiles(directory="static"), name="static") templates = Jinja2Templates(directory="templates") # ========== 路由 ========== @app.get("/", response_class=HTMLResponse) async def index(request: Request): """首页""" return templates.TemplateResponse("index.html", {"request": request}) @app.get("/api/status", response_model=StatusResponse) async def get_status(): """获取当前状态""" return StatusResponse( is_running=app_state.is_running, mode=app_state.mode, max_threads=app_state.max_threads, total_users=app_state.total_users, completed_users=app_state.completed_users, failed_users=app_state.failed_users, active_threads=app_state.active_threads, users_status=app_state.users_status ) @app.post("/api/start") async def start_scraping(request: StartRequest): """启动采集""" if app_state.is_running: raise HTTPException(status_code=400, detail="采集任务正在运行中") # 读取用户 ID 列表 users_file = os.path.join(config.DATA_DIR, "users.txt") if not os.path.exists(users_file): raise HTTPException(status_code=400, detail="用户文件不存在，请先创建 data/users.txt") with open(users_file, 'r', encoding='utf-8') as f: user_ids = [line.strip() for line in f if line.strip()] if not user_ids: raise HTTPException(status_code=400, detail="用户 ID 列表为空") # 验证模式 if request.mode not in ["single", "multi"]: raise HTTPException(status_code=400, detail="无效的模式，必须是 single 或 multi") # 验证线程数 if request.mode == "multi": if request.max_threads < 1 or request.max_threads > config.MAX_CONCURRENT_THREADS: raise HTTPException(status_code=400, detail=f"线程数必须在 1-{config.MAX_CONCURRENT_THREADS} 之间") # 重置状态 app_state.is_running = True app_state.mode = request.mode app_state.max_threads = request.max_threads if request.mode == "multi" else 1 app_state.total_users = len(user_ids) app_state.completed_users = 0 app_state.failed_users = 0 app_state.active_threads = 0 app_state.users_status = {} app_state.thread_status = {} # 初始化用户状态 for user_id in user_ids: app_state.users_status[user_id] = { "status": "pending", "note_count": 0, "thread_id": None } # 创建事件队列 app_state.event_queue = asyncio.Queue() # 创建采集器 app_state.scraper = XiaoHongShuScraper(app_state.event_queue, log_manager) # 启动采集任务 async def run_scraper(): try: await app_state.scraper.run( user_ids=user_ids, mode=request.mode, max_threads=request.max_threads ) finally: app_state.is_running = False await app_state.event_queue.put({ "type": "finished" }) app_state.scraper_task = asyncio.create_task(run_scraper()) return { "success": True, "message": f"采集任务已启动，模式: {request.mode}, 用户数: {len(user_ids)}", "total_users": len(user_ids), "mode": request.mode, "max_threads": app_state.max_threads } @app.post("/api/stop") async def stop_scraping(): """停止采集""" if not app_state.is_running: raise HTTPException(status_code=400, detail="没有正在运行的采集任务") if app_state.scraper: app_state.scraper.stop() # 取消任务 if app_state.scraper_task: app_state.scraper_task.cancel() try: await app_state.scraper_task except asyncio.CancelledError: pass app_state.is_running = False return { "success": True, "message": "采集任务已停止" } @app.post("/api/logs/clear") async def clear_logs(): """清空日志""" log_manager.clear() return { "success": True, "message": "日志已清空" } @app.get("/api/logs") async def get_logs(level: Optional[str] = None, limit: Optional[int] = None): """获取日志""" logs = log_manager.get_logs(level=level, limit=limit) return { "logs": logs, "stats": log_manager.get_stats() } @app.get("/api/logs/search") async def search_logs(keyword: str, limit: Optional[int] = None): """搜索日志""" logs = log_manager.search_logs(keyword, limit) return { "logs": logs, "count": len(logs) } @app.get("/api/logs/download") async def download_logs(): """下载日志文件""" text = log_manager.export_to_text() return StreamingResponse( iter([text]), media_type="text/plain", headers={ "Content-Disposition": f"attachment; filename=xiaohongshu_logs_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt" } ) @app.get("/api/users") async def get_users(): """获取用户列表""" users_file = os.path.join(config.DATA_DIR, "users.txt") if not os.path.exists(users_file): return {"users": []} with open(users_file, 'r', encoding='utf-8') as f: users = [{"id": line.strip()} for line in f if line.strip()] # 添加采集状态 for user in users: user_id = user["id"] if user_id in app_state.users_status: user.update(app_state.users_status[user_id]) else: user["status"] = "pending" user["note_count"] = 0 return {"users": users} @app.get("/api/data/{user_id}") async def get_user_data(user_id: str): """获取用户采集数据""" file_path = os.path.join(config.DATA_DIR, f"{user_id}.json") if not os.path.exists(file_path): raise HTTPException(status_code=404, detail="用户数据不存在") with open(file_path, 'r', encoding='utf-8') as f: data = json.load(f) return data @app.get("/api/events") async def event_stream(): """SSE 事件流""" async def event_generator(): while True: try: # 等待事件，超时 1 秒发送心跳 event = await asyncio.wait_for( app_state.event_queue.get(), timeout=1.0 ) # 更新状态 if event["type"] == "log": # 日志不需要更新状态 pass elif event["type"] == "progress": app_state.completed_users = event.get("completed", app_state.completed_users) app_state.failed_users = event.get("failed", app_state.failed_users) app_state.active_threads = event.get("active_threads", app_state.active_threads) elif event["type"] == "user_status": if event["user_id"] in app_state.users_status: app_state.users_status[event["user_id"]]["status"] = event["status"] app_state.users_status[event["user_id"]]["note_count"] = event.get("note_count", 0) app_state.users_status[event["user_id"]]["thread_id"] = event.get("thread_id") elif event["type"] == "thread_status": app_state.thread_status[event["thread_id"]] = { "status": event["status"], "current_user": event.get("current_user") } # 更新活跃线程数 app_state.active_threads = sum( 1 for s in app_state.thread_status.values() if s.get("status") == "busy" ) # 发送事件 yield f"data: {json.dumps(event)}\n\n" except asyncio.TimeoutError: # 发送心跳 yield ": heartbeat\n\n" except Exception as e: # 错误事件 yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n" return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "Access-Control-Allow-Origin": "*" } ) # ========== Git 相关 API ========== @app.get("/api/git/status") async def git_status(): """获取 Git 状态""" if not os.path.exists("/app/.git"): return { "enabled": False, "message": "不是 Git 仓库" } try: # 获取当前分支 result = subprocess.run( ["git", "branch", "--show-current"], capture_output=True, text=True, cwd="/app" ) current_branch = result.stdout.strip() # 获取最新提交 result = subprocess.run( ["git", "log", "-1", "--format=%h %s %cr"], capture_output=True, text=True, cwd="/app" ) latest_commit = result.stdout.strip() # 检查是否有未提交的更改 result = subprocess.run( ["git", "status", "--porcelain"], capture_output=True, text=True, cwd="/app" ) has_changes = len(result.stdout.strip()) > 0 return { "enabled": True, "branch": current_branch, "latest_commit": latest_commit, "has_changes": has_changes, "auto_pull_enabled": config.ENABLE_GIT_PULL, "auto_pull_interval": config.AUTO_PULL_INTERVAL } except Exception as e: return { "enabled": False, "error": str(e) } @app.post("/api/git/pull") async def git_pull(): """手动拉取最新代码""" if not os.path.exists("/app/.git"): raise HTTPException(status_code=400, detail="不是 Git 仓库") try: # 获取当前分支 result = subprocess.run( ["git", "branch", "--show-current"], capture_output=True, text=True, cwd="/app" ) current_branch = result.stdout.strip() # 拉取最新代码 result = subprocess.run( ["git", "fetch", "origin", current_branch], capture_output=True, text=True, cwd="/app" ) if result.returncode != 0: raise HTTPException(status_code=500, detail=f"拉取失败: {result.stderr}") # 检查是否有更新 result = subprocess.run( ["git", "rev-list", "HEAD..origin/" + current_branch, "--count"], capture_output=True, text=True, cwd="/app" ) commits_behind = int(result.stdout.strip()) if commits_behind == 0: return { "success": True, "message": "已是最新代码", "needs_restart": False } # 重置到最新代码 result = subprocess.run( ["git", "reset", "--hard", "origin/" + current_branch], capture_output=True, text=True, cwd="/app" ) if result.returncode != 0: raise HTTPException(status_code=500, detail=f"重置失败: {result.stderr}") # 检查 Python 文件是否有变化 result = subprocess.run( ["git", "diff", "--name-only", "HEAD@{1}", "HEAD"], capture_output=True, text=True, cwd="/app" ) changed_files = result.stdout.strip().split('\n') python_changed = any(f.endswith('.py') for f in changed_files if f) return { "success": True, "message": f"已拉取 {commits_behind} 个新提交", "changed_files": changed_files, "needs_restart": python_changed, "restart_hint": "Python 代码有变化，建议重启容器以应用更新" } except HTTPException: raise except Exception as e: raise HTTPException(status_code=500, detail=f"拉取失败: {str(e)}") @app.post("/api/git/restart") async def restart_container(): """ 请求重启容器 注意：这需要在容器外配置重启策略 """ return { "success": True, "message": "重启请求已接收", "methods": [ "docker restart xiaohongshu-scraper", "docker-compose restart" ] } # ========== 启动事件 ========== @app.on_event("startup") async def startup_event(): """启动时初始化""" # 确保数据目录存在 os.makedirs(config.DATA_DIR, exist_ok=True) os.makedirs(config.SCREENSHOT_DIR, exist_ok=True) # 初始化事件队列 app_state.event_queue = asyncio.Queue() print("=" * 60) print("小红书数据采集器 Web 服务启动") print(f"访问地址: http://{config.WEB_HOST}:{config.WEB_PORT}") print("=" * 60) @app.on_event("shutdown") async def shutdown_event(): """关闭时清理""" if app_state.scraper: app_state.scraper.stop() if app_state.scraper_task: app_state.scraper_task.cancel() print("Web 服务已关闭") if __name__ == "__main__": import uvicorn uvicorn.run( "app.web:app", host=config.WEB_HOST, port=config.WEB_PORT, reload=True )