""" 潜在客户导入服务 后台任务：分页调用 company API，逐个获取详情并存储到本地数据库 """ import time import threading from datetime import datetime from typing import Optional, Dict, Any from sqlalchemy.orm import Session from database import SessionLocal from models import ImportTask, PotentialCustomer, PotentialContact from fumamx_cli import FumamxCLIWrapper class ImportService: """导入服务单例类""" _instance = None _lock = threading.Lock() def __new__(cls): if cls._instance is None: with cls._lock: if cls._instance is None: cls._instance = super().__new__(cls) cls._instance._init() return cls._instance def _init(self): """初始化""" self.cli = FumamxCLIWrapper() self.current_task_id: Optional[int] = None self._cancel_event = threading.Event() self._thread: Optional[threading.Thread] = None def get_running_task(self) -> Optional[Dict]: """获取当前运行的任务""" if self.current_task_id is None: return None db = SessionLocal() try: task = db.query(ImportTask).filter( ImportTask.id == self.current_task_id ).first() if task and task.status == "running": return task.to_dict() else: # 任务已完成或不存在 self.current_task_id = None return None finally: db.close() def get_latest_task(self) -> Optional[Dict]: """获取最后一次任务（无论状态）""" db = SessionLocal() try: task = db.query(ImportTask).order_by( ImportTask.created_at.desc() ).first() return task.to_dict() if task else None finally: db.close() def start_import(self, keyword: str, page_size: int = 10) -> Dict: """启动导入任务""" # 检查是否有运行中的任务 if self.current_task_id is not None: running = self.get_running_task() if running: return { "success": False, "error": "已有导入任务正在运行", "task": running } # 创建新任务 db = SessionLocal() try: # 先测试搜索，获取总数量 search_result = self.cli.search_company(keyword, page=1, page_size=page_size) if search_result.get("code") != "0": return { "success": False, "error": f"搜索失败: {search_result.get('msg', '未知错误')}" } total_count = int(search_result.get("data", {}).get("total", 0)) if total_count == 0: return { "success": False, "error": "未找到匹配的企业" } total_pages = (total_count + page_size - 1) // page_size # 创建任务记录 task = ImportTask( keyword=keyword, status="pending", total_pages=total_pages, page_size=page_size, total_count=total_count, processed_count=0, success_count=0, failed_count=0, duplicate_count=0 ) db.add(task) db.commit() db.refresh(task) # 启动后台线程 self.current_task_id = task.id self._cancel_event.clear() self._thread = threading.Thread( target=self._import_worker, args=(task.id, keyword, total_pages, page_size), daemon=True ) self._thread.start() return { "success": True, "task": task.to_dict() } except Exception as e: return { "success": False, "error": f"启动失败: {str(e)}" } finally: db.close() def cancel_import(self) -> Dict: """取消当前导入任务""" if self.current_task_id is None: return { "success": False, "error": "没有正在运行的任务" } self._cancel_event.set() # 更新任务状态 db = SessionLocal() try: task = db.query(ImportTask).filter( ImportTask.id == self.current_task_id ).first() if task: task.status = "cancelled" task.completed_at = datetime.now() db.commit() return { "success": True, "task": task.to_dict() } finally: db.close() return {"success": True} def _import_worker(self, task_id: int, keyword: str, total_pages: int, page_size: int): """后台导入工作线程""" db = SessionLocal() error_logs = [] # 收集错误日志 try: # 更新任务状态为运行中 task = db.query(ImportTask).filter(ImportTask.id == task_id).first() if not task: return task.status = "running" task.started_at = datetime.now() db.commit() # 逐页处理 for page in range(1, total_pages + 1): # 检查取消标志 if self._cancel_event.is_set(): task.status = "cancelled" task.completed_at = datetime.now() task.error_msg = "\n".join(error_logs[-10:]) if error_logs else None # 保留最后10条错误 db.commit() break # 更新当前页码 task.current_page = page db.commit() try: # 1. 搜索当前页企业列表 search_result = self.cli.search_company(keyword, page=page, page_size=page_size) if search_result.get("code") != "0": error_msg = f"第{page}页搜索失败: {search_result.get('msg', '未知错误')}" error_logs.append(error_msg) print(error_msg) task.failed_count += page_size task.processed_count += page_size db.commit() continue companies = search_result.get("data", {}).get("results", []) # 逐个处理企业 for company in companies: # 检查取消标志 if self._cancel_event.is_set(): break domain = company.get('domain', 'unknown') try: self._process_company(db, task, company, keyword) task.success_count += 1 except Exception as e: task.failed_count += 1 error_detail = f"[{domain}] {type(e).__name__}: {str(e)}" error_logs.append(error_detail) print(f"处理企业失败: {error_detail}") # 实时更新错误信息到数据库（最多保留最近10条） task.error_msg = "\n".join(error_logs[-10:]) task.processed_count += 1 db.commit() # API 限速：每家企业处理完后等待 4 秒（1秒搜索 + 3个详情API各1秒） time.sleep(4) if self._cancel_event.is_set(): break except Exception as e: page_error = f"第{page}页处理失败: {type(e).__name__}: {str(e)}" error_logs.append(page_error) print(page_error) task.failed_count += page_size # 粗略估算 task.processed_count += page_size task.error_msg = "\n".join(error_logs[-10:]) db.commit() # 每页之间额外等待 time.sleep(2) # 任务完成 if task.status == "running": task.status = "completed" task.completed_at = datetime.now() task.error_msg = "\n".join(error_logs[-10:]) if error_logs else None db.commit() except Exception as e: task.status = "failed" error_logs.append(f"任务异常: {type(e).__name__}: {str(e)}") task.error_msg = "\n".join(error_logs[-10:]) task.completed_at = datetime.now() db.commit() finally: self.current_task_id = None db.close() def _process_company(self, db: Session, task: ImportTask, company: Dict, keyword: str): """处理单个企业""" external_id = company.get("id") domain = company.get("domain") if not external_id or not domain: return # 检查是否已存在（去重） existing = db.query(PotentialCustomer).filter( PotentialCustomer.domain == domain, PotentialCustomer.external_id == external_id ).first() if existing: task.duplicate_count += 1 return # 2. 获取企业详情 info_result = self.cli.get_company_info(external_id, domain) info_data = info_result.get("data", {}) if info_result.get("code") == 0 else {} # 3. 获取电话/图片信息（可选，失败不影响主流程） try: phone_result = self.cli.get_company_phones(domain) except: phone_result = {} # 4. 获取联系人列表 contacts_result = self.cli.get_company_contacts(external_id, domain) contacts_data = contacts_result.get("data", {}) if contacts_result.get("code") == 0 else {} contacts_list = contacts_data.get("list") or [] # 创建企业记录 customer = PotentialCustomer( external_id=external_id, domain=domain, name=company.get("title") or info_data.get("name"), country=company.get("country") or info_data.get("country"), country_en=company.get("countryEn") or info_data.get("countryEn"), country_cn=company.get("countryCnName") or info_data.get("countryCn"), industry=info_data.get("industry"), employees=info_data.get("employees"), founded=info_data.get("founded"), address=info_data.get("address"), phone=info_data.get("phone"), description=company.get("description") or info_data.get("description"), website=company.get("website"), linkedin_url=info_data.get("linkedin_url"), procurement_frequency=company.get("procurementFrequency"), annual_purchase_amount=company.get("annualPurchaseAmount"), hs_codes=company.get("hsCode") or [], hs_entities=company.get("hsEntity") or [], keyword=keyword, source_score=company.get("score"), source_db=company.get("sourceDb"), import_task_id=task.id, contact_count=len(contacts_list) ) db.add(customer) db.flush() # 获取 customer.id # 创建联系人记录 for contact in contacts_list: pc = PotentialContact( customer_id=customer.id, name=contact.get("name"), position=contact.get("position"), email=contact.get("email"), phone=contact.get("phone") or contact.get("realWhatsApp"), linkedin=contact.get("linkedin"), facebook=contact.get("facebook"), twitter=contact.get("twitter"), whatsapp=contact.get("whatsapp"), score=contact.get("score"), is_executive=bool(contact.get("isExecutive")), is_kp=bool(contact.get("isKp")), is_purchase=bool(contact.get("isPurchase")), email_check_status=contact.get("emailCheckStatus"), whatsapp_status=contact.get("whatsAppStatus") ) db.add(pc) db.commit() # 全局导入服务实例 import_service = ImportService()