jettask 0.2.18__py3-none-any.whl → 0.2.20__py3-none-any.whl

jettask/{backend/data_access.py → persistence/base.py}

@@ -10,8 +10,13 @@ from datetime import datetime, timedelta, timezone
 from typing import Dict, List, Optional, Tuple
 import redis.asyncio as redis
 from sqlalchemy import text, bindparam
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.asyncio import AsyncSession
+
+# 导入统一的数据库连接工具
+from ..utils.db_connector import (
+    get_dual_mode_async_redis_client,
+    get_pg_engine_and_factory
+)
 
 # 设置日志
 logger = logging.getLogger(__name__)
@@ -62,67 +67,60 @@ class PostgreSQLConfig:
 
 
 class JetTaskDataAccess:
-    """JetTask数据访问类"""
-    
+    """JetTask数据访问类（使用统一的数据库连接工具）"""
+
     def __init__(self):
         self.redis_config = RedisConfig.from_env()
         self.pg_config = PostgreSQLConfig.from_env()
         # Redis前缀可以从环境变量配置，默认为 "jettask"
         self.redis_prefix = os.environ.get('JETTASK_REDIS_PREFIX', 'jettask')
+
+        # 使用全局单例连接池
+        self._text_redis_client: Optional[redis.Redis] = None
+        self._binary_redis_client: Optional[redis.Redis] = None
+
+        # PostgreSQL 相关
         self.async_engine = None
         self.AsyncSessionLocal = None
-        self._redis_pool = None
-        self._binary_redis_pool = None  # 用于处理包含二进制数据的Stream
         
     async def initialize(self):
-        """初始化数据库连接"""
+        """初始化数据库连接（使用全局单例）"""
         try:
