fast-mq-task 1.0.0__tar.gz

fast-mq-task-1.0.0/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.1
+Name: fast-mq-task
+Version: 1.0.0
+Summary: fast rabbitmq task for zsodata
+Home-page: http://www.zsodata.com
+Author: zsodata
+Author-email: team@zso.io
+License: BSD License
+Platform: all
+Requires-Python: >=3.7
+
+pip uninstall fast-mq-task
+
+
+pip install --upgrade fast-mq-task -i https://pypi.python.org/simple
+pip show fast-mq-task

fast-mq-task-1.0.0/README.rst

@@ -0,0 +1,5 @@
+pip uninstall fast-mq-task
+
+
+pip install --upgrade fast-mq-task -i https://pypi.python.org/simple
+pip show fast-mq-task

fast-mq-task-1.0.0/fast_mq_task/client.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*-
+import logging
+from typing import Dict, Union
+from aio_pika import connect, Exchange, ExchangeType, Channel, Queue
+from aio_pika.abc import AbstractExchange, AbstractQueue
+
+from .keys import get_default_exchange_name, get_default_routing_key, get_default_queue_name
+
+
+# --------------------------
+# 客户端基础模块 (client.py)
+# --------------------------
+class RabbitMQClient:
+    def __init__(self, amqp_url: str):
+        self.amqp_url = amqp_url
+        self.connection = None
+        self.channels: Dict[str, Channel] = {}
+        self.exchanges: Dict[str, Union[Exchange, AbstractExchange]] = {}
+        self.queues: Dict[str, Union[Queue, AbstractQueue]] = {}
+
+    async def connect(self):
+        self.connection = await connect(self.amqp_url)
+        return self
+
+    async def close(self):
+        for channel in self.channels.values():
+            await channel.close()
+        await self.connection.close()
+
+    async def get_channel(self, task_type: str, prefetch: int = 4) -> Channel:
+        task_type = task_type or 'default'
+        if task_type not in self.channels:
+            channel = await self.connection.channel()
+            await channel.set_qos(
+                prefetch_count=prefetch
+            )
+            logging.info(f"信道定义: task_type={task_type}, prefetch={prefetch}")
+            self.channels[task_type] = channel
+
+        return self.channels[task_type]
+
+    async def get_exchange(self, channel: Channel, exchange_name: str = None, exchange_type: str = "topic") -> Exchange:
+        # 统一交换器类型为小写
+        exchange_type = exchange_type.lower()
+
+        # 添加交换器类型校验
+        if exchange_type not in ["direct", "topic", "fanout", "headers"]:
+            raise ValueError(f"Invalid exchange type: {exchange_type}")
+        if not exchange_name:
+            exchange_name = get_default_exchange_name()
+
+        if exchange_name not in self.exchanges:
+            self.exchanges[exchange_name] = await channel.declare_exchange(
+                name=exchange_name,
+                type=exchange_type,
+                durable=True,  # 持久化交换器
+                auto_delete=False,
+                arguments={
+                    "x-queue-type": "quorum"  # 使用高可用队列类型
+                }
+            )
+            logging.info(f"交换机定义: exchange={exchange_name}, type={exchange_type}")
+
+        return self.exchanges[exchange_name]
+
+    async def ensure_declare(self, task_type: str, prefetch: int = 10):
+        # 创建专属channel避免相互影响
+        channel = await self.get_channel(task_type, prefetch)
+
+        exchange_key = get_default_exchange_name()
+        exchange_type = ExchangeType.TOPIC
+        routing_key = get_default_routing_key(task_type)
+        queue_key = get_default_queue_name(task_type)
+
+        # 声明交换器和队列
+        exchange = await self.get_exchange(channel=channel, exchange_name=exchange_key, exchange_type=exchange_type)
+        if queue_key not in self.queues:
+            queue = await channel.declare_queue(
+                name=queue_key,
+                durable=True,
+                arguments={
+                    'x-max-priority': 10  # 支持优先级队列
+                }
+            )
+            await queue.bind(exchange, routing_key)
+            self.queues[queue_key] = queue
+
+        return self.queues[queue_key], exchange

fast-mq-task-1.0.0/fast_mq_task/consumer.py

