crawlo 1.1.1__py3-none-any.whl → 1.1.2__py3-none-any.whl

tests/test_scheduler.py

@@ -0,0 +1,242 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+测试修复后的 Scheduler 分布式队列功能
+"""
+import asyncio
+import sys
+from unittest.mock import Mock
+from crawlo.core.scheduler import Scheduler
+from crawlo.network.request import Request
+from crawlo.utils.log import get_logger
+
+
+class MockCrawler:
+    """模拟 Crawler 对象"""
+    def __init__(self, use_redis=True):
+        self.settings = MockSettings(use_redis)
+        self.stats = Mock()
+
+
+class MockSettings:
+    """模拟 Settings 对象"""
+    def __init__(self, use_redis=True):
+        self.use_redis = use_redis
+        
+    def get(self, key, default=None):
+        config = {
+            'FILTER_CLASS': 'crawlo.filters.memory_filter.MemoryFilter',
+            'LOG_LEVEL': 'INFO',
+            'DEPTH_PRIORITY': 1,
+            'SCHEDULER_MAX_QUEUE_SIZE': 100,
+            'SCHEDULER_QUEUE_NAME': 'test:crawlo:requests',
+            'FILTER_DEBUG': False,
+            'PROJECT_NAME': 'test',
+        }
+        if self.use_redis:
+            config['REDIS_URL'] = 'redis://localhost:6379/0'
+        
+        return config.get(key, default)
+    
+    def get_int(self, key, default=0):
+        value = self.get(key, default)
+        return int(value) if value is not None else default
+        
+    def get_bool(self, key, default=False):
+        value = self.get(key, default)
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            return value.lower() in ('true', '1', 'yes')
+        return bool(value) if value is not None else default
+
+
+class MockFilter:
+    """模拟去重过滤器"""
+    def __init__(self):
+        self.seen = set()
+        
+    @classmethod
+    def create_instance(cls, crawler):
+        return cls()
+    
+    async def requested(self, request):
+        if request.url in self.seen:
+            return True
+        self.seen.add(request.url)
+        return False
+    
+    def log_stats(self, request):
+        pass
+
+
+async def test_memory_scheduler():
+    """测试内存调度器"""
+    print("🔍 测试内存调度器...")
+    
+    crawler = MockCrawler(use_redis=False)
+    scheduler = Scheduler.create_instance(crawler)
+    
+    # 模拟去重过滤器
+    scheduler.dupe_filter = MockFilter()
+    
+    scheduler.open()
+    
+    # 测试入队
+    request1 = Request(url="https://example1.com")
+    request2 = Request(url="https://example2.com")
+    
+    success1 = await scheduler.enqueue_request(request1)
+    success2 = await scheduler.enqueue_request(request2)
+    
+    print(f"   📤 入队结果: {success1}, {success2}")
+    print(f"   📊 队列大小: {len(scheduler)}")
+    
+    # 测试出队
+    req1 = await scheduler.next_request()
+    req2 = await scheduler.next_request()
+    
+    print(f"   📥 出队结果: {req1.url if req1 else None}, {req2.url if req2 else None}")
+    print(f"   📊 剩余大小: {len(scheduler)}")
+    
+    await scheduler.close()
+    print("   ✅ 内存调度器测试完成")
+
+
+async def test_redis_scheduler():
+    """测试 Redis 调度器"""
+    print("🔍 测试 Redis 调度器...")
+    
+    try:
+        crawler = MockCrawler(use_redis=True)
+        scheduler = Scheduler.create_instance(crawler)
+        
+        # 模拟去重过滤器
+        scheduler.dupe_filter = MockFilter()
+        
+        scheduler.open()
+        
+        # 测试入队
+        request1 = Request(url="https://redis-test1.com", priority=5)
+        request2 = Request(url="https://redis-test2.com", priority=3)
+        request3 = Request(url="https://redis-test3.com", priority=8)
+        
+        success1 = await scheduler.enqueue_request(request1)
+        success2 = await scheduler.enqueue_request(request2)
+        success3 = await scheduler.enqueue_request(request3)
+        
+        print(f"   📤 入队结果: {success1}, {success2}, {success3}")
+        print(f"   📊 队列大小: {len(scheduler)}")
+        
+        # 等待一小段时间让 Redis 操作完成
+        await asyncio.sleep(0.5)
+        
+        # 测试出队（应该按优先级排序）
+        req1 = await scheduler.next_request()
+        req2 = await scheduler.next_request()
+        req3 = await scheduler.next_request()
+        
+        print("   📥 出队结果（按优先级）:")
+        if req1:
+            print(f"      {req1.url} (优先级: {getattr(req1, 'priority', 0)})")
+        if req2:
+            print(f"      {req2.url} (优先级: {getattr(req2, 'priority', 0)})")
+        if req3:
+            print(f"      {req3.url} (优先级: {getattr(req3, 'priority', 0)})")
+            
+        print(f"   📊 剩余大小: {len(scheduler)}")
+        
+        await scheduler.close()
+        print("   ✅ Redis 调度器测试完成")
+        
+    except Exception as e:
+        print(f"   ❌ Redis 调度器测试失败: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+async def test_concurrent_redis():
+    """测试并发 Redis 操作"""
+    print("🔍 测试并发 Redis 操作...")
+    
+    async def producer(scheduler, name, count):
+        """生产者"""
+        for i in range(count):
+            request = Request(url=f"https://{name}-{i}.com", priority=i % 10)
+            await scheduler.enqueue_request(request)
+            await asyncio.sleep(0.01)
+        print(f"   ✅ 生产者 {name} 完成 ({count} 个请求)")
+    
+    async def consumer(scheduler, name, count):
+        """消费者"""
+        consumed = 0
+        for _ in range(count):
+            request = await scheduler.next_request()
+            if request:
+                consumed += 1
+                await asyncio.sleep(0.005)
+            else:
+                break
+        print(f"   ✅ 消费者 {name} 处理了 {consumed} 个请求")
+    
+    try:
+        crawler = MockCrawler(use_redis=True)
+        scheduler = Scheduler.create_instance(crawler)
+        scheduler.dupe_filter = MockFilter()
+        scheduler.open()
+        
+        # 并发运行生产者和消费者
+        tasks = [
+            producer(scheduler, "producer-1", 5),
+            producer(scheduler, "producer-2", 5),
+            consumer(scheduler, "consumer-1", 3),
+            consumer(scheduler, "consumer-2", 3),
+            consumer(scheduler, "consumer-3", 4),
+        ]
+        
+        await asyncio.gather(*tasks, return_exceptions=True)
+        
+        print(f"   📊 最终队列大小: {len(scheduler)}")
+        
+        await scheduler.close()
+        print("   ✅ 并发测试完成")
+        
+    except Exception as e:
+        print(f"   ❌ 并发测试失败: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+async def main():
+    """主测试函数"""
+    print("🚀 开始测试修复后的 Scheduler...")
+    print("=" * 50)
+    
+    try:
+        # 1. 测试内存调度器
+        await test_memory_scheduler()
+        print()
+        
+        # 2. 测试 Redis 调度器
+        await test_redis_scheduler()
+        print()
+        
+        # 3. 测试并发操作
+        await test_concurrent_redis()
+        
+        print("=" * 50)
+        print("🎉 所有 Scheduler 测试完成！")
+        
+    except Exception as e:
+        print("=" * 50)
+        print(f"❌ 测试失败: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    # 设置日志级别避免过多输出
+    import logging
+    logging.getLogger('crawlo').setLevel(logging.WARNING)
+    
+    asyncio.run(main())

crawlo/pipelines/mysql_batch_pipline.py

@@ -1,273 +0,0 @@
-# -*- coding: utf-8 -*-
-import asyncio
-import aiomysql
-from typing import Optional, List, Dict
-from asyncmy import create_pool
-from crawlo.utils.log import get_logger
-from crawlo.exceptions import ItemDiscard
-from crawlo.utils.tools import make_insert_sql, logger
-
-
-class AsyncmyMySQLPipeline:
-    def __init__(self, crawler):
-        self.crawler = crawler
-        self.settings = crawler.settings
-        self.logger = get_logger(self.__class__.__name__, self.settings.get('LOG_LEVEL'))
-
-        # 配置参数
-        self.table_name = (
-            self.settings.get('MYSQL_TABLE') or
-            getattr(crawler.spider, 'mysql_table', None) or
-            f"{crawler.spider.name}_items"
-        )
-        self.batch_size = self.settings.getint('MYSQL_BATCH_SIZE', 100)
-        self.flush_interval = self.settings.getfloat('MYSQL_FLUSH_INTERVAL', 3.0)  # 秒
-
-        # 连接池相关
-        self._pool_lock = asyncio.Lock()
-        self._pool_initialized = False
-        self.pool = None
-
-        # 缓冲区与锁
-        self.items_buffer: List[Dict] = []
-        self.buffer_lock = asyncio.Lock()
-
-        # 后台任务
-        self.flush_task: Optional[asyncio.Task] = None
-
-        # 注册关闭事件
-        crawler.subscriber.subscribe(self.spider_closed, event='spider_closed')
-
-    @classmethod
-    def from_crawler(cls, crawler):
-        return cls(crawler)
-
-    async def _ensure_pool(self):
-        """确保连接池已初始化（线程安全）"""
-        if self._pool_initialized:
-            return
-
-        async with self._pool_lock:
-            if not self._pool_initialized:
-                try:
-                    self.pool = await create_pool(
-                        host=self.settings.get('MYSQL_HOST', 'localhost'),
-                        port=self.settings.get_int('MYSQL_PORT', 3306),
-                        user=self.settings.get('MYSQL_USER', 'root'),
-                        password=self.settings.get('MYSQL_PASSWORD', ''),
-                        db=self.settings.get('MYSQL_DB', 'scrapy_db'),
-                        minsize=self.settings.get_int('MYSQL_POOL_MIN', 3),
-                        maxsize=self.settings.get_int('MYSQL_POOL_MAX', 10),
-                        echo=self.settings.get_bool('MYSQL_ECHO', False)
-                    )
-                    self._pool_initialized = True
-                    self.logger.debug(f"MySQL连接池初始化完成（表: {self.table_name}）")
-                except Exception as e:
-                    self.logger.error(f"MySQL连接池初始化失败: {e}")
-                    raise
-
-    async def open_spider(self, spider):
-        """爬虫启动时初始化后台刷新任务"""
-        await self._ensure_pool()
-        self.flush_task = asyncio.create_task(self._flush_loop())
-
-    async def _flush_loop(self):
-        """后台循环：定期检查是否需要刷新缓冲区"""
-        while True:
-            await asyncio.sleep(self.flush_interval)
-            if len(self.items_buffer) > 0:
-                await self._flush_buffer()
-
-    async def _flush_buffer(self):
-        """将缓冲区中的数据批量写入数据库"""
-        async with self.buffer_lock:
-            if not self.items_buffer:
-                return
-
-            items_to_insert = self.items_buffer.copy()
-            self.items_buffer.clear()
-
-        try:
-            await self._ensure_pool()
-            first_item = items_to_insert[0]
-            sql = make_insert_sql(table=self.table_name, data=first_item, many=True)
-
-            values = [list(item.values()) for item in items_to_insert]
-
-            async with self.pool.acquire() as conn:
-                async with conn.cursor() as cursor:
-                    affected_rows = await cursor.executemany(sql, values)
-                    await conn.commit()
-
-            spider_name = getattr(self.crawler.spider, 'name', 'unknown')
-            self.logger.info(f"批量插入 {affected_rows} 条记录到 {self.table_name}")
-            self.crawler.stats.inc_value('mysql/insert_success_batch', len(items_to_insert))
-
-        except Exception as e:
-            self.logger.error(f"批量插入失败: {e}")
-            self.crawler.stats.inc_value('mysql/insert_failed_batch', len(items_to_insert))
-            # 可选：重试或丢弃
-            raise ItemDiscard(f"批量插入失败: {e}")
-
-    async def process_item(self, item, spider, kwargs=None) -> dict:
-        """将 item 添加到缓冲区，触发批量插入"""
-        item_dict = dict(item)
-
-        async with self.buffer_lock:
-            self.items_buffer.append(item_dict)
-            if len(self.items_buffer) >= self.batch_size:
-                # 达到批量阈值，立即刷新
-                await self._flush_buffer()
-
-        return item
-
-    async def spider_closed(self):
-        """关闭爬虫时，确保所有剩余数据被写入"""
-        if self.flush_task:
-            self.flush_task.cancel()
-            try:
-                await self.flush_task
-            except asyncio.CancelledError:
-                pass
-
-        # 刷最后一批数据
-        if self.items_buffer:
-            await self._flush_buffer()
-
-        # 关闭连接池
-        if self.pool:
-            self.pool.close()
-            await self.pool.wait_closed()
-            self.logger.info("MySQL连接池已关闭")
-
-
-class AiomysqlMySQLPipeline:
-    def __init__(self, crawler):
-        self.crawler = crawler
-        self.settings = crawler.settings
-        self.logger = get_logger(self.__class__.__name__, self.settings.get('LOG_LEVEL'))
-
-        # 配置
-        self.table_name = (
-            self.settings.get('MYSQL_TABLE') or
-            getattr(crawler.spider, 'mysql_table', None) or
-            f"{crawler.spider.name}_items"
-        )
-        self.batch_size = self.settings.getint('MYSQL_BATCH_SIZE', 100)
-        self.flush_interval = self.settings.getfloat('MYSQL_FLUSH_INTERVAL', 3.0)
-
-        # 连接池
-        self._pool_lock = asyncio.Lock()
-        self._pool_initialized = False
-        self.pool = None
-
-        # 缓冲
-        self.items_buffer: List[Dict] = []
-        self.buffer_lock = asyncio.Lock()
-
-        # 后台任务
-        self.flush_task: Optional[asyncio.Task] = None
-
-        crawler.subscriber.subscribe(self.spider_closed, event='spider_closed')
-
-    @classmethod
-    def create_instance(cls, crawler):
-        return cls(crawler)
-
-    async def _init_pool(self):
-        """延迟初始化连接池（线程安全）"""
-        if self._pool_initialized:
-            return
-
-        async with self._pool_lock:
-            if not self._pool_initialized:
-                try:
-                    self.pool = await aiomysql.create_pool(
-                        host=self.settings.get('MYSQL_HOST', 'localhost'),
-                        port=self.settings.getint('MYSQL_PORT', 3306),
-                        user=self.settings.get('MYSQL_USER', 'root'),
-                        password=self.settings.get('MYSQL_PASSWORD', ''),
-                        db=self.settings.get('MYSQL_DB', 'scrapy_db'),
-                        minsize=self.settings.getint('MYSQL_POOL_MIN', 3),
-                        maxsize=self.settings.getint('MYSQL_POOL_MAX', 10),
-                        cursorclass=aiomysql.DictCursor,
-                        autocommit=False
-                    )
-                    self._pool_initialized = True
-                    self.logger.debug(f"aiomysql连接池已初始化（表: {self.table_name}）")
-                except Exception as e:
-                    self.logger.error(f"aiomysql连接池初始化失败: {e}")
-                    raise
-
-    async def open_spider(self, spider):
-        """爬虫启动时创建后台刷新任务"""
-        await self._init_pool()
-        self.flush_task = asyncio.create_task(self._flush_loop())
-
-    async def _flush_loop(self):
-        """定期刷新缓冲区"""
-        while True:
-            await asyncio.sleep(self.flush_interval)
-            if len(self.items_buffer) > 0:
-                await self._flush_buffer()
-
-    async def _flush_buffer(self):
-        """执行批量插入"""
-        async with self.buffer_lock:
-            if not self.items_buffer:
-                return
-            items_to_insert = self.items_buffer.copy()
-            self.items_buffer.clear()
-
-        try:
-            await self._init_pool()
-            keys = items_to_insert[0].keys()
-            placeholders = ', '.join(['%s'] * len(keys))
-            columns = ', '.join([f'`{k}`' for k in keys])
-            sql = f"INSERT INTO `{self.table_name}` ({columns}) VALUES ({placeholders})"
-
-            values = [list(item.values()) for item in items_to_insert]
-
-            async with self.pool.acquire() as conn:
-                async with conn.cursor() as cursor:
-                    result = await cursor.executemany(sql, values)
-                    await conn.commit()
-
-            spider_name = getattr(self.crawler.spider, 'name', 'unknown')
-            self.logger.info(f"【{spider_name}】批量插入 {result} 条记录到 {self.table_name}")
-            self.crawler.stats.inc_value('mysql/insert_success_batch', len(items_to_insert))
-
-        except aiomysql.Error as e:
-            self.logger.error(f"aiomysql批量插入失败: {e}")
-            self.crawler.stats.inc_value('mysql/insert_failed_batch', len(items_to_insert))
-            raise ItemDiscard(f"MySQL错误: {e.args[1]}")
-        except Exception as e:
-            self.logger.error(f"未知错误: {e}")
-            raise ItemDiscard(f"处理失败: {e}")
-
-    async def process_item(self, item, spider) -> dict:
-        item_dict = dict(item)
-
-        async with self.buffer_lock:
-            self.items_buffer.append(item_dict)
-            if len(self.items_buffer) >= self.batch_size:
-                await self._flush_buffer()
-
-        return item
-
-    async def spider_closed(self):
-        """清理资源并提交剩余数据"""
-        if self.flush_task:
-            self.flush_task.cancel()
-            try:
-                await self.flush_task
-            except asyncio.CancelledError:
-                pass
-
-        if self.items_buffer:
-            await self._flush_buffer()
-
-        if self.pool:
-            self.pool.close()
-            await self.pool.wait_closed()
-            self.logger.info("aiomysql连接池已释放")

crawlo/utils/pqueue.py

@@ -1,174 +0,0 @@
-# -*- coding:UTF-8 -*-
-import sys
-import asyncio
-import warnings
-from urllib.parse import urlparse
-from asyncio import PriorityQueue
-from redis.asyncio import from_url
-from typing import Any, Optional, Dict, Annotated
-from pydantic import (
-    BaseModel,
-    Field,
-    model_validator
-)
-
-from crawlo import Request
-from crawlo.settings.default_settings import REDIS_URL
-
-
-class SpiderPriorityQueue(PriorityQueue):
-    """带超时功能的异步优先级队列"""
-
-    def __init__(self, maxsize: int = 0) -> None:
-        """初始化队列，maxsize为0表示无大小限制"""
-        super().__init__(maxsize)
-
-    async def get(self, timeout: float = 0.1) -> Optional[Request]:
-        """
-        异步获取队列元素，带超时功能
-
-        Args:
-            timeout: 超时时间（秒），默认0.1秒
-
-        Returns:
-            队列元素(优先级, 值)或None(超时)
-        """
-        try:
-            # 根据Python版本选择超时实现方式
-            if sys.version_info >= (3, 11):
-                async with asyncio.timeout(timeout):
-                    return await super().get()
-            else:
-                return await asyncio.wait_for(super().get(), timeout=timeout)
-        except asyncio.TimeoutError:
-            return None
-
-
-class TaskModel(BaseModel):
-    """爬虫任务数据模型 (完全兼容Pydantic V2)"""
-    url: Annotated[str, Field(min_length=1, max_length=2000, examples=["https://example.com"])]
-    meta: Dict[str, Any] = Field(default_factory=dict)
-    priority: Annotated[int, Field(default=0, ge=0, le=10, description="0=最高优先级")]
-
-    @classmethod
-    def validate_url(cls, v: str) -> str:
-        """验证URL格式"""
-        if not v.startswith(('http://', 'https://')):
-            raise ValueError('URL必须以 http:// 或 https:// 开头')
-
-        parsed = urlparse(v)
-        if not parsed.netloc:
-            raise ValueError('URL缺少有效域名')
-
-        return v.strip()
-
-    @model_validator(mode='after')
-    def validate_priority_logic(self) -> 'TaskModel':
-        """跨字段验证示例"""
-        if 'admin' in self.url and self.priority > 5:
-            self.priority = 5  # 自动调整管理页面的优先级
-        return self
-
-
-class DistributedPriorityQueue:
-    def __init__(
-            self,
-            redis_url: str,
-            queue_name: str = "spider_queue",
-            max_connections: int = 10,
-            health_check_interval: int = 30
-    ):
-        """
-        Args:
-            redis_url: redis://[:password]@host:port[/db]
-            queue_name: Redis有序集合键名
-            max_connections: 连接池大小
-            health_check_interval: 连接健康检查间隔(秒)
-        """
-        self.redis = from_url(
-            redis_url,
-            max_connections=max_connections,
-            health_check_interval=health_check_interval,
-            socket_keepalive=True,
-            decode_responses=True
-        )
-        self.queue_name = queue_name
-
-    async def put(self, task: TaskModel) -> bool:
-        """
-        添加任务到队列（使用Pydantic V2的model_dump_json）
-
-        Args:
-            task: 已验证的TaskModel实例
-
-        Returns:
-            bool: 是否成功添加 (Redis的ZADD返回添加数量)
-        """
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore", category=DeprecationWarning)
-            task_str = task.model_dump_json()  # 正确使用V2的序列化方法
-        return await self.redis.zadd(
-            self.queue_name,
-            {task_str: task.priority}
-        ) > 0
-
-    async def get(self, timeout: float = 1.0) -> Optional[TaskModel]:
-        """
-        获取优先级最高的任务（自动验证）
-
-        Args:
-            timeout: 阻塞超时时间(秒)
-
-        Returns:
-            TaskModel实例或None(超时/队列空)
-        """
-        try:
-            result = await self.redis.bzpopmax(
-                self.queue_name,
-                timeout=timeout
-            )
-            if result:
-                _, task_str, _ = result
-                with warnings.catch_warnings():
-                    warnings.simplefilter("ignore", category=DeprecationWarning)
-                    return TaskModel.model_validate_json(task_str)  # 正确使用V2的反序列化方法
-        except Exception as e:
-            print(f"任务获取失败: {type(e).__name__}: {e}")
-        return None
-
-    async def aclose(self):
-        """安全关闭连接"""
-        await self.redis.aclose()
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.aclose()
-
-
-# 使用示例
-async def demo():
-    async with DistributedPriorityQueue(
-            REDIS_URL,
-            max_connections=20,
-            health_check_interval=10
-    ) as queue:
-        # 添加任务（自动触发验证）
-        task = TaskModel(
-            url="https://example.com/1",
-            priority=1,
-            meta={"depth": 2}
-        )
-
-        if await queue.put(task):
-            print(f"任务添加成功: {task.url}")
-
-        # 获取任务
-        if result := await queue.get(timeout=2.0):
-            print(f"获取任务: {result.url} (优先级={result.priority})")
-            print(f"元数据: {result.meta}")
-
-
-if __name__ == "__main__":
-    asyncio.run(demo())