onestep 0.5.0__py3-none-any.whl

onestep/__init__.py

@@ -0,0 +1,67 @@
+from .onestep import step
+from .cron import Cron
+from .retry import (
+    BaseRetry, BaseErrorCallback, NackErrorCallBack,
+    NeverRetry, AlwaysRetry, TimesRetry, RetryIfException, AdvancedRetry
+)
+from .broker import (
+    BaseBroker, BaseConsumer,
+    MemoryBroker, RabbitMQBroker, WebHookBroker, CronBroker, RedisStreamBroker, RedisPubSubBroker, SQSBroker
+)
+from .middleware import (
+    BaseMiddleware, BaseConfigMiddleware,
+    NacosPublishConfigMiddleware, NacosConsumeConfigMiddleware,
+    RedisPublishConfigMiddleware, RedisConsumeConfigMiddleware,
+    UniqueMiddleware, MemoryUniqueMiddleware,
+)
+from .exception import (
+    StopMiddleware, DropMessage,
+    RetryException, RetryInQueue, RetryInLocal
+)
+
+__all__ = [
+    'step', 'Cron',
+
+    # broker
+    'BaseBroker',
+    'BaseConsumer',
+    'MemoryBroker',
+    'RabbitMQBroker',
+    'WebHookBroker',
+    'CronBroker',
+    'RedisStreamBroker',
+    'RedisPubSubBroker',
+    'SQSBroker',
+
+    # retry
+    'BaseRetry',
+    'NeverRetry',
+    'AlwaysRetry',
+    'TimesRetry',
+    'RetryIfException',
+    'AdvancedRetry',
+    # error callback
+    'BaseErrorCallback',
+    'NackErrorCallBack',
+
+    # middleware
+    'BaseMiddleware',
+    'BaseConfigMiddleware',
+    'NacosPublishConfigMiddleware',
+    'NacosConsumeConfigMiddleware',
+    'RedisPublishConfigMiddleware',
+    'RedisConsumeConfigMiddleware',
+    'UniqueMiddleware',
+    'MemoryUniqueMiddleware',
+
+    # exception
+    'StopMiddleware',
+    'DropMessage',
+    'RetryException',
+    'RetryInQueue',
+    'RetryInLocal',
+
+    '__version__'
+]
+
+__version__ = '0.4.4'

onestep/_utils.py