@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+import asyncio
+import inspect
+import logging
+import traceback
+from functools import partial
+from typing import Callable, Dict
+from aio_pika import IncomingMessage
+from .client import RabbitMQClient
+from .decorators import TaskRegistry
+from .models import TaskMessage, RabbitMeta
+from concurrent.futures import ThreadPoolExecutor
+
+
+# --------------------------
+# 消费者模块 (consumer.py)
+# --------------------------
+class TaskConsumer(RabbitMQClient):
+    def __init__(
+            self,
+            amqp_url: str,
+            task_config: Dict[str, Dict] = None,  # 新增任务专属配置
+            default_prefetch: int = 10
+    ):
+        super().__init__(amqp_url)
+        self.task_config = task_config or {}
+        self.default_prefetch = default_prefetch
+        self._running = False
+        self._executors: Dict[str, ThreadPoolExecutor] = {}  # 任务专属线程池
+        self._semaphores: Dict[str, asyncio.Semaphore] = {}  # 异步信号量控制
+
+    async def _init_resources(self):
+        """初始化每个任务的执行资源"""
+        for idx, task_type in enumerate(TaskRegistry._handlers):
+            # 创建独立线程池
+            config = self.task_config.get(task_type, {})
+            max_workers = config.get('max_workers', 4)
+            self._executors[task_type] = ThreadPoolExecutor(
+                max_workers=max_workers,
+                thread_name_prefix=f"Worker-{task_type}-"
+            )
+
+            # 创建异步信号量
+            prefetch = config.get('prefetch', self.default_prefetch)
+            self._semaphores[task_type] = asyncio.Semaphore(prefetch)
+            logger.info(f'消费者初始化[{idx}]: task_type={task_type}, prefetch={prefetch}, max_workers={max_workers}')
+
+    async def _process_message(self, rabbit_meta: RabbitMeta, msg: IncomingMessage):
+        """处理消息的优化版本"""
+        async with msg.process():
+            try:
+                handler: Callable = rabbit_meta.handler_func
+                # 获取任务配置
+                task_type = rabbit_meta.task_type
+                semaphore = self._semaphores[task_type]
+                executor = self._executors[task_type]
+
+                # 使用专属信号量控制并发
+                async with semaphore:
+                    task = TaskMessage.parse_raw(msg.body)
+
+                    if inspect.iscoroutinefunction(handler):
+                        # 异步任务直接await
+                        await handler(task)
+                    else:
+                        # 同步任务放入专属线程池
+                        loop = asyncio.get_event_loop()
+                        await loop.run_in_executor(
+                            executor,
+                            partial(handler, task)
+                        )
+            except Exception as e:
+                logging.error(f"处理失败: {traceback.format_exc()}")
+
+    async def _setup_queues(self):
+        """队列初始化优化"""
+        for rabbit_meta in TaskRegistry._handlers.values():
+            queue, _ = await self.ensure_declare(task_type=rabbit_meta.task_type, prefetch=rabbit_meta.prefetch or self.default_prefetch)
+
+            # 为每个队列创建独立消费者
+            await queue.consume(
+                partial(self._process_message, rabbit_meta)
+            )
+
+    async def start(self):
+        """启动消费者"""
+        await self._init_resources()
+        await self._setup_queues()
+        self._running = True
+        while self._running:
+            await asyncio.sleep(1)
+
+    async def stop(self):
+        """优雅关闭"""
+        self._running = False
+        for executor in self._executors.values():
+            executor.shutdown(wait=True)
+        await super().close()

fast-mq-task-1.0.0/fast_mq_task/decorators.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+import logging
+from typing import Callable, Dict
+from .models import RabbitMeta
+
+
+# --------------------------
+# 装饰器模块 (decorators.py)
+# --------------------------
+class TaskRegistry:
+    _handlers: Dict[str, RabbitMeta] = {}
+
+    @classmethod
+    def register(cls, task_type: str) -> Callable:
+        # 新增参数校验
+        if not task_type:
+            raise ValueError("`task_type`不能全部为空")
+
+        def decorator(func: Callable) -> Callable:
+            meta = RabbitMeta(
+                task_type=task_type,
+                handler_func=func
+            )
+            cls._handlers[task_type] = meta
+            logging.info(f"处理器注册: task_type={task_type}, exchange={meta.queue_key}, func={func}")
+
+            return func
+
+        return decorator
+
+
+task_handler = TaskRegistry.register

fast-mq-task-1.0.0/fast_mq_task/keys.py

@@ -0,0 +1,19 @@
+import os
+
+
+def get_default_exchange_name():
+    return f"{os.getenv('RABBIT_EXCHANGE_PREFIX') or 'fast'}.{os.getenv('APP_CODE') or 'default'}"
+
+
+def get_default_routing_key(task_type: str):
+    if not task_type:
+        raise ValueError(f'task_type must not none: task_type={task_type}')
+
+    return f"{os.getenv('RABBIT_ROUTING_KEY_PREFIX') or 'fast'}.{os.getenv('APP_CODE') or 'default'}.{task_type}"
+
+
+def get_default_queue_name(task_type: str):
+    if not task_type:
+        raise ValueError(f'task_type must not none: task_type={task_type}')
+
+    return f"{os.getenv('RABBIT_QUEUE_PREFIX') or 'fast'}.{os.getenv('APP_CODE') or 'default'}.{task_type}"