-            # 初始化PostgreSQL引擎
-            dsn = self.pg_config.dsn
-            if dsn.startswith('postgresql://'):
-                dsn = dsn.replace('postgresql://', 'postgresql+psycopg://', 1)
-            
-            self.async_engine = create_async_engine(
-                dsn,
+            # 构建 Redis URL
+            redis_url = f"redis://"
+            if self.redis_config.password:
+                redis_url += f":{self.redis_config.password}@"
+            redis_url += f"{self.redis_config.host}:{self.redis_config.port}/{self.redis_config.db}"
+
+            # 构建 PostgreSQL 配置
+            pg_config = {
+                'host': self.pg_config.host,
+                'port': self.pg_config.port,
+                'user': self.pg_config.user,
+                'password': self.pg_config.password,
+                'database': self.pg_config.database
+            }
+
+            # 初始化 PostgreSQL 连接（使用全局单例）
+            self.async_engine, self.AsyncSessionLocal = get_pg_engine_and_factory(
+                config=pg_config,
                 pool_size=10,
                 max_overflow=5,
                 pool_pre_ping=True,
                 echo=False
             )
-            
-            self.AsyncSessionLocal = sessionmaker(
-                bind=self.async_engine,
-                class_=AsyncSession,
-                expire_on_commit=False
-            )
-            
-            # 初始化Redis连接池（用于普通操作）
-            self._redis_pool = redis.ConnectionPool(
-                host=self.redis_config.host,
-                port=self.redis_config.port,
-                db=self.redis_config.db,
-                password=self.redis_config.password,
-                encoding='utf-8',
-                decode_responses=True,
-                socket_keepalive=True,
-                socket_connect_timeout=5,
-                retry_on_timeout=True
-            )
-            
-            # 初始化二进制Redis连接池（用于Stream操作）
-            self._binary_redis_pool = redis.ConnectionPool(
-                host=self.redis_config.host,
-                port=self.redis_config.port,
-                db=self.redis_config.db,
-                password=self.redis_config.password,
-                decode_responses=False,  # 不解码，因为Stream包含msgpack二进制数据
+
+            # 初始化 Redis 连接（使用全局单例，双模式）
+            self._text_redis_client, self._binary_redis_client = get_dual_mode_async_redis_client(
+                redis_url=redis_url,
+                max_connections=50,
                 socket_keepalive=True,
                 socket_connect_timeout=5,
                 retry_on_timeout=True
             )
-            
+
             logger.info("数据库连接初始化成功")
-            
+
         except Exception as e:
             logger.error(f"数据库连接初始化失败: {e}")
             raise
@@ -132,21 +130,24 @@ class JetTaskDataAccess:
         return self.AsyncSessionLocal()
     
     async def close(self):
-        """关闭数据库连接"""
-        if self.async_engine:
-            await self.async_engine.dispose()
-        if self._redis_pool:
-            await self._redis_pool.disconnect()
-        if self._binary_redis_pool:
-            await self._binary_redis_pool.disconnect()
-    
+        """关闭数据库连接（由于使用全局单例，这里只重置状态）"""
+        # 注意：连接池由全局单例管理，这里只清理引用
+        self._text_redis_client = None
+        self._binary_redis_client = None
+        self.async_engine = None
+        self.AsyncSessionLocal = None
+
     async def get_redis_client(self):
-        """获取Redis客户端"""
-        return redis.Redis(connection_pool=self._redis_pool)
-    
+        """获取 Redis 客户端（使用全局单例）"""
+        if not self._text_redis_client:
+            raise RuntimeError("Redis client not initialized")
+        return self._text_redis_client
+
     async def get_binary_redis_client(self):
-        """获取二进制Redis客户端（用于Stream操作）"""
-        return redis.Redis(connection_pool=self._binary_redis_pool)
+        """获取二进制 Redis 客户端（用于Stream操作，使用全局单例）"""
+        if not self._binary_redis_client:
+            raise RuntimeError("Binary Redis client not initialized")
+        return self._binary_redis_client
     
     async def fetch_queues_data(self) -> List[Dict]:
         """获取队列数据（基于Redis Stream）"""
@@ -2046,7 +2047,7 @@ class JetTaskDataAccess:
     async def _sync_task_to_redis(self, task_id: str, enabled: bool):
         """同步任务状态到 Redis"""
         try:
-            if not self._redis_pool:
+            if not self._redis_connector:
                 logger.debug("Redis not configured, skipping sync")
                 return
                 

jettask/persistence/consumer.py

@@ -0,0 +1,315 @@
+"""PostgreSQL消费者主模块
+
+协调各个子模块，提供统一的消费者接口。
+"""
+
+import asyncio
+import logging
+import os
+import socket
+import traceback
+from typing import Optional
+
+import redis.asyncio as redis
+from redis.asyncio import Redis
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import text
+
+from jettask.webui.config import PostgreSQLConfig, RedisConfig
+from jettask.worker.manager import ConsumerManager
+
+from .backlog_monitor import BacklogMonitor
+from .task_updater import TaskUpdater
+from .offline_recovery import OfflineRecoveryHandler
+from .task_persistence import TaskPersistence
+from .queue_discovery import QueueDiscovery
+from .message_consumer import MessageConsumer
+from .maintenance import DatabaseMaintenance
+
+logger = logging.getLogger(__name__)
+
+
+class PostgreSQLConsumer:
+    """PostgreSQL消费者，从Redis队列消费任务并持久化到PostgreSQL
+
+    支持多租户（命名空间）隔离
+
+    架构说明：
+    - BacklogMonitor: 监控Stream积压情况
+    - TaskUpdater: 更新任务状态（从TASK_CHANGES流）
+    - OfflineRecoveryHandler: 恢复离线worker的消息
+    - TaskPersistence: 解析并持久化任务数据
+    - QueueDiscovery: 发现和管理队列
+    - MessageConsumer: 消费队列消息
+    - DatabaseMaintenance: 数据库维护任务
+    """
+
+    def __init__(
+        self,
+        pg_config: PostgreSQLConfig,
+        redis_config: RedisConfig,
+        prefix: str = "jettask",
+        node_id: str = None,
+        # consumer_strategy 参数已移除，现在只使用 HEARTBEAT 策略
+        namespace_id: str = None,
+        namespace_name: str = None,
+        enable_backlog_monitor: bool = True,
+        backlog_monitor_interval: int = 1
+    ):
+        """初始化PostgreSQL消费者
+
+        Args:
+            pg_config: PostgreSQL配置
+            redis_config: Redis配置
+            prefix: Redis键前缀
+            node_id: 节点ID
+            consumer_strategy: 消费者策略
+            namespace_id: 命名空间ID
+            namespace_name: 命名空间名称
+            enable_backlog_monitor: 是否启用积压监控
+            backlog_monitor_interval: 积压监控间隔（秒）
+        """
+        self.pg_config = pg_config
+        self.redis_config = redis_config
+        self.prefix = prefix
+
+        # 命名空间支持
+        self.namespace_id = namespace_id
+        self.namespace_name = namespace_name or "default"
+
+        # 节点标识
+        hostname = socket.gethostname()
+        self.node_id = node_id or f"{hostname}_{os.getpid()}"
+
+        # 消费者配置
+        # consumer_strategy 已移除，现在只使用 HEARTBEAT 策略
+        self.consumer_group = f"{prefix}_pg_consumer"
+
+        # Redis和数据库连接（将在start时初始化）
+        self.redis_client: Optional[Redis] = None
+        self.async_engine = None
+        self.AsyncSessionLocal = None
+
+        # ConsumerManager（将在start时初始化）
+        self.consumer_manager = None
+        self.consumer_id = None
+
+        # 各个子模块（将在start时初始化）
+        self.backlog_monitor: Optional[BacklogMonitor] = None
+        self.task_updater: Optional[TaskUpdater] = None
+        self.offline_recovery: Optional[OfflineRecoveryHandler] = None
+        self.task_persistence: Optional[TaskPersistence] = None
+        self.queue_discovery: Optional[QueueDiscovery] = None
+        self.message_consumer: Optional[MessageConsumer] = None
+        self.database_maintenance: Optional[DatabaseMaintenance] = None
+
+        # 积压监控配置
+        self.enable_backlog_monitor = enable_backlog_monitor
+        self.backlog_monitor_interval = backlog_monitor_interval
+
+        self._running = False
+
+    async def start(self):
+        """启动消费者"""
+        logger.info(f"Starting PostgreSQL consumer (modular) on node: {self.node_id}")
+
+        # 1. 连接Redis（使用全局客户端实例）
+        from jettask.utils.db_connector import get_async_redis_client, get_sync_redis_client
+
+        # 构建 Redis URL
+        redis_url = f"redis://"
+        if self.redis_config.password:
+            redis_url += f":{self.redis_config.password}@"
+        redis_url += f"{self.redis_config.host}:{self.redis_config.port}/{self.redis_config.db}"
+
+        self.redis_client = get_async_redis_client(
+            redis_url=redis_url,
+            decode_responses=False  # 保持二进制模式
+        )
+
+        # 2. 初始化ConsumerManager（需要同步的Redis客户端）
+        sync_redis_client = get_sync_redis_client(
+            redis_url=redis_url,
+            decode_responses=True
+        )
+
+        # 配置ConsumerManager
+        initial_queues = ['TASK_CHANGES']  # TASK_CHANGES是固定的
+        consumer_config = {
+            'redis_prefix': self.prefix,
+            'queues': initial_queues,
+            'worker_prefix': 'PG_CONSUMER',  # 使用不同的前缀，与task worker区分开
+        }
+
+        self.consumer_manager = ConsumerManager(
+            redis_client=sync_redis_client,
+            # strategy 参数已移除，现在只使用 HEARTBEAT 策略
+            config=consumer_config
+        )
+
+        # 获取稳定的consumer_id
+        self.consumer_id = self.consumer_manager.get_consumer_name('TASK_CHANGES')
+        logger.debug(f"Using consumer_id: {self.consumer_id} with strategy: HEARTBEAT")
+
+        # 3. 创建SQLAlchemy异步引擎
+        if self.pg_config.dsn.startswith('postgresql://'):
+            dsn = self.pg_config.dsn.replace('postgresql://', 'postgresql+asyncpg://', 1)
+        else:
+            dsn = self.pg_config.dsn
+
+        self.async_engine = create_async_engine(
+            dsn,
+            pool_size=50,
+            max_overflow=20,
+            pool_pre_ping=True,
+            pool_recycle=300,
+            echo=False
+        )
+
+        # 预热连接池
+        logger.debug("Pre-warming database connection pool...")
+        async with self.async_engine.begin() as conn:
+            await conn.execute(text("SELECT 1"))
+
+        # 创建异步会话工厂
+        self.AsyncSessionLocal = sessionmaker(
+            self.async_engine,
+            class_=AsyncSession,
+            expire_on_commit=False
+        )
+
+        # 4. 初始化各个子模块
+        # 任务持久化处理器
+        self.task_persistence = TaskPersistence(
+            async_session_local=self.AsyncSessionLocal,
+            namespace_id=self.namespace_id,
+            namespace_name=self.namespace_name
+        )
+
+        # 队列发现器
+        self.queue_discovery = QueueDiscovery(
+            redis_client=self.redis_client,
+            redis_prefix=self.prefix,
+            consumer_group=self.consumer_group,
+            consumer_manager=self.consumer_manager
+        )
+
+        # 先进行一次队列发现，确保ConsumerManager有正确的队列列表
+        await self.queue_discovery.initial_queue_discovery()
+
+        # 消息消费器
+        self.message_consumer = MessageConsumer(
+            redis_client=self.redis_client,
+            redis_prefix=self.prefix,
+            consumer_group=self.consumer_group,
+            consumer_id=self.consumer_id,
+            task_persistence=self.task_persistence,
+            queue_discovery=self.queue_discovery
+        )
+
+        # 任务状态更新器
+        self.task_updater = TaskUpdater(
+            redis_client=self.redis_client,
+            async_session_local=self.AsyncSessionLocal,
+            redis_prefix=self.prefix,
+            consumer_id=self.consumer_id
+        )
+
+        # 离线恢复处理器
+        self.offline_recovery = OfflineRecoveryHandler(
+            redis_client=self.redis_client,
+            redis_prefix=self.prefix,
+            consumer_id=self.consumer_id,
+            task_updater=self.task_updater
+        )
+        # 延迟初始化（需要consumer_manager）
+        self.offline_recovery.set_consumer_manager(self.consumer_manager)
+
+        # 数据库维护
+        self.database_maintenance = DatabaseMaintenance(
+            async_session_local=self.AsyncSessionLocal
+        )
+
+        # 积压监控器
+        self.backlog_monitor = BacklogMonitor(
+            redis_client=self.redis_client,
+            async_session_local=self.AsyncSessionLocal,
+            redis_prefix=self.prefix,
+            namespace_name=self.namespace_name,
+            node_id=self.node_id,
+            enable_monitor=self.enable_backlog_monitor,
+            monitor_interval=self.backlog_monitor_interval
+        )
+
+        # 5. 启动所有子模块
+        self._running = True
+
+        # 启动队列发现
+        await self.queue_discovery.start_discovery()
+
+        # 启动消息消费
+        await self.message_consumer.start()
+
+        # 启动任务状态更新
+        await self.task_updater.start()
+
+        # 启动离线恢复
+        await self.offline_recovery.start()
+
+        # 启动数据库维护
+        await self.database_maintenance.start()
+
+        # 启动积压监控
+        if self.enable_backlog_monitor:
+            await self.backlog_monitor.start()
+            logger.info(f"Stream backlog monitor enabled with {self.backlog_monitor_interval}s interval")
+
+        # 如果使用HEARTBEAT策略，ConsumerManager会自动管理心跳
+        if self.consumer_manager:
+            logger.debug("Heartbeat is managed by ConsumerManager")
+
+        logger.debug("PostgreSQL consumer started successfully")
+
+    async def stop(self):
+        """停止消费者"""
+        logger.debug("Stopping PostgreSQL consumer...")
+        self._running = False
+
+        # 停止所有子模块
+        if self.backlog_monitor:
+            await self.backlog_monitor.stop()
+
+        if self.database_maintenance:
+            await self.database_maintenance.stop()
+
+        if self.offline_recovery:
+            await self.offline_recovery.stop()
+
+        if self.task_updater:
+            await self.task_updater.stop()
+
+        if self.message_consumer:
+            await self.message_consumer.stop()
+
+        if self.queue_discovery:
+            await self.queue_discovery.stop_discovery()
+
+        # 清理ConsumerManager
+        if self.consumer_manager:
+            try:
+                self.consumer_manager.cleanup()
+                logger.debug(f"Cleaned up ConsumerManager for consumer: {self.consumer_id}")
+            except Exception as e:
+                logger.error(f"Error cleaning up ConsumerManager: {e}")
+                logger.error(traceback.format_exc())
+
+        # 关闭连接
+        if self.redis_client:
+            await self.redis_client.close()
+
+        if self.async_engine:
+            await self.async_engine.dispose()
+
+        logger.debug("PostgreSQL consumer stopped")

jettask/{core → persistence}/db_manager.py

@@ -55,8 +55,8 @@ class ConnectionPool:
         self.config = config
         
         # Redis连接池
-        self._redis_pool: Optional[redis.ConnectionPool] = None
-        self._binary_redis_pool: Optional[redis.ConnectionPool] = None
+        self._redis_client: Optional[redis.Redis] = None
+        self._binary_redis_client: Optional[redis.Redis] = None
         
         # PostgreSQL连接池
         self._pg_pool: Optional[asyncpg.Pool] = None
@@ -92,25 +92,26 @@ class ConnectionPool:
                 raise
     
     async def _init_redis(self):
-        """初始化Redis连接池"""
+        """初始化Redis客户端（使用全局客户端实例）"""
+        from jettask.utils.db_connector import get_async_redis_client
+
         url = self.config.redis_url
-        
-        # 文本连接池（decode_responses=True）
-        self._redis_pool = redis.ConnectionPool.from_url(
-            url,
+
+        # 文本客户端（decode_responses=True）
+        self._redis_client = get_async_redis_client(
+            redis_url=url,
             max_connections=20,
-            decode_responses=True,
-            encoding='utf-8'
+            decode_responses=True
         )
-        
-        # 二进制连接池（decode_responses=False）
-        self._binary_redis_pool = redis.ConnectionPool.from_url(
-            url,
+
+        # 二进制客户端（decode_responses=False）
+        self._binary_redis_client = get_async_redis_client(
+            redis_url=url,
             max_connections=20,
             decode_responses=False
         )
-        
-        logger.debug(f"Redis连接池已创建: {self.namespace}")
+
+        logger.debug(f"Redis客户端已就绪: {self.namespace}")
     
     async def _init_postgres(self):
         """初始化PostgreSQL连接池"""
@@ -160,8 +161,8 @@ class ConnectionPool:
         if not self.config.has_redis():
             raise ValueError(f"命名空间 '{self.namespace}' 未配置Redis")
         
-        pool = self._redis_pool if decode else self._binary_redis_pool
-        return redis.Redis(connection_pool=pool)
+        client = self._redis_client if decode else self._binary_redis_client
+        return client
     
     @asynccontextmanager
     async def get_pg_connection(self) -> AsyncGenerator[asyncpg.Connection, None]:
@@ -197,11 +198,11 @@ class ConnectionPool:
     async def close(self):
         """关闭所有连接"""
         try:
-            if self._redis_pool:
-                await self._redis_pool.aclose()
-            
-            if self._binary_redis_pool:
-                await self._binary_redis_pool.aclose()
+            if self._redis_client:
+                await self._redis_client.aclose()
+
+            if self._binary_redis_client:
+                await self._binary_redis_client.aclose()
             
             if self._pg_pool:
                 await self._pg_pool.close()

jettask/persistence/maintenance.py

@@ -0,0 +1,81 @@
+"""数据库维护模块
+
+负责定期执行数据库维护任务，如ANALYZE等。
+"""
+
+import asyncio
+import logging
+import time
+import traceback
+
+from sqlalchemy import text
+from sqlalchemy.orm import sessionmaker
+
+logger = logging.getLogger(__name__)
+
+
+class DatabaseMaintenance:
+    """数据库维护处理器
+
+    职责：
+    - 定期执行ANALYZE命令更新统计信息
+    - 优化查询性能
+    """
+
+    def __init__(self, async_session_local: sessionmaker):
+        """初始化数据库维护处理器
+
+        Args:
+            async_session_local: SQLAlchemy会话工厂
+        """
+        self.AsyncSessionLocal = async_session_local
+
+        # 维护配置
+        self.analyze_interval = 7200  # 每2小时执行一次ANALYZE
+        self.last_analyze_time = 0
+
+        self._running = False
+        self._maintenance_task = None
+
+    async def start(self):
+        """启动维护任务"""
+        self._running = True
+        self._maintenance_task = asyncio.create_task(self._maintenance_loop())
+        logger.debug("DatabaseMaintenance started")
+
+    async def stop(self):
+        """停止维护任务"""
+        self._running = False
+
+        if self._maintenance_task:
+            self._maintenance_task.cancel()
+            try:
+                await self._maintenance_task
+            except asyncio.CancelledError:
+                pass
+
+        logger.debug("DatabaseMaintenance stopped")
+
+    async def _maintenance_loop(self):
+        """维护循环"""
+        while self._running:
+            try:
+                current_time = time.time()
+
+                # 执行ANALYZE
+                if current_time - self.last_analyze_time > self.analyze_interval:
+                    async with self.AsyncSessionLocal() as session:
+                        logger.debug("Running ANALYZE on tasks and task_runs tables...")
+                        await session.execute(text("ANALYZE tasks"))
+                        await session.execute(text("ANALYZE task_runs"))
+                        await session.commit()
+                        logger.debug("ANALYZE completed successfully for both tables")
+                        self.last_analyze_time = current_time
+
+                # 每5分钟检查一次
+                await asyncio.sleep(300)
+
+            except Exception as e:
+                logger.error(f"Error in database maintenance: {e}")
+                logger.error(traceback.format_exc())
+                await asyncio.sleep(60)