@@ -0,0 +1,21 @@
+import functools
+import logging
+
+
+def catch_error(return_val=None):
+    """
+    捕获异常装饰器
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                logging.debug(e)
+                return return_val
+
+        return wrapper
+
+    return decorator

onestep/broker/__init__.py

@@ -0,0 +1,22 @@
+from .base import (
+    BaseBroker, BaseConsumer
+)
+from .memory import MemoryBroker, MemoryConsumer
+from .webhook import WebHookBroker
+from .rabbitmq import RabbitMQBroker, RabbitMQConsumer
+from .mysql import MysqlBroker
+from .sqs import SQSBroker, SQSConsumer, SNSBroker
+from .redis import RedisStreamBroker, RedisPubSubBroker
+from .cron import CronBroker
+
+
+__all__ = [
+    "BaseBroker", "BaseConsumer",
+    "MemoryBroker", "MemoryConsumer",
+    "WebHookBroker",
+    "RabbitMQBroker", "RabbitMQConsumer",
+    "MysqlBroker",
+    "SQSBroker", "SQSConsumer", "SNSBroker",
+    "RedisStreamBroker", "RedisPubSubBroker",
+    "CronBroker"
+]

onestep/broker/base.py

@@ -0,0 +1,144 @@
+# -*- coding: utf-8 -*-
+import abc
+import logging
+from queue import Queue, Empty
+from typing import Any, Optional, List, Callable
+
+from onestep.middleware import BaseMiddleware
+from onestep.exception import StopMiddleware
+from onestep.message import Message
+
+logger = logging.getLogger(__name__)
+
+
+class BaseBroker:
+    message_cls = Message
+
+    def __init__(self,
+                 name: Optional[str] = None,
+                 queue: Optional[Queue] = None,
+                 middlewares: Optional[List[BaseMiddleware]] = None,
+                 once: bool = False,
+                 cancel_consume: Optional[Callable] = None):
+        """
+        @param name: broker name
+        @param queue: broker queue name
+        @param middlewares: broker middlewares
+        @param once: just run once
+            if once is True, when broker receive a message, it will shutdown
+        @param cancel_consume: cancel consume
+            cancel_consume(message) -> bool
+            if cancel_consume return True, broker will shutdown
+        """
+        self.queue = queue
+        self.name = name or "broker"
+        self.middlewares = []
+        self.once = once
+        self.cancel_consume = cancel_consume
+
+        if middlewares:
+            for middleware in middlewares:
+                self.add_middleware(middleware)
+
+    def add_middleware(self, middleware: BaseMiddleware):
+        if not isinstance(middleware, BaseMiddleware):
+            raise TypeError(f"middleware must be BaseMiddleware instance, not {type(middleware)}")
+        self.middlewares.append(middleware)
+
+    def send(self, message, *args, **kwargs):
+        """对消息进行预处理，然后再发送"""
+        if not isinstance(message, Message):
+            message = self.message_cls(body=message)
+        self.before_emit("send", message=message, step=kwargs.get("step"))
+        # TODO: 对消息发送进行N次重试，确保消息发送成功。
+        result = self.publish(message.to_json(), *args, **kwargs)
+        self.after_emit("send", message=message, step=kwargs.get("step"))
+        return result
+
+    @abc.abstractmethod
+    def publish(self, message: Any, *args, **kwargs):
+        """
+        将消息原样发布到 broker 中。如果当前Broker是Job的to_broker, 则必须实现此方法
+        """
+        raise NotImplementedError('Please implement in subclasses.')
+
+    @abc.abstractmethod
+    def consume(self, *args, **kwargs):
+        """
+        消费消息。如果当前Broker是Job的from_broker, 则必须实现此方法
+        """
+        raise NotImplementedError('Please implement in subclasses.')
+
+    @abc.abstractmethod
+    def confirm(self, message: Message):
+        """确认消息"""
+        raise NotImplementedError('Please implement in subclasses.')
+
+    @abc.abstractmethod
+    def reject(self, message: Message):
+        """拒绝消息"""
+        raise NotImplementedError('Please implement in subclasses.')
+
+    @abc.abstractmethod
+    def requeue(self, message: Message, is_source=False):
+        """
+        重发消息：先拒绝 再 重入
+        is_source = False 重入使用消息的当前状态
+        is_source = True 重入使用消息的初始状态
+        """
+        raise NotImplementedError('Please implement in subclasses.')
+
+    def before_emit(self, signal, **kwargs):
+        signal = "before_" + signal
+        self._emit(signal, **kwargs)
+
+    def after_emit(self, signal, **kwargs):
+        signal = "after_" + signal
+        self._emit(signal, **kwargs)
+
+    def _emit(self, signal, **kwargs):
+        for middleware in self.middlewares:
+            if not hasattr(middleware, signal):
+                continue
+            try:
+                params = dict(kwargs)
+                params.setdefault("broker", self)
+                params.setdefault("step", None)
+                getattr(middleware, signal)(**params)
+            except StopMiddleware:
+                break
+
+    def shutdown(self):
+        """关闭Broker"""
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} {self.name}>"
+
+    def __str__(self):
+        return self.name
+
+
+class BaseConsumer:
+
+    def __init__(self, broker: BaseBroker, *args, **kwargs):
+        self.queue = broker.queue
+        self.message_cls = broker.message_cls or Message
+        self.timeout = kwargs.pop("timeout", 1000)
+
+    def __next__(self):
+        try:
+            q = self.queue
+            if q is None:
+                return None
+            timeout_ms = self.timeout
+            if isinstance(timeout_ms, (int, float)) and timeout_ms > 0:
+                data = q.get(timeout=timeout_ms / 1000.0)
+            else:
+                # 当超时为0、负数或非数字时，使用非阻塞获取以避免ValueError
+                data = q.get_nowait()
+            return self.message_cls.from_broker(broker_message=data)
+        except (Empty, ValueError):
+            return None
+
+    def __iter__(self):
+        return self

onestep/broker/cron.py

@@ -0,0 +1,51 @@
+"""
+使用 CRON 表达式触发任务执行
+
+支持人性化 DSL 与宏：
+- DSL：Cron.every(minutes=5)、Cron.daily(at="09:00")、Cron.weekly(on=["mon","fri"], at="10:30") 等
+- 宏：@hourly/@daily/@weekly/@monthly/@yearly 以及扩展 @workdays/@weekends/@every 5m/2h/3d/1mo
+"""
+import logging
+import threading
+from datetime import datetime
+from typing import Any
+
+from croniter import croniter
+from ..cron import resolve_cron
+
+from .memory import MemoryBroker, MemoryConsumer
+
+logger = logging.getLogger(__name__)
+
+
+class CronBroker(MemoryBroker):
+
+    def __init__(self, cron, name=None, middlewares=None, body: Any = None, start_time=None, *args, **kwargs):
+        super().__init__(name=name, middlewares=middlewares, *args, **kwargs)
+        # 支持 DSL/宏/原始字符串：统一解析为标准表达式
+        self.cron = resolve_cron(cron)
+        self.start_time = start_time or datetime.now()
+        self.itr = croniter(self.cron, self.start_time)
+        self.next_fire_time = self.itr.get_next(datetime)
+        self.body = body
+        self._thread = None
+
+    def _scheduler(self):
+        if self.next_fire_time <= datetime.now():
+            self.next_fire_time = self.itr.get_next(datetime)
+            self.publish(self.body)
+
+        self._thread = threading.Timer(interval=1, function=self._scheduler)
+        self._thread.start()
+
+    def consume(self, *args, **kwargs):
+        self._scheduler()
+        return CronConsumer(self, *args, **kwargs)
+
+    def shutdown(self):
+        if self._thread:
+            self._thread.cancel()
+
+
+class CronConsumer(MemoryConsumer):
+    ...

onestep/broker/memory.py

@@ -0,0 +1,75 @@
+import logging
+import json
+from queue import Queue, Full as FullException
+from typing import Any
+
+from .base import BaseBroker, BaseConsumer
+
+from ..message import Message
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryMessage(Message):
+
+    @classmethod
+    def from_broker(cls, broker_message: Any):
+        if isinstance(broker_message, (str, bytes, bytearray)):
+            try:
+                message = json.loads(broker_message)
+            except json.JSONDecodeError:
+                message = {"body": broker_message}
+        else:
+            message = broker_message
+        if not isinstance(message, dict):
+            message = {"body": message}
+        if "body" not in message:
+            # 来自 外部的消息 可能没有 body, 故直接认为都是 message.body
+            message = {"body": message}
+        extra_val = message.get("extra")
+        if not isinstance(extra_val, dict):
+            extra_val = None
+        return cls(body=message.get("body"), extra=extra_val, message=broker_message)
+
+
+class MemoryBroker(BaseBroker):
+    message_cls = MemoryMessage
+
+    def __init__(self, maxsize=0, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.queue = Queue(maxsize)
+
+    def publish(self, message: Any, *args, **kwargs):
+        try:
+            self.queue.put_nowait(message)
+        except FullException:
+            logger.warning("CronBroker queue is full, skip this task, "
+                           "you can increase maxsize with `maxsize` argument")
+
+    def consume(self, *args, **kwargs):
+        return MemoryConsumer(self, *args, **kwargs)
+
+    def confirm(self, message: Message):
+        """确认消息"""
+        pass
+
+    def reject(self, message: Message):
+        """拒绝消息"""
+        pass
+
+    def requeue(self, message: Message, is_source=False):
+        """重发消息：先拒绝 再 重入"""
+        if is_source:
+            self.publish(message.message)
+        else:
+            self.send(message)
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} {self.name}>"
+
+    def __str__(self):
+        return self.name
+
+
+class MemoryConsumer(BaseConsumer):
+    ...

onestep/broker/mysql.py

@@ -0,0 +1,63 @@
+from queue import Queue
+from typing import Optional, Dict, Any, cast
+
+try:
+    from use_mysql import MysqlStore, Model, SQLModel
+except ImportError:
+    MysqlStore = None
+
+from .base import BaseBroker, BaseConsumer
+from ..message import Message
+
+
+class MysqlBroker(BaseBroker):
+
+    def __init__(
+        self,
+        params: Optional[Dict] = None,
+        auto_create: Optional[bool] = True,
+        *args,
+        **kwargs
+    ):
+        if MysqlStore is None or Model is None:
+            raise ImportError("未安装 use-mysql 依赖")
+
+        super().__init__(*args, **kwargs)
+        self.queue = Queue()
+        params = params or {}
+        auto_create = auto_create or False
+        self.client = MysqlStore(generate_schemas=auto_create, **params)
+        self.client.init()
+        self._shutdown = False
+
+    def send(self, message, *args, **kwargs):
+        if not isinstance(message, Message):
+            message = self.message_cls(body=message)
+        return self.publish(message, *args, **kwargs)
+
+
+    def publish(self, message: Any, **kwargs):
+        body_raw = message.body if isinstance(message, Message) else message
+        if not isinstance(body_raw, SQLModel):
+            raise TypeError("MySQL Broker 仅支持 Model 类型的消息体")
+
+        body = cast(SQLModel, body_raw)
+        return self.client.create(body.__class__, **body.model_dump())
+
+  
+    def confirm(self, message: Message):
+        return None
+
+    def reject(self, message: Message):
+        return None
+
+    def requeue(self, message: Message, is_source=False):
+        return None
+
+    def shutdown(self):
+        self._shutdown = True
+        self.client.shutdown()
+
+
+class MysqlConsumer(BaseConsumer):
+    ...

onestep/broker/rabbitmq.py

@@ -0,0 +1,153 @@
+import json
+import threading
+from queue import Queue
+from typing import Optional, Dict, Any
+
+try:
+    import amqpstorm
+    from use_rabbitmq import useRabbitMQ as RabbitMQStore
+except ImportError:
+    amqpstorm = None
+    RabbitMQStore = None
+
+from .base import BaseBroker, BaseConsumer
+from ..message import Message
+
+
+class _RabbitMQMessage(Message):
+
+    @classmethod
+    def from_broker(cls, broker_message):
+        try:
+            message = json.loads(broker_message.body)
+        except json.JSONDecodeError:
+            message = {"body": broker_message.body}
+        if not isinstance(message, dict):
+            message = {"body": message}
+        if "body" not in message:
+            # 来自 外部的消息 可能没有 body, 故直接认为都是 message.body
+            message = {"body": message}
+
+        return cls(body=message.get("body"), extra=message.get("extra"), message=broker_message)
+
+
+class RabbitMQBroker(BaseBroker):
+    message_cls = _RabbitMQMessage
+
+    def __init__(
+        self,
+        queue_name,
+        params: Optional[Dict] = None,
+        prefetch: Optional[int] = 1,
+        auto_create: Optional[bool]=True,
+        queue_params: Optional[Dict]=None,
+        *args,
+        **kwargs
+        ):
+        """
+        Initializes the RabbitMQ broker.
+
+        Args:
+            queue_name (str): The name of the queue.
+            params (Optional[Dict], optional): Parameters for RabbitMQStore. Defaults to None.
+            prefetch (Optional[int], optional): Number of messages to prefetch. Defaults to 1.
+            auto_create (Optional[bool], optional): Whether to automatically create the queue. Defaults to True.
+            queue_params (Optional[Dict], optional): Parameters for queue declaration. Defaults to None.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Attributes:
+            queue_name (str): The name of the queue.
+            queue (Queue): The queue instance.
+            client (RabbitMQStore): The RabbitMQ client instance.
+            prefetch (int): Number of messages to prefetch.
+            threads (list): List of threads.
+        """
+        
+        if RabbitMQStore is None:
+            raise ImportError("RabbitMQ dependencies not installed. Please install 'use-rabbitmq' package.")
+        
+        super().__init__(*args, **kwargs)
+        self.queue_name = queue_name
+        self.queue = Queue()
+        params = params or {}
+        self.client = RabbitMQStore(**params)
+        if auto_create:
+            self.client.declare_queue(self.queue_name, **(queue_params or {}))
+        self.prefetch = prefetch
+        self.threads = []
+        self._consuming_started = False
+        self._consume_lock = threading.Lock()
+
+    def _consume(self, *args, **kwargs):
+        def callback(message):
+            self.queue.put(message)
+
+        prefetch = kwargs.pop("prefetch", self.prefetch)
+        self.client.start_consuming(queue_name=self.queue_name, callback=callback, prefetch=prefetch, **kwargs)
+
+    def consume(self, *args, **kwargs):
+        daemon = kwargs.pop('daemon', True)
+        with self._consume_lock:
+            if not self._consuming_started:
+                thread = threading.Thread(target=self._consume, args=args, kwargs=kwargs)
+                thread.daemon = daemon
+                thread.start()
+                self.threads.append(thread)
+                self._consuming_started = True
+        return RabbitMQConsumer(self)
+
+    def publish(self, message: Any, properties: Optional[dict] = None, **kwargs):
+        """发布消息"""
+        self.client.send(self.queue_name, message, properties=properties, **kwargs)
+
+    def confirm(self, message: Message):
+        """确认消息"""
+        broker_msg = getattr(message, "message", None)
+        if broker_msg is not None and hasattr(broker_msg, "ack"):
+            broker_msg.ack()
+        else:
+            # 无可确认的原始消息，忽略确认
+            return
+
+    def reject(self, message: Message):
+        """拒绝消息"""
+        broker_msg = getattr(message, "message", None)
+        if broker_msg is not None and hasattr(broker_msg, "reject"):
+            broker_msg.reject(requeue=False)
+        else:
+            # 无法拒绝（无原始消息对象），忽略
+            return
+
+    def requeue(self, message: Message, is_source=False):
+        """
+        重发消息：先拒绝 再 重入
+        
+        :param message: 消息
+        :param is_source: 是否是原始消息，True: 使用原始消息重入当前队列，False: 使用消息的最新数据重入当前队列
+        """
+        broker_msg = getattr(message, "message", None)
+        if broker_msg is not None and hasattr(broker_msg, "reject"):
+            if is_source:
+                broker_msg.reject(requeue=True)
+            else:
+                broker_msg.reject(requeue=False)
+                self.send(message)
+        else:
+            # 没有原始消息控制能力，回退为直接发送当前消息状态
+            if is_source and broker_msg is not None and hasattr(broker_msg, "body"):
+                # 尝试使用原始消息体重入队列
+                self.client.send(self.queue_name, broker_msg.body)
+            else:
+                self.send(message)
+
+    def shutdown(self):
+        self.client.shutdown()
+        for thread in self.threads:
+            thread.join()
+        self._consuming_started = False
+        self.threads = []
+
+
+class RabbitMQConsumer(BaseConsumer):
+    ...

onestep/broker/redis/__init__.py

@@ -0,0 +1,9 @@
+from .stream import RedisStreamBroker, RedisStreamConsumer
+from .pubsub import RedisPubSubBroker, RedisPubSubConsumer
+
+__all__ = [
+    "RedisStreamBroker",
+    "RedisStreamConsumer",
+    "RedisPubSubBroker",
+    "RedisPubSubConsumer"
+]