jettask 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

jettask/pg_consumer/pg_consumer_v2.py

@@ -0,0 +1,403 @@
+"""
+PostgreSQL Consumer V2 - 支持多消费者组的双表结构
+"""
+import logging
+import json
+from typing import Optional, Dict, Any, List
+from datetime import datetime
+import asyncio
+import psycopg2
+from psycopg2.extras import RealDictCursor, Json
+import asyncpg
+
+logger = logging.getLogger(__name__)
+
+
+class PGConsumerV2:
+    """PostgreSQL消费者V2 - 支持多消费者组"""
+    
+    def __init__(self, pg_url: str, redis_prefix: str = 'jettask'):
+        self.pg_url = pg_url
+        self.redis_prefix = redis_prefix
+        self.conn = None
+        self.async_conn = None
+        
+    def connect(self):
+        """建立同步数据库连接"""
+        if not self.conn or self.conn.closed:
+            self.conn = psycopg2.connect(self.pg_url)
+            self.conn.autocommit = False
+            
+    async def async_connect(self):
+        """建立异步数据库连接"""
+        if not self.async_conn:
+            self.async_conn = await asyncpg.connect(self.pg_url)
+            
+    def close(self):
+        """关闭同步连接"""
+        if self.conn:
+            self.conn.close()
+            self.conn = None
+            
+    async def async_close(self):
+        """关闭异步连接"""
+        if self.async_conn:
+            await self.async_conn.close()
+            self.async_conn = None
+            
+    def ensure_tables(self):
+        """确保数据表存在"""
+        from .sql_utils import execute_sql_file
+        self.connect()
+        try:
+            with self.conn.cursor() as cursor:
+                # 检查新表是否存在
+                cursor.execute("""
+                    SELECT EXISTS (
+                        SELECT 1 FROM information_schema.tables 
+                        WHERE table_name = 'tasks'
+                    )
+                """)
+                if not cursor.fetchone()[0]:
+                    # 使用新的SQL执行函数
+                    sql_path = '/home/yuyang/easy-task/jettask/pg_consumer/sql/create_new_tables.sql'
+                    execute_sql_file(self.conn, sql_path)
+                    logger.info("Created new table structure for multi-consumer group support")
+        except Exception as e:
+            self.conn.rollback()
+            logger.error(f"Error ensuring tables: {e}")
+            raise
+            
+    async def async_ensure_tables(self):
+        """异步确保数据表存在"""
+        from .sql_utils import split_sql_statements
+        await self.async_connect()
+        try:
+            # 检查新表是否存在，并且有正确的列
+            exists = await self.async_conn.fetchval("""
+                SELECT EXISTS (
+                    SELECT 1 FROM information_schema.columns 
+                    WHERE table_name = 'tasks' AND column_name = 'stream_id'
+                )
+            """)
+            if not exists:
+                logger.info("Creating new table structure...")
+                # 执行创建表的SQL
+                with open('/home/yuyang/easy-task/jettask/pg_consumer/sql/create_new_tables.sql', 'r') as f:
+                    sql_content = f.read()
+                    # 使用智能分割函数
+                    statements = split_sql_statements(sql_content)
+                    for i, stmt in enumerate(statements, 1):
+                        try:
+                            logger.debug(f"执行第 {i}/{len(statements)} 个SQL语句")
+                            await self.async_conn.execute(stmt)
+                        except Exception as e:
+                            # 忽略已存在的对象错误
+                            if 'already exists' not in str(e):
+                                logger.warning(f"Error executing statement {i}: {e}")
+                                logger.debug(f"Statement: {stmt[:100]}...")
+                logger.info("Created new table structure for multi-consumer group support")
+        except Exception as e:
+            logger.error(f"Error ensuring tables: {e}")
+            raise
+    
+    def record_task(self, event_id: str, event_data: Dict[str, Any], 
+                    queue: str, consumer_group: str) -> int:
+        """
+        记录任务信息（同步方法）
+        返回task_id
+        """
+        self.connect()
+        try:
+            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
+                # 1. 插入或获取任务基础信息
+                task_name = event_data.get('_task_name', event_data.get('name', 'unknown'))
+                
+                # 尝试插入任务，如果已存在则忽略
+                cursor.execute("""
+                    INSERT INTO tasks (
+                        stream_id, queue, task_name, task_type, 
+                        payload, priority, source, metadata
+                    ) VALUES (
+                        %s, %s, %s, %s, %s, %s, %s, %s
+                    )
+                    ON CONFLICT (stream_id) DO UPDATE
+                    SET updated_at = CURRENT_TIMESTAMP
+                    RETURNING id
+                """, (
+                    event_id,
+                    queue,
+                    task_name,
+                    event_data.get('event_type', 'task'),
+                    Json(event_data),
+                    event_data.get('priority', 0),
+                    event_data.get('source', 'stream'),
+                    Json({'consumer_group': consumer_group})
+                ))
+                
+                task_id = cursor.fetchone()['id']
+                
+                # 2. 插入任务运行记录
+                consumer_name = event_data.get('consumer', '')
+                worker_id = event_data.get('worker_id', consumer_name)
+                
+                cursor.execute("""
+                    INSERT INTO task_runs (
+                        task_id, stream_id, consumer_group, 
+                        consumer_name, worker_id, status
+                    ) VALUES (
+                        %s, %s, %s, %s, %s, %s
+                    )
+                    ON CONFLICT (task_id, consumer_group) DO UPDATE
+                    SET 
+                        status = EXCLUDED.status,
+                        updated_at = CURRENT_TIMESTAMP
+                    RETURNING id
+                """, (
+                    task_id,
+                    event_id,
+                    consumer_group,
+                    consumer_name,
+                    worker_id,
+                    'pending'
+                ))
+                
+                self.conn.commit()
+                return task_id
+                
+        except Exception as e:
+            self.conn.rollback()
+            logger.error(f"Error recording task: {e}")
+            raise
+            
+    async def async_record_task(self, event_id: str, event_data: Dict[str, Any],
+                                queue: str, consumer_group: str) -> int:
+        """
+        异步记录任务信息
+        返回task_id
+        """
+        await self.async_connect()
+        try:
+            task_name = event_data.get('_task_name', event_data.get('name', 'unknown'))
+            
+            # 使用事务
+            async with self.async_conn.transaction():
+                # 1. 插入或获取任务基础信息
+                task_id = await self.async_conn.fetchval("""
+                    INSERT INTO tasks (
+                        stream_id, queue, namespace, scheduled_task_id,
+                        payload, priority, source, metadata
+                    ) VALUES (
+                        $1, $2, $3, $4, $5, $6, $7, $8
+                    )
+                    ON CONFLICT (stream_id) DO UPDATE
+                    SET metadata = EXCLUDED.metadata
+                    RETURNING id
+                """, 
+                    event_id,
+                    queue,
+                    namespace,
+                    scheduled_task_id,
+                    json.dumps(event_data),
+                    event_data.get('priority', 0),
+                    event_data.get('source', 'stream'),
+                    json.dumps({'consumer_group': consumer_group})
+                )
+                
+                # 2. 插入任务运行记录
+                consumer_name = event_data.get('consumer', '')
+                worker_id = event_data.get('worker_id', consumer_name)
+                
+                await self.async_conn.execute("""
+                    INSERT INTO task_runs (
+                        task_id, stream_id, task_name, consumer_group,
+                        consumer_name, worker_id, status, created_at
+                    ) VALUES (
+                        $1, $2, $3, $4, $5, $6, $7, CURRENT_TIMESTAMP
+                    )
+                    ON CONFLICT (task_id, consumer_group) DO UPDATE
+                    SET 
+                        status = EXCLUDED.status,
+                        updated_at = CURRENT_TIMESTAMP
+                """,
+                    task_id,
+                    event_id,
+                    task_name,
+                    consumer_group,
+                    consumer_name,
+                    worker_id,
+                    'pending'
+                )
+                
+            return task_id
+            
+        except Exception as e:
+            logger.error(f"Error recording task async: {e}")
+            raise
+            
+    def update_task_status(self, event_id: str, consumer_group: str, 
+                          status: str, **kwargs) -> bool:
+        """
+        更新任务执行状态（同步方法）
+        """
+        self.connect()
+        try:
+            with self.conn.cursor() as cursor:
+                # 构建更新字段
+                update_fields = ['status = %s', 'updated_at = CURRENT_TIMESTAMP']
+                update_values = [status]
+                
+                # 处理可选字段
+                if 'end_time' in kwargs:
+                    update_fields.append('end_time = %s')
+                    update_values.append(kwargs['end_time'])
+                    
+                if 'error_message' in kwargs:
+                    update_fields.append('error_message = %s')
+                    update_values.append(kwargs['error_message'])
+                    
+                if 'error_details' in kwargs:
+                    update_fields.append('error_details = %s')
+                    update_values.append(Json(kwargs['error_details']))
+                    
+                if 'result' in kwargs:
+                    update_fields.append('result = %s')
+                    update_values.append(Json(kwargs['result']))
+                    
+                if 'retry_count' in kwargs:
+                    update_fields.append('retry_count = %s')
+                    update_values.append(kwargs['retry_count'])
+                
+                # 添加WHERE条件的值
+                update_values.extend([event_id, consumer_group])
+                
+                # 更新task_runs表
+                cursor.execute(f"""
+                    UPDATE task_runs
+                    SET {', '.join(update_fields)}
+                    WHERE stream_id = %s AND consumer_group = %s
+                """, update_values)
+                
+                # tasks表已经没有status字段，不需要更新总体状态
+                
+                self.conn.commit()
+                return cursor.rowcount > 0
+                
+        except Exception as e:
+            self.conn.rollback()
+            logger.error(f"Error updating task status: {e}")
+            return False
+            
+    async def async_update_task_status(self, event_id: str, consumer_group: str,
+                                      status: str, **kwargs) -> bool:
+        """
+        异步更新任务执行状态
+        """
+        await self.async_connect()
+        try:
+            async with self.async_conn.transaction():
+                # 构建更新语句
+                set_clauses = ['status = $1', 'updated_at = CURRENT_TIMESTAMP']
+                values = [status]
+                param_count = 1
+                
+                # 添加可选字段
+                for field, db_field in [
+                    ('end_time', 'end_time'),
+                    ('error_message', 'error_message'),
+                    ('retry_count', 'retry_count')
+                ]:
+                    if field in kwargs:
+                        param_count += 1
+                        set_clauses.append(f'{db_field} = ${param_count}')
+                        values.append(kwargs[field])
+                        
+                # JSON字段
+                for field, db_field in [
+                    ('error_details', 'error_details'),
+                    ('result', 'result')
+                ]:
+                    if field in kwargs:
+                        param_count += 1
+                        set_clauses.append(f'{db_field} = ${param_count}::jsonb')
+                        values.append(json.dumps(kwargs[field]))
+                
+                # 添加WHERE条件的值
+                values.extend([event_id, consumer_group])
+                
+                # 更新task_runs表
+                await self.async_conn.execute(f"""
+                    UPDATE task_runs
+                    SET {', '.join(set_clauses)}
+                    WHERE stream_id = ${param_count + 1} 
+                    AND consumer_group = ${param_count + 2}
+                """, *values)
+                
+                # tasks表已经没有status字段，不需要更新总体状态
+                    
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error updating task status async: {e}")
+            return False
+            
+    def get_task_info(self, event_id: str) -> Optional[Dict]:
+        """获取任务信息（包括所有消费者组的执行情况）"""
+        self.connect()
+        try:
+            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
+                cursor.execute("""
+                    SELECT 
+                        t.*,
+                        array_agg(
+                            json_build_object(
+                                'consumer_group', tr.consumer_group,
+                                'status', tr.status,
+                                'start_time', tr.start_time,
+                                'end_time', tr.end_time,
+                                'duration_ms', tr.duration_ms,
+                                'error_message', tr.error_message,
+                                'retry_count', tr.retry_count
+                            )
+                        ) FILTER (WHERE tr.id IS NOT NULL) as runs
+                    FROM tasks t
+                    LEFT JOIN task_runs tr ON t.id = tr.task_id
+                    WHERE t.stream_id = %s
+                    GROUP BY t.id
+                """, (event_id,))
+                
+                return cursor.fetchone()
+                
+        except Exception as e:
+            logger.error(f"Error getting task info: {e}")
+            return None
+            
+    async def async_get_task_info(self, event_id: str) -> Optional[Dict]:
+        """异步获取任务信息"""
+        await self.async_connect()
+        try:
+            row = await self.async_conn.fetchrow("""
+                SELECT 
+                    t.*,
+                    array_agg(
+                        json_build_object(
+                            'consumer_group', tr.consumer_group,
+                            'status', tr.status,
+                            'start_time', tr.start_time,
+                            'end_time', tr.end_time,
+                            'duration_ms', tr.duration_ms,
+                            'error_message', tr.error_message,
+                            'retry_count', tr.retry_count
+                        )
+                    ) FILTER (WHERE tr.id IS NOT NULL) as runs
+                FROM tasks t
+                LEFT JOIN task_runs tr ON t.id = tr.task_id
+                WHERE t.stream_id = $1
+                GROUP BY t.id
+            """, event_id)
+            
+            return dict(row) if row else None
+            
+        except Exception as e:
+            logger.error(f"Error getting task info async: {e}")
+            return None