fast-mq-task-1.0.0/fast_mq_task/models.py

@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+import time
+import uuid
+from typing import Any, Dict, Callable
+from aio_pika import ExchangeType
+from pydantic import BaseModel
+
+from . import keys
+
+
+# --------------------------
+# 模型类模块 (models.py)
+# --------------------------
+class RabbitMeta(BaseModel):
+    """
+    消息发布元数据
+    """
+    # 交换机键
+    task_type: str = None
+    # 交换机键
+    exchange_key: str = None
+    # 交换机类型
+    exchange_type: str = None
+    # 绑定关系（路由）键
+    routing_key: str = None
+    # 队列名称
+    queue_key: str = None
+    # 队列名称
+    prefetch: int = None
+    # 队列名称
+    max_workers: int = None
+
+    # 处理函数
+    handler_func: Callable = None
+
+    def __init__(self, task_type: str = None, prefetch: int = None, max_workers: int = None, handler_func: Callable = None, **data: Any):
+        super().__init__(**data)
+        self.task_type = task_type
+
+        self.prefetch = prefetch
+        self.max_workers = max_workers
+        self.handler_func = handler_func
+
+        self.exchange_key = keys.get_default_exchange_name()
+        self.exchange_type = ExchangeType.TOPIC
+        self.routing_key = keys.get_default_routing_key(task_type)
+        self.queue_key = keys.get_default_queue_name(task_type)
+
+
+class TaskMessage(BaseModel):
+    task_type: str
+    data: Dict[str, Any]
+    task_id: str = str(uuid.uuid4())
+    created_at: int = int(time.time())
+
+    def get_routing_key(self):
+        return keys.get_default_routing_key(self.task_type)

fast-mq-task-1.0.0/fast_mq_task/producer.py

@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+import time
+from aio_pika import Message, DeliveryMode
+from .client import RabbitMQClient
+from .models import TaskMessage
+
+
+# --------------------------
+# 生产者模块 (producer.py)
+# --------------------------
+class TaskProducer(RabbitMQClient):
+    async def publish_task(self, task: TaskMessage):
+        # 确保队列定义
+        _, exchange = await self.ensure_declare(task_type=task.task_type)
+
+        message = Message(
+            body=task.json().encode(),
+            content_type="application/json",
+            headers={
+                "x-task-type": task.task_type,
+                "x-created-at": time.time()
+            },
+            delivery_mode=DeliveryMode.PERSISTENT  # 添加消息持久化
+        )
+        await exchange.publish(
+            message=message,
+            routing_key=task.get_routing_key()
+        )

fast-mq-task-1.0.0/fast_mq_task.egg-info/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.1
+Name: fast-mq-task
+Version: 1.0.0
+Summary: fast rabbitmq task for zsodata
+Home-page: http://www.zsodata.com
+Author: zsodata
+Author-email: team@zso.io
+License: BSD License
+Platform: all
+Requires-Python: >=3.7
+
+pip uninstall fast-mq-task
+
+
+pip install --upgrade fast-mq-task -i https://pypi.python.org/simple
+pip show fast-mq-task

fast-mq-task-1.0.0/fast_mq_task.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+MANIFEST.in
+README.rst
+setup.py
+fast_mq_task/__init__.py
+fast_mq_task/client.py
+fast_mq_task/consumer.py
+fast_mq_task/decorators.py
+fast_mq_task/keys.py
+fast_mq_task/models.py
+fast_mq_task/producer.py
+fast_mq_task.egg-info/PKG-INFO
+fast_mq_task.egg-info/SOURCES.txt
+fast_mq_task.egg-info/dependency_links.txt
+fast_mq_task.egg-info/top_level.txt

fast-mq-task-1.0.0/fast_mq_task.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

fast-mq-task-1.0.0/fast_mq_task.egg-info/top_level.txt

@@ -0,0 +1 @@
+fast_mq_task

fast-mq-task-1.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

fast-mq-task-1.0.0/setup.py

@@ -0,0 +1,21 @@
+from distutils.core import setup
+from setuptools import find_packages
+
+with open('README.rst', 'r', encoding="utf-8") as f:
+    long_description = f.read()
+
+setup(name='fast-mq-task',
+      version='1.0.0',
+      description='fast rabbitmq task for zsodata',
+      long_description=long_description,
+      author='zsodata',
+      author_email='team@zso.io',
+      url='http://www.zsodata.com',
+      install_requires=[
+      ],
+      python_requires='>=3.7',
+      license='BSD License',
+      packages=find_packages(),
+      platforms=['all'],
+      include_package_data=True
+      )