import asyncio import json import random from datetime import datetime from typing import Any, Dict, List, Optional from loguru import logger from app.core.clickhouse import clickhouse_manager from app.models.token import BossToken from app.services.crawler.boss import BossService from app.services.crawler.qcwy import QcwyService from app.services.crawler.zhilian import ZhilianService class CompanyCleaner: def __init__(self): self.boss_service = BossService() self.qcwy_service = QcwyService() self.zhilian_service = ZhilianService() self._boss_token_loaded = False def _apply_proxy(self, proxy: Optional[str]) -> None: self.boss_service.set_proxy(proxy) self.qcwy_service.set_proxy(proxy) self.zhilian_service.set_proxy(proxy) async def _ensure_boss_token_loaded(self) -> None: if self._boss_token_loaded and self.boss_service.login_data.get("mpt"): return token_obj = await BossToken.filter(is_active=True).order_by("-updated_at").first() if not token_obj: logger.warning("BossToken not found or inactive in CompanyCleaner") return self.boss_service.set_login_data(token_obj.mpt or "", "") self._boss_token_loaded = True async def collect_pending_companies(self, limit: int = 1000, source: Optional[str] = None): client = await clickhouse_manager.get_client() logger.info(f"Starting to collect pending companies (limit={limit}, source={source or 'all'})...") if source is None or source == "zhilian": await self._collect_zhilian(client, limit) if source is None or source == "qcwy": await self._collect_qcwy(client, limit) if source is None or source == "boss": await self._collect_boss(client, limit) logger.info("Finished collecting pending companies.") async def _collect_zhilian(self, client, limit: int): logger.info("Collecting Zhilian companies...") # 优化：先获取已存在的公司ID，避免在子查询中读取json_data # 使用PREWHERE提前过滤时间范围，减少需要读取的数据量 # 检查90天内已处理的公司，避免重复请求 days_back_existing = 90 # 查询最近90天的数据，避免重复请求已处理过的公司 existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'companyNumber') as cid FROM job_data.zhilian_company PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'companyNumber') != '' LIMIT 50000 """ # 添加重试机制 existing_result = None existing_cids = set() # 默认使用空集合 for attempt in range(3): try: logger.info(f"Querying existing Zhilian companies (attempt {attempt+1})...") existing_result = await client.query(existing_companies_query) existing_cids = {row[0] for row in existing_result.result_rows if row[0]} break except Exception as e: error_str = str(e).lower() if "memory" in error_str or "memory_limit" in error_str: if attempt == 0: days_back_existing = 1 existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'companyNumber') as cid FROM job_data.zhilian_company PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'companyNumber') != '' LIMIT 5000 """ logger.warning(f"Memory error, reducing time range to {days_back_existing} days") elif attempt == 1: existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'companyNumber') as cid FROM job_data.zhilian_company SAMPLE 0.1 PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'companyNumber') != '' LIMIT 2000 """ logger.warning(f"Memory error persists, using SAMPLE 0.1") else: logger.error(f"Failed to query existing companies after {attempt+1} attempts: {e}") logger.warning("Using empty set for existing_cids, continuing with collection...") existing_cids = set() break else: logger.error(f"Non-memory error while querying existing companies: {e}") raise pending_query = "SELECT DISTINCT company_id FROM job_data.pending_company WHERE source = 'zhilian'" pending_result = await client.query(pending_query) pending_cids = {row[0] for row in pending_result.result_rows if row[0]} # 构建排除列表 exclude_cids = existing_cids | pending_cids # 优化：添加时间范围过滤，只查询最近30天的数据，减少扫描量 # 使用 PREWHERE 提前过滤时间范围，避免读取大量历史数据的 json_data # 增加 LIMIT 以便在 Python 中过滤后仍有足够的数据 query = f""" SELECT DISTINCT JSONExtractString(json_data, 'companyNumber') as cid, JSONExtractString(json_data, 'companyName') as cname FROM job_data.zhilian_job PREWHERE created_at > now() - INTERVAL 30 DAY WHERE json_data != '' AND JSONExtractString(json_data, 'companyNumber') != '' LIMIT {limit * 2} """ logger.info(f"Executing SQL for Zhilian (limit={limit * 2}): {query[:500]}...") result = await client.query(query) if not result.result_rows: return # 在 Python 中过滤掉已存在的和待处理的 rows: List[Dict[str, Any]] = [] for cid, cname in result.result_rows: if not cid or cid in exclude_cids: continue if len(rows) >= limit: break rows.append( { "source": "zhilian", "company_id": cid, "company_name": cname, "status": "pending", "created_at": datetime.now(), "updated_at": datetime.now(), } ) await self._insert_pending(client, rows) logger.info(f"Added {len(rows)} Zhilian companies to pending.") async def _collect_qcwy(self, client, limit: int): logger.info("Collecting QCWY companies...") # 优化：先获取已存在的公司ID，避免在子查询中读取json_data # 使用PREWHERE提前过滤时间范围，减少需要读取的数据量 # 检查90天内已处理的公司，避免重复请求 days_back_existing = 90 # 查询最近90天的数据，避免重复请求已处理过的公司 existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'companyId') as cid FROM job_data.qcwy_company PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'companyId') != '' LIMIT 50000 """ # 添加重试机制 existing_result = None existing_cids = set() # 默认使用空集合 for attempt in range(3): try: logger.info(f"Querying existing QCWY companies (attempt {attempt+1})...") existing_result = await client.query(existing_companies_query) # 查询成功，提取结果 existing_cids = {row[0] for row in existing_result.result_rows if row[0]} break except Exception as e: error_str = str(e).lower() if "memory" in error_str or "memory_limit" in error_str: if attempt == 0: # 第一次失败：进一步减少时间范围 days_back_existing = 1 existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'companyId') as cid FROM job_data.qcwy_company PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'companyId') != '' LIMIT 5000 """ logger.warning(f"Memory error, reducing time range to {days_back_existing} days") elif attempt == 1: # 第二次失败：使用采样 existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'companyId') as cid FROM job_data.qcwy_company SAMPLE 0.1 PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'companyId') != '' LIMIT 2000 """ logger.warning(f"Memory error persists, using SAMPLE 0.1") else: # 最后一次尝试也失败，使用空集合继续执行（避免阻塞整个流程） logger.error(f"Failed to query existing companies after {attempt+1} attempts: {e}") logger.warning("Using empty set for existing_cids, continuing with collection...") existing_cids = set() break else: # 其他错误直接抛出 logger.error(f"Non-memory error while querying existing companies: {e}") raise pending_query = "SELECT DISTINCT company_id FROM job_data.pending_company WHERE source = 'qcwy'" pending_result = await client.query(pending_query) pending_cids = {row[0] for row in pending_result.result_rows if row[0]} # 构建排除列表 exclude_cids = existing_cids | pending_cids # 优化策略： # 1. 减少时间范围：从30天减少到7天，大幅减少扫描的数据量 # 2. 减少LIMIT：从limit*2减少到更小的值，减少内存占用 # 3. 使用更严格的PREWHERE条件：先过滤时间，再过滤空json_data和超大JSON # 4. 限制JSON大小：过滤掉过大的json_data（可能包含大量嵌套数据） # 5. 分批查询：如果limit较大，分批处理，每次查询更少的数据 days_back = 7 # 从30天减少到7天，减少扫描量 # 注意：不使用length(json_data)检查，因为它需要读取整个列来计算长度 query_limit = min(limit * 2, 100) # 限制最大查询数量，避免内存超限 # 分批查询策略：如果limit较大，分批处理 result = None for attempt in range(3): # 最多尝试3次 try: # 根据尝试次数调整参数 if attempt == 1: # 第一次失败后：减少时间范围到3天 days_back = 3 query_limit = min(query_limit, 50) logger.warning(f"Retry {attempt}: Reducing time range to {days_back} days and limit to {query_limit}") elif attempt == 2: # 第二次失败后：使用采样 query = f""" SELECT DISTINCT JSONExtractString(json_data, 'coId') as cid, JSONExtractString(json_data, 'companyName') as cname FROM job_data.qcwy_job SAMPLE 0.1 PREWHERE created_at > now() - INTERVAL {days_back} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'coId') != '' LIMIT {query_limit} """ logger.warning(f"Retry {attempt}: Using SAMPLE 0.1 to reduce memory usage") result = await client.query(query) break # 正常查询或第一次重试 query = f""" SELECT DISTINCT JSONExtractString(json_data, 'coId') as cid, JSONExtractString(json_data, 'companyName') as cname FROM job_data.qcwy_job PREWHERE created_at > now() - INTERVAL {days_back} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'coId') != '' LIMIT {query_limit} """ logger.info(f"Executing SQL for QCWY (limit={query_limit}, days={days_back}, attempt={attempt+1}): {query[:400]}...") result = await client.query(query) break except Exception as e: error_str = str(e).lower() # 如果查询失败（可能是内存超限），继续重试 if "memory" in error_str or "memory_limit" in error_str: if attempt < 2: logger.warning(f"Memory error on attempt {attempt+1}: {e}") continue else: # 最后一次尝试也失败，抛出异常 logger.error(f"Query failed after {attempt+1} attempts: {e}") raise else: # 其他错误直接抛出 logger.error(f"Query failed with non-memory error: {e}") raise if not result or not result.result_rows: logger.info("No QCWY companies found in query result.") return # 在 Python 中过滤掉已存在的和待处理的 rows: List[Dict[str, Any]] = [] for cid, cname in result.result_rows: if not cid or cid in exclude_cids: continue if len(rows) >= limit: break rows.append( { "source": "qcwy", "company_id": cid, "company_name": cname, "status": "pending", "created_at": datetime.now(), "updated_at": datetime.now(), } ) if rows: await self._insert_pending(client, rows) logger.info(f"Added {len(rows)} QCWY companies to pending.") else: logger.info("No new QCWY companies found after filtering.") async def _collect_boss(self, client, limit: int): logger.info("Collecting Boss companies...") # 优化：先获取已存在的公司ID，避免在子查询中读取json_data # 使用PREWHERE提前过滤时间范围，减少需要读取的数据量 # 检查90天内已处理的公司，避免重复请求 days_back_existing = 90 # 查询最近90天的数据，避免重复请求已处理过的公司 existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'brandId') as cid FROM job_data.boss_company PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'brandId') != '' LIMIT 50000 """ # 添加重试机制 existing_result = None existing_cids = set() # 默认使用空集合 for attempt in range(3): try: logger.info(f"Querying existing Boss companies (attempt {attempt+1})...") existing_result = await client.query(existing_companies_query) existing_cids = {row[0] for row in existing_result.result_rows if row[0]} break except Exception as e: error_str = str(e).lower() if "memory" in error_str or "memory_limit" in error_str: if attempt == 0: days_back_existing = 1 existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'brandId') as cid FROM job_data.boss_company PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'brandId') != '' """ logger.warning(f"Memory error, reducing time range to {days_back_existing} days") elif attempt == 1: existing_companies_query = f""" SELECT DISTINCT JSONExtractString(json_data, 'brandId') as cid FROM job_data.boss_company SAMPLE 0.1 PREWHERE updated_at > now() - INTERVAL {days_back_existing} DAY AND json_data != '' WHERE JSONExtractString(json_data, 'brandId') != '' """ logger.warning(f"Memory error persists, using SAMPLE 0.1") else: logger.error(f"Failed to query existing companies after {attempt+1} attempts: {e}") logger.warning("Using empty set for existing_cids, continuing with collection...") existing_cids = set() break else: logger.error(f"Non-memory error while querying existing companies: {e}") raise pending_query = "SELECT DISTINCT company_id FROM job_data.pending_company WHERE source = 'boss'" pending_result = await client.query(pending_query) pending_cids = {row[0] for row in pending_result.result_rows if row[0]} # 构建排除列表 exclude_cids = existing_cids | pending_cids # 优化：添加时间范围过滤，只查询最近30天的数据，减少扫描量 # 使用 PREWHERE 提前过滤时间范围，避免读取大量历史数据的 json_data # 增加 LIMIT 以便在 Python 中过滤后仍有足够的数据 query = f""" SELECT DISTINCT JSONExtractString(json_data, 'brandId') as cid, JSONExtractString(json_data, 'brandName') as cname FROM job_data.boss_job PREWHERE created_at > now() - INTERVAL 30 DAY WHERE json_data != '' AND JSONExtractString(json_data, 'brandId') != '' LIMIT {limit * 2} """ logger.info(f"Executing SQL for Boss (limit={limit * 2}): {query[:500]}...") result = await client.query(query) if not result.result_rows: return # 在 Python 中过滤掉已存在的和待处理的 rows: List[Dict[str, Any]] = [] for cid, cname in result.result_rows: if not cid or cid in exclude_cids: continue if len(rows) >= limit: break rows.append( { "source": "boss", "company_id": cid, "company_name": cname, "status": "pending", "created_at": datetime.now(), "updated_at": datetime.now(), } ) await self._insert_pending(client, rows) logger.info(f"Added {len(rows)} Boss companies to pending.") async def _insert_pending(self, client, rows: List[Dict[str, Any]]): if not rows: return data: List[List[Any]] = [] for r in rows: data.append( [ r["source"], r["company_id"], r["company_name"], r["status"], "", r["created_at"], r["updated_at"], 1, ] ) await client.insert( "job_data.pending_company", data, column_names=[ "source", "company_id", "company_name", "status", "error_msg", "created_at", "updated_at", "version", ], ) async def process_single_company( self, source: str, company_id: str, proxy: Optional[str] = None, max_delay_seconds: int = 5, ) -> Dict[str, Any]: client = await clickhouse_manager.get_client() if proxy: self._apply_proxy(proxy) delay = 0 if max_delay_seconds and max_delay_seconds > 0: delay = random.randint(1, max_delay_seconds) if delay > 0: await asyncio.sleep(delay) query = f""" SELECT source, company_id, company_name, version FROM job_data.pending_company FINAL WHERE source = '{source}' AND company_id = '{company_id}' ORDER BY version DESC LIMIT 1 """ result = await client.query(query) if result.result_rows: source_value, cid, cname, version = result.result_rows[0] else: source_value = source cid = company_id cname = "" version = 0 try: success = await self._fetch_and_save(source_value, cid) status = "done" if success else "failed" error_msg = "" if success else "Fetch failed" except Exception as e: logger.error(f"Error processing {source_value} {cid}: {e}") status = "failed" error_msg = str(e) await client.insert( "job_data.pending_company", [ [ source_value, cid, cname, status, error_msg.replace("'", "''"), datetime.now(), datetime.now(), int(version) + 1, ] ], column_names=[ "source", "company_id", "company_name", "status", "error_msg", "created_at", "updated_at", "version", ], ) return { "success": status == "done", "source": source_value, "company_id": cid, "company_name": cname, "status": status, "error_msg": error_msg, "version": int(version) + 1, } async def process_pending_companies( self, limit: int = 100, source: Optional[str] = None, proxy: Optional[str] = None, max_delay_seconds: int = 0, ): client = await clickhouse_manager.get_client() logger.info(f"Processing pending companies (limit={limit}, source={source or 'all'})...") if proxy: self._apply_proxy(proxy) where_clause = "WHERE status = 'pending'" if source: where_clause += f" AND source = '{source}'" query = f""" SELECT source, company_id, company_name, version FROM job_data.pending_company FINAL {where_clause} ORDER BY created_at ASC LIMIT {limit} """ result = await client.query(query) if not result.result_rows: logger.info("No pending companies to process.") return for source_value, cid, cname, version in result.result_rows: logger.info(f"Processing {source_value} company: {cname} ({cid})") try: if max_delay_seconds and max_delay_seconds > 0: delay = random.randint(1, max_delay_seconds) if delay > 0: await asyncio.sleep(delay) success = await self._fetch_and_save(source_value, cid) status = "done" if success else "failed" error_msg = "" if success else "Fetch failed" except Exception as e: logger.error(f"Error processing {source_value} {cid}: {e}") status = "failed" error_msg = str(e) await client.insert( "job_data.pending_company", [ [ source_value, cid, cname, status, error_msg.replace("'", "''"), datetime.now(), datetime.now(), int(version) + 1, ] ], column_names=[ "source", "company_id", "company_name", "status", "error_msg", "created_at", "updated_at", "version", ], ) async def _fetch_and_save(self, source: str, company_id: str) -> bool: data: Optional[Dict[str, Any]] = None target_table = "" if source == "zhilian": data = self.zhilian_service.get_company_detail(company_id) target_table = "zhilian_company" elif source == "qcwy": data = self.qcwy_service.get_company_info(company_id) target_table = "qcwy_company" elif source == "boss": await self._ensure_boss_token_loaded() data = self.boss_service.get_company_detail_by_id(company_id) target_table = "boss_company" if not data: logger.error(f"No data returned from source={source} company_id={company_id}") return False try: logger.info( f"Raw company data from source={source} company_id={company_id}: " f"{json.dumps(data, ensure_ascii=False)[:2000]}" ) except Exception as e: logger.error(f"Failed to log raw company data for source={source} company_id={company_id}: {e}") client = await clickhouse_manager.get_client() name = "" if source == "zhilian": name = data.get("companyBase", {}).get("companyName", "") elif source == "qcwy": name = data.get("companyName", "") elif source == "boss": name = data.get("name", "") json_str = json.dumps(data, ensure_ascii=False) await client.insert( f"job_data.{target_table}", [[0, json_str, name, datetime.now(), datetime.now()]], column_names=["id", "json_data", "company_name", "created_at", "updated_at"], ) return True async def cleanup_old_records(self): """ 清理已完成或失败的记录 (每日调用) """ client = await clickhouse_manager.get_client() logger.info("Starting cleanup of processed pending companies...") # ClickHouse mutations are async, but lightweight for this purpose query = "ALTER TABLE job_data.pending_company DELETE WHERE status IN ('done', 'failed')" try: await client.command(query) logger.info("Cleanup command executed successfully.") except Exception as e: logger.error(f"Cleanup failed: {e}") company_cleaner = CompanyCleaner()