jettask/pg_consumer/sql_utils.py

@@ -0,0 +1,182 @@
+"""SQL 文件处理工具"""
+import re
+from typing import List
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def split_sql_statements(sql_content: str) -> List[str]:
+    """
+    智能分割 SQL 语句，正确处理 PL/pgSQL 函数定义
+    
+    Args:
+        sql_content: SQL 文件内容
+        
+    Returns:
+        分割后的 SQL 语句列表
+    """
+    statements = []
+    current_statement = []
+    in_function = False
+    in_string = False
+    in_dollar_quote = False
+    dollar_quote_tag = None
+    
+    lines = sql_content.split('\n')
+    
+    for line in lines:
+        # 跳过纯注释行
+        stripped = line.strip()
+        if stripped.startswith('--') and not current_statement:
+            continue
+            
+        # 检测函数开始
+        if not in_function and not in_string and not in_dollar_quote:
+            # 检查是否是函数定义的开始
+            if re.match(r'^\s*CREATE\s+(OR\s+REPLACE\s+)?FUNCTION', line, re.IGNORECASE):
+                in_function = True
+                logger.debug(f"进入函数定义: {line[:50]}")
+        
+        # 处理美元引号（$$）
+        if not in_string:
+            # 查找美元引号
+            dollar_matches = re.findall(r'\$([^$]*)\$', line)
+            for match in dollar_matches:
+                if not in_dollar_quote:
+                    # 开始美元引号
+                    in_dollar_quote = True
+                    dollar_quote_tag = match
+                    logger.debug(f"进入美元引号: ${match}$")
+                elif dollar_quote_tag == match:
+                    # 结束美元引号
+                    in_dollar_quote = False
+                    dollar_quote_tag = None
+                    logger.debug(f"退出美元引号: ${match}$")
+        
+        # 处理普通字符串（单引号）
+        if not in_dollar_quote:
+            # 简单的单引号检测（不处理转义）
+            quote_count = line.count("'") - line.count("\\'")
+            if quote_count % 2 == 1:
+                in_string = not in_string
+        
+        current_statement.append(line)
+        
+        # 检查语句是否结束
+        if stripped.endswith(';') and not in_string and not in_dollar_quote:
+            # 对于函数，需要特殊处理
+            if in_function:
+                # 检查是否是函数结束（以 language 'xxx' 或 $$ language 结尾）
+                if re.search(r"(language\s+['\"]?\w+['\"]?\s*;?\s*$|\$\$\s*language\s+['\"]?\w+['\"]?\s*;?\s*$)", 
+                           stripped, re.IGNORECASE):
+                    in_function = False
+                    logger.debug(f"函数定义结束: {line[:50]}")
+                    # 完整的函数定义
+                    full_statement = '\n'.join(current_statement)
+                    if full_statement.strip():
+                        statements.append(full_statement)
+                    current_statement = []
+                # else: 仍在函数内部，继续累积
+            else:
+                # 普通语句结束
+                full_statement = '\n'.join(current_statement)
+                if full_statement.strip():
+                    statements.append(full_statement)
+                current_statement = []
+    
+    # 处理最后可能未完成的语句
+    if current_statement:
+        full_statement = '\n'.join(current_statement)
+        if full_statement.strip():
+            statements.append(full_statement)
+    
+    logger.info(f"SQL 文件分割完成，共 {len(statements)} 个语句")
+    
+    # 清理语句
+    cleaned_statements = []
+    for stmt in statements:
+        stmt = stmt.strip()
+        if stmt and not stmt.startswith('--'):
+            # 移除尾部多余的分号（保留一个）
+            while stmt.endswith(';;'):
+                stmt = stmt[:-1]
+            cleaned_statements.append(stmt)
+    
+    return cleaned_statements
+
+
+def execute_sql_file(connection, file_path: str):
+    """
+    执行 SQL 文件（同步版本）
+    
+    Args:
+        connection: 数据库连接
+        file_path: SQL 文件路径
+    """
+    with open(file_path, 'r', encoding='utf-8') as f:
+        sql_content = f.read()
+    
+    statements = split_sql_statements(sql_content)
+    
+    cursor = connection.cursor()
+    try:
+        for i, stmt in enumerate(statements, 1):
+            try:
+                logger.debug(f"执行第 {i}/{len(statements)} 个语句")
+                cursor.execute(stmt)
+            except Exception as e:
+                if 'already exists' not in str(e):
+                    logger.warning(f"语句 {i} 执行失败: {e}")
+                    logger.debug(f"失败的语句: {stmt[:100]}...")
+        connection.commit()
+        logger.info(f"SQL 文件执行完成: {file_path}")
+    except Exception as e:
+        connection.rollback()
+        logger.error(f"SQL 文件执行失败: {e}")
+        raise
+    finally:
+        cursor.close()
+
+
+async def execute_sql_file_async(async_session, file_path: str):
+    """
+    执行 SQL 文件（异步版本）
+    
+    Args:
+        async_session: 异步数据库会话
+        file_path: SQL 文件路径
+    """
+    with open(file_path, 'r', encoding='utf-8') as f:
+        sql_content = f.read()
+    
+    statements = split_sql_statements(sql_content)
+    
+    # 为每个语句创建单独的事务，避免一个失败影响其他语句
+    for i, stmt in enumerate(statements, 1):
+        try:
+            logger.debug(f"异步执行第 {i}/{len(statements)} 个语句")
+            await async_session.execute(stmt)
+            await async_session.commit()  # 立即提交每个成功的语句
+        except Exception as e:
+            # 静默处理一些预期的错误
+            error_str = str(e)
+            if any(x in error_str for x in [
+                'already exists',
+                'InFailedSqlTransaction',
+                'DeadlockDetected',
+                'duplicate key',
+                'already exists'
+            ]):
+                logger.debug(f"语句 {i} 跳过（预期错误）: {type(e).__name__}")
+            else:
+                logger.warning(f"语句 {i} 执行失败: {e}")
+                logger.debug(f"失败的语句: {stmt[:100]}...")
+            
+            # 始终尝试回滚失败的事务
+            try:
+                await async_session.rollback()
+            except:
+                pass
+    
+    logger.info(f"SQL 文件异步执行完成: {file_path}")

