from datetime import datetime import json from typing import Optional from fastapi import APIRouter, File, UploadFile, Form, Body, Query from tortoise.expressions import Q from tortoise.functions import Count, Sum from app.models.cleaning import CleaningTask from app.models.company import CompanyCleaningQueue from app.schemas import Success, SuccessExtra from app.services.cleaning import CleaningService from app.services.company_cleaner import company_cleaner from app.services.company_storage import company_storage, normalize_company_id from app.controllers.cleaning import cleaning_controller router = APIRouter() cleaning_service = CleaningService() COMPANY_SOURCES = ("boss", "qcwy", "zhilian") @router.get("/stats", summary="获取公司清洗统计信息") async def get_stats(): """获取 MySQL 中待处理公司的统计信息""" pending_count = await CompanyCleaningQueue.filter(status="pending").count() today_count = await CompanyCleaningQueue.filter( status="done", updated_at__gte=datetime.now().replace(hour=0, minute=0, second=0, microsecond=0), ).count() dist_rows = await CompanyCleaningQueue.annotate(total=Count("id")).group_by("source", "status").values( "source", "status", "total", ) stats = { "total_pending": pending_count, "today_processed": today_count, "details": [] } source_stats = { source: { "pending": 0, "processing": 0, "done": 0, "failed": 0, "total": 0, "jobs_fetched": 0, "jobs_stored": 0, "jobs_duplicate": 0, "jobs_failed": 0, } for source in COMPANY_SOURCES } for row in dist_rows: source = row["source"] status = row["status"] count = row["total"] if source not in source_stats: continue if status in source_stats[source]: source_stats[source][status] = count source_stats[source]["total"] += count job_dist_rows = await CompanyCleaningQueue.annotate( jobs_fetched_sum=Sum("jobs_fetched"), jobs_stored_sum=Sum("jobs_stored"), jobs_duplicate_sum=Sum("jobs_duplicate"), jobs_failed_sum=Sum("jobs_failed"), ).group_by("source").values( "source", "jobs_fetched_sum", "jobs_stored_sum", "jobs_duplicate_sum", "jobs_failed_sum", ) for row in job_dist_rows: source = row["source"] if source not in source_stats: continue source_stats[source]["jobs_fetched"] = int(row["jobs_fetched_sum"] or 0) source_stats[source]["jobs_stored"] = int(row["jobs_stored_sum"] or 0) source_stats[source]["jobs_duplicate"] = int(row["jobs_duplicate_sum"] or 0) source_stats[source]["jobs_failed"] = int(row["jobs_failed_sum"] or 0) stats["details"] = [{"source": k, **v} for k, v in source_stats.items()] return Success(data=stats) @router.get("/companies", summary="获取公司清洗列表") async def get_companies_list( page: int = Query(1, ge=1), page_size: int = Query(20, ge=1, le=100), source: Optional[str] = Query(None), status: Optional[str] = Query(None) ): """分页获取待处理公司列表详情""" offset = (page - 1) * page_size VALID_SOURCES = {"boss", "qcwy", "zhilian"} VALID_STATUSES = {"pending", "processing", "done", "failed"} queue_query = CompanyCleaningQueue.all() if source: if source not in VALID_SOURCES: return SuccessExtra(data=[], total=0, page=page, page_size=page_size) queue_query = queue_query.filter(source=source) if status: if status not in VALID_STATUSES: return SuccessExtra(data=[], total=0, page=page, page_size=page_size) queue_query = queue_query.filter(status=status) total = await queue_query.count() rows = await queue_query.order_by("-updated_at").offset(offset).limit(page_size) data = [] for row in rows: data.append({ "source": row.source, "company_id": row.company_id, "company_name": row.company_name, "status": row.status, "error_msg": row.error_msg, "retry_count": row.retry_count, "started_at": row.started_at.isoformat() if row.started_at else None, "finished_at": row.finished_at.isoformat() if row.finished_at else None, "jobs_fetched": row.jobs_fetched, "jobs_stored": row.jobs_stored, "jobs_duplicate": row.jobs_duplicate, "jobs_failed": row.jobs_failed, "jobs_error_msg": row.jobs_error_msg, "created_at": row.created_at.isoformat() if row.created_at else None, "updated_at": row.updated_at.isoformat() if row.updated_at else None }) return SuccessExtra(data=data, total=total, page=page, page_size=page_size) @router.get("/company-detail", summary="获取公司清洗详情") async def get_company_cleaning_detail( source: str = Query(..., description="数据源"), company_id: str = Query(..., description="公司ID"), company_name: Optional[str] = Query(None, description="公司名称"), ): if source not in {"boss", "qcwy", "zhilian"}: return Success(code=400, msg="不支持的数据源") model = company_storage.company_model(source) normalized_id = normalize_company_id(source, company_id) row = await model.get_or_none(source_company_id=normalized_id) if not row and company_name: row = await model.get_or_none(company_name=company_name) if not row: return Success(code=404, msg="未找到公司清洗结果") data = row.raw_json if isinstance(data, str): try: data = json.loads(data) except Exception: data = {"raw": data} return Success( data={ "source": source, "company_id": row.source_company_id, "company_name": row.company_name, "created_at": row.created_at.isoformat() if row.created_at else None, "updated_at": row.updated_at.isoformat() if row.updated_at else None, "data": data, } ) @router.post("/collect-pending", summary="分析待处理数据") async def collect_pending_companies_api( limit: int = Body(1000, embed=True, ge=1, le=10000), source: Optional[str] = Body(None, embed=True) ): """ 分析招聘数据，收集待处理的公司ID到 MySQL 队列表 """ summary = await company_cleaner.collect_pending_companies(limit=limit, source=source) return Success( msg=f"已完成数据分析，本次新增 {summary['total_created']} 条待处理公司", data=summary, ) @router.post("/run-pending", summary="手动执行待处理公司清洗") async def run_pending_companies( limit: int = Body(100, embed=True, ge=1, le=5000), source: Optional[str] = Body(None, embed=True), proxy: Optional[str] = Body(None, embed=True), max_delay_seconds: int = Body(5, embed=True), ): """ 手动触发待处理公司清洗任务 仅会处理当前状态为 pending 的记录，已经处理过的记录不会重复执行。 """ await company_cleaner.process_pending_companies( limit=limit, source=source, proxy=proxy, max_delay_seconds=max_delay_seconds, ) return Success(msg=f"已触发执行最近 {limit} 条待处理公司清洗任务") @router.post("/crawl-execute", summary="爬取并执行待处理公司清洗") async def crawl_execute_pending( limit: int = Body(100, embed=True, ge=1, le=5000), source: Optional[str] = Body(None, embed=True), proxy: Optional[str] = Body(None, embed=True), max_delay_seconds: int = Body(5, embed=True), ): await company_cleaner.collect_pending_companies(limit=limit, source=source) await company_cleaner.process_pending_companies( limit=limit, source=source, proxy=proxy, max_delay_seconds=max_delay_seconds, ) return Success(msg=f"已触发爬取并执行最近 {limit} 条待处理公司清洗任务") @router.post("/update-company-status", summary="更新公司爬取状态（爬虫端调用）") async def update_company_status( source: str = Body(..., embed=True), company_id: str = Body(..., embed=True), status: str = Body(..., embed=True), error_message: str = Body("", embed=True), ): """爬虫完成公司数据抓取后，调用此接口更新 MySQL 队列状态""" VALID_STATUSES = {"done", "failed"} if status not in VALID_STATUSES: return Success(msg=f"无效状态: {status}，仅支持 {VALID_STATUSES}", code=400) normalized_id = normalize_company_id(source, company_id) queue, _ = await CompanyCleaningQueue.get_or_create( source=source, company_id=normalized_id, defaults={ "company_name": "", "status": "pending", "error_msg": "", "retry_count": 0, "started_at": None, "finished_at": None, }, ) queue.status = status queue.error_msg = error_message or "" queue.finished_at = datetime.now() if status == "failed": queue.retry_count += 1 await queue.save() return Success(msg="状态更新成功", data={"source": source, "company_id": company_id, "status": status}) @router.post("/process-company", summary="执行单个公司清洗任务") async def process_single_company_api( source: str = Body(..., embed=True), company_id: str = Body(..., embed=True), proxy: Optional[str] = Body(None, embed=True), max_delay_seconds: int = Body(5, embed=True), ): result = await company_cleaner.process_single_company( source=source, company_id=company_id, proxy=proxy, max_delay_seconds=max_delay_seconds, ) success = bool(result.get("success")) msg = "任务执行成功" if success else "任务执行失败" return Success(msg=msg, data=result) @router.post("/upload", summary="上传文件并保存任务") async def upload_file( file: UploadFile = File(...), clean_type: str = Form("auto"), platform: str = Form("auto"), proxy: Optional[str] = Form(None) ): targets = await cleaning_service.parse_file(file) tasks = [ CleaningTask( target=t, clean_type=clean_type, platform=platform, proxy=proxy, status="pending" ) for t in targets ] if tasks: await CleaningTask.bulk_create(tasks) return Success(msg=f"Successfully imported {len(tasks)} tasks") @router.get("/list", summary="获取清洗任务列表") async def list_tasks( page: int = Query(1, description="页码"), page_size: int = Query(10, description="每页数量"), target: str = Query(None, description="目标搜索"), status: str = Query(None, description="状态筛选"), clean_type: str = Query(None, description="清洗类型筛选") ): q = Q() if target: q &= Q(target__contains=target) if status: q &= Q(status=status) if clean_type: q &= Q(clean_type=clean_type) total, tasks = await cleaning_controller.list(page=page, page_size=page_size, search=q, order=["-created_at"]) data = [await t.to_dict() for t in tasks] return SuccessExtra(data=data, total=total, page=page, page_size=page_size) @router.post("/process/{task_id}", summary="处理单个任务") async def process_task(task_id: int): task = await cleaning_controller.get(id=task_id) if not task: return Success(code=404, msg="Task not found") task.status = "processing" await task.save() result = await cleaning_service.process_single_item( target=task.target, clean_type=task.clean_type, platform=task.platform, proxy=task.proxy, ) task.status = "success" if result.get("success") else "fail" task.storage_status = result.get("storage_status", "unknown") task.remote_sent = result.get("remote_sent", False) task.result_summary = result.get("data_summary") task.error_msg = result.get("error") await task.save() return Success(data=await task.to_dict(), msg="Task processed") @router.delete("/delete", summary="删除任务") async def delete_task( id: int = Query(..., description="任务ID") ): await cleaning_controller.remove(id=id) return Success(msg="Deleted Successfully") @router.post("/clear", summary="清空所有任务") async def clear_tasks(): await CleaningTask.all().delete() return Success(msg="All tasks cleared")