Support keword

core/task_executor.py

@@ -0,0 +1,235 @@
+"""
+爬取任务执行器
+负责执行爬取任务并更新任务状态
+"""
+
+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()
+        
+        # 准备参数
+        websites = task.websites if task.websites else 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()
+        
+        # 准备参数
+        websites = task.websites if task.websites else 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()
+        
+        # 准备参数
+        websites = task.websites if task.websites else 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()