first commit
This commit is contained in:
+516
@@ -0,0 +1,516 @@
|
||||
"""
|
||||
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
|
||||
)
|
||||
Reference in New Issue
Block a user