jettask/scheduler/__init__.py

@@ -0,0 +1,17 @@
+"""
+定时任务调度模块
+支持 Redis + PostgreSQL 双存储方案
+"""
+
+from .models import ScheduledTask, TaskExecutionHistory
+from .scheduler import TaskScheduler
+from .loader import TaskLoader
+from .manager import ScheduledTaskManager
+
+__all__ = [
+    'ScheduledTask',
+    'TaskExecutionHistory', 
+    'TaskScheduler',
+    'TaskLoader',
+    'ScheduledTaskManager'
+]

jettask/scheduler/add_execution_count.sql

@@ -0,0 +1,11 @@
+-- 为scheduled_tasks表添加执行次数字段
+ALTER TABLE scheduled_tasks 
+ADD COLUMN IF NOT EXISTS execution_count INTEGER DEFAULT 0;
+
+-- 添加注释
+COMMENT ON COLUMN scheduled_tasks.execution_count IS '任务执行次数';
+
+-- 为现有记录设置初始值（可选，根据历史数据估算）
+UPDATE scheduled_tasks 
+SET execution_count = 0 
+WHERE execution_count IS NULL;

jettask/scheduler/add_priority_field.sql

@@ -0,0 +1,26 @@
+-- 添加priority字段到scheduled_tasks表
+-- 用于支持定时任务的优先级设置
+
+-- 检查是否已经存在priority字段，避免重复添加
+DO $$
+BEGIN
+    IF NOT EXISTS (
+        SELECT 1 
+        FROM information_schema.columns 
+        WHERE table_name = 'scheduled_tasks' 
+        AND column_name = 'priority'
+    ) THEN
+        ALTER TABLE scheduled_tasks 
+        ADD COLUMN priority INTEGER DEFAULT NULL;
+        
+        -- 添加注释
+        COMMENT ON COLUMN scheduled_tasks.priority IS '任务优先级 (1=最高, 数字越大优先级越低，NULL=默认最低)';
+        
+        -- 创建索引以提高查询性能
+        CREATE INDEX idx_scheduled_tasks_priority ON scheduled_tasks(priority);
+        
+        RAISE NOTICE 'Added priority column to scheduled_tasks table';
+    ELSE
+        RAISE NOTICE 'Priority column already exists in scheduled_tasks table';
+    END IF;
+END $$;

