onestep 0.5.0__py3-none-any.whl

onestep/broker/redis/pubsub.py

@@ -0,0 +1,93 @@
+import json
+import threading
+from queue import Queue
+from typing import Any
+
+try:
+    from use_redis import useRedis
+except ImportError:
+    ...
+
+from ..base import BaseBroker, BaseConsumer, Message
+
+
+class _RedisPubSubMessage(Message):
+
+    @classmethod
+    def from_broker(cls, broker_message: Any):
+        if "channel" in broker_message:
+            try:
+                message = json.loads(broker_message.get("data"))  # 已转换的 message
+            except (json.JSONDecodeError, TypeError):
+                message = {"body": broker_message.get("data")}  # 未转换的 message
+        else:
+            # 来自 外部的消息 直接认为都是 message.body
+            message = {"body": broker_message.body}
+
+        yield cls(body=message.get("body"), extra=message.get("extra"), message=broker_message)
+
+
+class RedisPubSubBroker(BaseBroker):
+    """ Redis PubSub Broker """
+    message_cls = _RedisPubSubMessage
+
+    def __init__(self, channel: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.channel = channel
+        self.queue = Queue()
+
+        self.threads = []
+
+        self.client = useRedis(**kwargs).connection
+
+    def _consume(self):
+        def callback(message: dict):
+            if message.get('type') != 'message':
+                return
+            self.queue.put(message)
+
+        ps = self.client.pubsub()
+        ps.subscribe(self.channel)
+        for message in ps.listen():
+            callback(message)
+
+    def consume(self, *args, **kwargs):
+        daemon = kwargs.pop('daemon', True)
+        thread = threading.Thread(target=self._consume, *args, **kwargs)
+        thread.daemon = daemon
+        thread.start()
+        self.threads.append(thread)
+        return RedisPubSubConsumer(self)
+
+    def send(self, message: Any, *args, **kwargs):
+        """Publish message to the Redis channel"""
+        if not isinstance(message, Message):
+            message = self.message_cls(body=message)
+
+        self.client.publish(self.channel, message.to_json(), *args, **kwargs)
+
+    publish = send
+
+    def confirm(self, message: Message):
+        pass
+
+    def reject(self, message: Message):
+        pass
+
+    def requeue(self, message: Message, is_source=False):
+        """
+         重发消息：先拒绝 再 重入
+
+         :param message: 消息
+         :param is_source: 是否是原始消息消息，True: 使用原始消息重入当前队列，False: 使用消息的最新数据重入当前队列
+         """
+        self.reject(message)
+
+        if is_source:
+            self.client.publish(self.channel, message.message['data'])
+        else:
+            self.send(message)
+
+
+class RedisPubSubConsumer(BaseConsumer):
+    ...

onestep/broker/redis/stream.py

@@ -0,0 +1,114 @@
+import json
+import threading
+import uuid
+from queue import Queue
+from typing import Optional, Dict, Any
+
+try:
+    from use_redis import useRedisStreamStore, RedisStreamMessage
+except ImportError:
+    ...
+
+from ..base import BaseBroker, BaseConsumer, Message
+
+
+class _RedisStreamMessage(Message):
+    @classmethod
+    def from_broker(cls, broker_message: "RedisStreamMessage"):
+        if "_message" in broker_message.body:
+            # 来自 RedisStreamBroker.send 的消息，message.body 默认是存于 _message 字段中
+            try:
+                message = json.loads(broker_message.body.get("_message"))  # 已转换的 message
+            except (json.JSONDecodeError, TypeError):
+                message = {"body": broker_message.body.get("_message")}  # 未转换的 message
+        else:
+            # 来自 外部的消息 直接认为都是 message.body
+            message = {"body": broker_message.body}
+
+        yield cls(body=message.get("body"), extra=message.get("extra"), message=broker_message)
+
+
+class RedisStreamBroker(BaseBroker):
+    """ Redis Stream Broker """
+    message_cls = _RedisStreamMessage
+
+    def __init__(
+            self,
+            stream: str,
+            group: str = "onestep",
+            params: Optional[Dict] = None,
+            prefetch: Optional[int] = 1,
+            stream_max_entries: int = 0,
+            redeliver_timeout: int = 60000,
+            claim_interval: int = 1800000,
+            *args,
+            **kwargs):
+        super().__init__(*args, **kwargs)
+        self.stream = stream
+        self.group = group
+        self.prefetch = prefetch
+        self.queue = Queue()
+
+        self.threads = []
+
+        self.client = useRedisStreamStore(
+            stream=stream,
+            group=group,
+            stream_max_entries=stream_max_entries,
+            redeliver_timeout=redeliver_timeout,
+            claim_interval=claim_interval,
+            **(params or {})
+        )
+
+    def _consume(self, *args, **kwargs):
+        def callback(message):
+            self.queue.put(message)
+
+        prefetch = kwargs.pop("prefetch", self.prefetch)
+        self.client.start_consuming(consumer=uuid.uuid4().hex, callback=callback, prefetch=prefetch, **kwargs)
+
+    def consume(self, *args, **kwargs):
+        daemon = kwargs.pop('daemon', True)
+        thread = threading.Thread(target=self._consume, args=args, kwargs=kwargs)
+        thread.daemon = daemon
+        thread.start()
+        self.threads.append(thread)
+        return RedisStreamConsumer(self)
+
+    def send(self, message: Any, *args, **kwargs):
+        """对消息进行预处理，然后再发送"""
+        if not isinstance(message, Message):
+            message = self.message_cls(body=message)
+
+        self.client.send({"_message": message.to_json()}, *args, **kwargs)
+
+    publish = send
+
+    def confirm(self, message: Message):
+        broker_msg = getattr(message, "message", None)
+        if broker_msg is not None:
+            self.client.ack(broker_msg)
+
+    def reject(self, message: Message):
+        broker_msg = getattr(message, "message", None)
+        if broker_msg is not None:
+            self.client.reject(broker_msg)
+
+    def requeue(self, message: Message, is_source=False):
+        """
+         重发消息：先拒绝 再 重入
+
+         :param message: 消息
+         :param is_source: 是否是原始消息消息，True: 使用原始消息重入当前队列，False: 使用消息的最新数据重入当前队列
+         """
+        self.reject(message)
+
+        broker_msg = getattr(message, "message", None)
+        if is_source and broker_msg is not None and hasattr(broker_msg, "body"):
+            self.client.send(broker_msg.body)
+        else:
+            self.send(message)
+
+
+class RedisStreamConsumer(BaseConsumer):
+    ...

onestep/broker/sqs/__init__.py

@@ -0,0 +1,4 @@
+from .sqs import SQSBroker, SQSConsumer
+from .sns import SNSBroker
+
+__all__ = ["SQSBroker", "SQSConsumer", "SNSBroker"]

onestep/broker/sqs/sns.py

@@ -0,0 +1,53 @@
+from typing import Optional, Dict, Any
+
+from onestep.broker import BaseBroker
+
+try:
+    from use_sqs import SNSPublisher
+except ImportError:
+    SNSPublisher = None
+
+
+class SNSBroker(BaseBroker):
+    """SNS消息Broker实现"""
+
+    def __init__(
+        self,
+        topic_arn: str,
+        message_group_id: Optional[str] = None,
+        params: Optional[Dict] = None,
+        *args,
+        **kwargs,
+    ):
+        """
+        初始化SNS Broker
+
+        :param topic_arn: SNS主题ARN
+        :param message_group_id: 消息组ID (FIFO主题必需)
+        :param params: SNS连接参数
+        """
+        if SNSPublisher is None:
+            raise ImportError("Please install the `use-sqs` module to use SNSBroker")
+
+        super().__init__(*args, **kwargs)
+        self.topic_arn = topic_arn
+        self.message_group_id = message_group_id
+
+        # 创建SNSPublisher实例
+        self.publisher = SNSPublisher(**(params or {}))
+
+    def publish(self, message: Any, **kwargs):
+        """发布消息"""
+        # 如果初始化时设置了message_group_id，且调用时未指定，则使用初始化的值
+        if self.message_group_id and "message_group_id" not in kwargs:
+            kwargs["message_group_id"] = self.message_group_id
+
+        self.publisher.publish(
+            self.topic_arn,
+            message=message,
+            **kwargs,
+        )
+
+    def consume(self, *args, **kwargs):
+        """消费消息 - SNS不支持直接消费"""
+        raise NotImplementedError("SNS Broker does not support consumption")

onestep/broker/sqs/sqs.py

@@ -0,0 +1,181 @@
+import json
+import threading
+from queue import Queue
+from typing import Optional, Dict, Any, Callable
+
+from onestep.broker import BaseBroker, BaseConsumer
+from onestep.message import Message
+try:
+    from use_sqs import SQSStore
+except ImportError:
+    SQSStore = None
+    boto3 = None
+
+
+
+class _SQSMessage(Message):
+    """SQS消息类"""
+
+    @classmethod
+    def from_broker(cls, broker_message: "boto3.resources.factory.sqs.Message"):
+        """从SQS消息创建Message对象"""
+        required_attrs = ("body", "delete", "message_id")
+        if not all(hasattr(broker_message, attr) for attr in required_attrs):
+            raise TypeError(
+                f"Message object missing required SQS attributes: {required_attrs}"
+            )
+
+        try:
+            # 如果body已经是dict，直接使用；如果是字符串，则解析JSON
+            if isinstance(broker_message.body, dict):
+                message = broker_message.body
+            else:
+                message = json.loads(broker_message.body)
+        except (json.JSONDecodeError, TypeError):
+            message = {"body": broker_message.body}
+
+        if not isinstance(message, dict):
+            message = {"body": message}
+        if "body" not in message:
+            message = {"body": message}
+
+        return cls(
+            body=message.get("body"), extra=message.get("extra"), message=broker_message
+        )
+
+
+class SQSBroker(BaseBroker):
+    """SQS消息队列Broker实现"""
+
+    message_cls = _SQSMessage
+
+    def __init__(
+        self,
+        queue_name: str,
+        message_group_id: str,
+        message_deduplication_id_func: Optional[Callable] = None,
+        params: Optional[Dict] = None,
+        prefetch: Optional[int] = 1,
+        auto_create: bool = True,
+        queue_params: Optional[Dict] = None,
+        *args,
+        **kwargs,
+    ):
+        """
+        初始化SQS Broker
+
+        :param queue_name: 队列名称
+        :param message_group_id: 消息组ID (FIFO队列必需)
+        :param message_deduplication_id_func: 接收 msg_body(json_str) 作为参数，返回 str 类型的 message_deduplication_id
+        :param params: SQS连接参数
+        :param prefetch: 预取消息数量
+        :param auto_create: 是否自动创建队列
+        :param queue_params: 队列参数
+        """
+        super().__init__(*args, **kwargs)
+        self.queue_name = queue_name
+        self.queue = Queue()
+        self.prefetch = prefetch
+        self.message_group_id = message_group_id
+        self.message_deduplication_id_func = message_deduplication_id_func
+        self.threads = []
+        self._shutdown = False
+        self._consuming_started = False
+        self._consume_lock = threading.Lock()
+
+        # 创建SQSStore实例
+        if SQSStore is None:
+            raise ImportError("Please install the `use-sqs` module to use SQSBroker")
+        self.store = SQSStore(**(params or {}))
+
+        # 确保队列存在
+        if auto_create:
+            self.store.declare_queue(queue_name, attributes=queue_params)
+
+    def _consume(self, *args, **kwargs):
+        """消费消息的内部方法"""
+        prefetch = kwargs.pop("prefetch", self.prefetch)
+
+        def callback(message):
+            """处理接收到的消息"""
+            # 直接将原始SQS消息放入队列，保留完整的消息引用
+            self.queue.put(message)
+
+        self.store.start_consuming(
+            self.queue_name, callback=callback, prefetch=prefetch, **kwargs
+        )
+
+    def consume(self, *args, **kwargs) -> "SQSConsumer":
+        """启动消费者"""
+        daemon = kwargs.pop("daemon", True)
+        timeout = kwargs.pop("timeout", 1000)
+        with self._consume_lock:
+            if not self._consuming_started:
+                thread_kwargs = kwargs.copy()
+                thread = threading.Thread(target=self._consume, args=args, kwargs=thread_kwargs)
+                thread.daemon = daemon
+                thread.start()
+                self.threads.append(thread)
+                self._consuming_started = True
+        return SQSConsumer(self, timeout=timeout)
+
+    def publish(self, message: Any, **kwargs):
+        """发布消息"""
+        if self.message_deduplication_id_func:
+            message_deduplication_id = self.message_deduplication_id_func(message)
+            if (
+                isinstance(message_deduplication_id, str)
+                and message_deduplication_id.strip()
+                and len(message_deduplication_id.strip()) <= 128
+            ):
+                kwargs["message_deduplication_id"] = message_deduplication_id
+
+        self.store.send(
+            self.queue_name,
+            message=message,
+            message_group_id=self.message_group_id,
+            **kwargs,
+        )
+
+    def confirm(self, message: Message):
+        """确认消息"""
+        message.message.delete()
+
+    def reject(self, message: Message):
+        """拒绝消息"""
+        pass  # 拒绝消息时不删除，等待下一次消费 # 以支持 DLQ 机制
+
+    def requeue(self, message: Message, is_source: bool = False):
+        """
+        重新入队消息
+
+        :param message: 消息对象
+        :param is_source: 是否使用原始消息
+        """
+        broker_msg = getattr(message, "message", None)
+        # 确认原始消息，重新入队处理后的消息，以实现将消息放在队列尾部并等待处理
+        if is_source and broker_msg is not None and hasattr(broker_msg, "body"):
+            broker_msg.delete()
+            self.store.send(
+                self.queue_name,
+                broker_msg.body,
+                message_group_id=self.message_group_id,
+            )
+        else:
+            broker_msg.delete()
+            self.store.send(
+                self.queue_name, message.body, message_group_id=self.message_group_id
+            )
+
+    def shutdown(self):
+        """关闭Broker"""
+        self._shutdown = True
+        self.store.shutdown()
+        for thread in self.threads:
+            thread.join()
+        self.queue = Queue()
+        self.threads.clear()
+        self._consuming_started = False
+
+
+class SQSConsumer(BaseConsumer): ...

onestep/broker/webhook.py

@@ -0,0 +1,84 @@
+import logging
+import threading
+import collections
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+from .memory import MemoryBroker, MemoryConsumer
+
+logger = logging.getLogger(__name__)
+
+Server = collections.namedtuple("Server", ["path", "queue"])
+
+
+class WebHookServer(BaseHTTPRequestHandler):
+    servers = collections.defaultdict(list)
+
+    def do_POST(self):
+        """
+        接收WebHook请求
+        """
+        server_paths = WebHookServer.servers.get(self.server.server_address, [])
+        for server in server_paths:
+            if self.path == server.path:
+                queue = server.queue
+                break
+        else:
+            return self.send_error(404)
+
+        content_len = int(self.headers.get('content-length', 0))
+        post_body = self.rfile.read(content_len).decode("utf-8")
+        queue.put_nowait(post_body)
+        self.send_response(200)
+        self.send_header('Content-type', 'application/json')
+        self.end_headers()
+        self.wfile.write(b'{ "status": "ok" }')
+
+
+class WebHookBroker(MemoryBroker):
+    _servers = {}
+
+    def __init__(self,
+                 path: str,
+                 host: str = "0.0.0.0",
+                 port: int = 8090,
+                 *args,
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+        self.host = host
+        self.port = port
+        self.path = path
+        self.threads = []
+
+    def _create_server(self):
+
+        if (self.host, self.port) not in self._servers:
+            hs = ThreadingHTTPServer(
+                (self.host, self.port),
+                WebHookServer
+            )
+            self._servers[(self.host, self.port)] = hs
+            # 只有在创建新服务器时才启动线程
+            thread = threading.Thread(target=hs.serve_forever)
+            thread.daemon = True
+            thread.start()
+            self.threads.append(thread)
+        else:
+            hs = self._servers[(self.host, self.port)]
+
+        WebHookServer.servers[(self.host, self.port)].append(Server(self.path, self.queue))
+
+    def consume(self, *args, **kwargs):
+        self._create_server()
+        logger.debug(f"WebHookBroker: {self.host}:{self.port}{self.path}")
+        return WebHookConsumer(self, *args, **kwargs)
+
+    def shutdown(self):
+        hs = self._servers.get((self.host, self.port))
+        if hs:
+            hs.shutdown()
+        for thread in self.threads:
+            thread.join()
+
+
+class WebHookConsumer(MemoryConsumer):
+    ...

onestep/cli.py

@@ -0,0 +1,80 @@
+import argparse
+import importlib
+import logging
+import sys
+from onestep import step, __version__
+
+LOGFORMAT = "[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s"
+
+
+def setup_logging():
+    # 设置全局日志级别为INFO，避免第三方库的DEBUG日志输出
+    logging.basicConfig(level=logging.INFO, format=LOGFORMAT, stream=sys.stdout)
+
+    # exclude amqpstorm logs
+    logging.getLogger("amqpstorm").setLevel(logging.CRITICAL)
+    
+    # 获取onestep的logger并设置为DEBUG级别以便调试
+    onestep_logger = logging.getLogger("onestep")
+    onestep_logger.setLevel(logging.DEBUG)
+    
+    return onestep_logger
+
+
+logger = setup_logging()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='run onestep'
+    )
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "step", nargs='?',
+        help="the run step",
+    )
+    parser.add_argument(
+        "--group", "-G", default=None,
+        help="the run group",
+        type=str
+    )
+    parser.add_argument(
+        "--print",
+        action="store_true",
+        help="enable printing")
+    parser.add_argument(
+        "--path", "-P", default=".", nargs="*", type=str,
+        help="the step import path (default: current running directory)"
+    )
+    group.add_argument(
+        "--cron",
+        help="the cron expression to test",
+        type=str
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    for path in args.path:
+        sys.path.insert(0, path)
+    if args.cron:
+        from croniter import croniter
+        from datetime import datetime
+        cron = croniter(args.cron, datetime.now())
+        for _ in range(10):
+            print(cron.get_next(datetime))
+        return
+
+    logger.info(f"OneStep {__version__} is start up.")
+    try:
+        importlib.import_module(args.step)
+        step.start(group=args.group, block=True, print_jobs=args.print)
+    except ModuleNotFoundError:
+        logger.error(f"Module `{args.step}` not found.")
+    except KeyboardInterrupt:
+        step.shutdown()
+
+
+if __name__ == '__main__':
+    sys.exit(main())