""" 爬取任务执行器 负责执行爬取任务并更新任务状态 """ import threading import time from django.utils import timezone from django.db import transaction from core.models import CrawlTask from core.utils import crawl_by_keyword, crawl_historical_articles, full_site_crawler, WEBSITE_SEARCH_CONFIGS class TaskExecutor: """任务执行器""" def __init__(self): self.running_tasks = {} self.lock = threading.Lock() def start_task(self, task_id): """启动任务""" with self.lock: if task_id in self.running_tasks: return False, "任务已在运行中" try: task = CrawlTask.objects.get(id=task_id) if task.status != 'pending': return False, "任务状态不允许启动" # 更新任务状态 task.status = 'running' task.started_at = timezone.now() task.save() # 启动后台线程执行任务 thread = threading.Thread(target=self._execute_task, args=(task_id,)) thread.daemon = True thread.start() self.running_tasks[task_id] = thread return True, "任务已启动" except CrawlTask.DoesNotExist: return False, "任务不存在" except Exception as e: return False, f"启动任务失败: {e}" def cancel_task(self, task_id): """取消任务""" with self.lock: if task_id in self.running_tasks: # 标记任务为取消状态 try: task = CrawlTask.objects.get(id=task_id) task.status = 'cancelled' task.completed_at = timezone.now() task.save() # 移除运行中的任务 del self.running_tasks[task_id] return True, "任务已取消" except CrawlTask.DoesNotExist: return False, "任务不存在" else: return False, "任务未在运行中" def _execute_task(self, task_id): """执行任务的核心逻辑""" try: task = CrawlTask.objects.get(id=task_id) # 根据任务类型执行不同的爬取逻辑 if task.task_type == 'keyword': self._execute_keyword_task(task) elif task.task_type == 'historical': self._execute_historical_task(task) elif task.task_type == 'full_site': self._execute_full_site_task(task) else: raise ValueError(f"不支持的任务类型: {task.task_type}") # 任务完成 with transaction.atomic(): task = CrawlTask.objects.select_for_update().get(id=task_id) task.status = 'completed' task.completed_at = timezone.now() task.progress = 100 task.save() except Exception as e: # 任务失败 try: with transaction.atomic(): task = CrawlTask.objects.select_for_update().get(id=task_id) task.status = 'failed' task.completed_at = timezone.now() task.error_message = str(e) task.save() except: pass finally: # 清理运行中的任务记录 with self.lock: if task_id in self.running_tasks: del self.running_tasks[task_id] def _execute_keyword_task(self, task): """执行关键词搜索任务""" # 更新当前操作 task.current_action = "开始关键词搜索" task.save() # 准备参数 selected_websites = task.websites.all() if selected_websites: websites = [w.name for w in selected_websites] else: websites = list(WEBSITE_SEARCH_CONFIGS.keys()) start_date = task.start_date.strftime('%Y-%m-%d') if task.start_date else None end_date = task.end_date.strftime('%Y-%m-%d') if task.end_date else None # 执行爬取 results = crawl_by_keyword( keyword=task.keyword, website_names=websites, max_pages=task.max_pages, start_date=start_date, end_date=end_date, max_articles=task.max_articles ) # 更新结果 task.total_articles = results['total_articles'] task.success_count = results['success_count'] task.failed_count = results['failed_count'] task.result_details = results['website_results'] task.save() def _execute_historical_task(self, task): """执行历史文章任务""" # 更新当前操作 task.current_action = "开始历史文章爬取" task.save() # 准备参数 selected_websites = task.websites.all() if selected_websites: websites = [w.name for w in selected_websites] else: websites = list(WEBSITE_SEARCH_CONFIGS.keys()) start_date = task.start_date.strftime('%Y-%m-%d') if task.start_date else None end_date = task.end_date.strftime('%Y-%m-%d') if task.end_date else None # 执行爬取 results = crawl_historical_articles( website_names=websites, start_date=start_date, end_date=end_date, max_articles_per_site=task.max_articles ) # 更新结果 task.total_articles = results['total_articles'] task.success_count = results['success_count'] task.failed_count = results['failed_count'] task.result_details = results['website_results'] task.save() def _execute_full_site_task(self, task): """执行全站爬取任务""" # 更新当前操作 task.current_action = "开始全站爬取" task.save() # 准备参数 selected_websites = task.websites.all() if selected_websites: websites = [w.name for w in selected_websites] else: websites = list(WEBSITE_SEARCH_CONFIGS.keys()) total_websites = len(websites) completed_websites = 0 for website_name in websites: try: # 更新当前网站 task.current_website = website_name task.current_action = f"正在爬取 {website_name}" task.save() # 获取或创建网站对象 from core.models import Website website, created = Website.objects.get_or_create( name=website_name, defaults={ 'base_url': WEBSITE_SEARCH_CONFIGS[website_name]["search_url"], 'enabled': True } ) # 执行全站爬取 full_site_crawler( WEBSITE_SEARCH_CONFIGS[website_name]["search_url"], website, max_pages=task.max_pages ) completed_websites += 1 progress = int((completed_websites / total_websites) * 100) task.progress = progress task.save() except Exception as e: # 记录错误但继续处理其他网站 print(f"爬取网站 {website_name} 时出错: {e}") continue # 更新最终结果 task.total_articles = completed_websites # 这里可以改为实际爬取的文章数 task.success_count = completed_websites task.failed_count = total_websites - completed_websites task.save() def get_task_status(self, task_id): """获取任务状态""" try: task = CrawlTask.objects.get(id=task_id) return { 'status': task.status, 'progress': task.progress, 'current_website': task.current_website, 'current_action': task.current_action, 'total_articles': task.total_articles, 'success_count': task.success_count, 'failed_count': task.failed_count, 'error_message': task.error_message } except CrawlTask.DoesNotExist: return None # 全局任务执行器实例 task_executor = TaskExecutor()