jettask/scheduler/add_scheduler_id.sql

@@ -0,0 +1,25 @@
+-- Migration to add scheduler_id column to scheduled_tasks table
+-- This allows unique identification and deduplication of tasks
+
+-- Add the scheduler_id column if it doesn't exist
+DO $$ 
+BEGIN
+    IF NOT EXISTS (
+        SELECT 1 
+        FROM information_schema.columns 
+        WHERE table_name = 'scheduled_tasks' 
+        AND column_name = 'scheduler_id'
+    ) THEN
+        ALTER TABLE scheduled_tasks 
+        ADD COLUMN scheduler_id VARCHAR(255) UNIQUE;
+        
+        -- Add comment
+        COMMENT ON COLUMN scheduled_tasks.scheduler_id IS 
+        'Unique identifier for the task, used for deduplication';
+    END IF;
+END $$;
+
+-- Create index for scheduler_id if it doesn't exist
+CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_scheduler_id 
+ON scheduled_tasks(scheduler_id) 
+WHERE scheduler_id IS NOT NULL;

jettask/scheduler/add_scheduler_id_index.sql

@@ -0,0 +1,10 @@
+-- Migration to add index for scheduler_id field
+-- This index is critical for performance as we heavily rely on scheduler_id for lookups
+
+-- Create unique index on scheduler_id for fast lookups
+CREATE UNIQUE INDEX IF NOT EXISTS idx_scheduled_tasks_scheduler_id 
+ON scheduled_tasks(scheduler_id);
+
+-- Also add a comment to clarify the importance
+COMMENT ON INDEX idx_scheduled_tasks_scheduler_id IS 
+'Unique index on scheduler_id for fast task lookups and deduplication';