Python Redis Stream【生产者=》消费者模式】

import redis
import time
from typing import Dict, List, Tuple, Any, Optional

from config.model import settings
from pydantic import BaseModel


class StreamMessage(BaseModel):
    message_id: str
    message_data: Dict[str, Any]


class StreamMessages(BaseModel):
    stream_name: str
    messages: List[StreamMessage]


class RedisStreamManager:
    def __init__(self, stream_name: str, redis_url: str = "", timeout: int = 2, max_length: int = 100000):
        self.redis_url = redis_url or settings.OMS_REDIS_URL
        self.redis = self._connect_to_redis()
        self.stream_name = stream_name
        self.timeout = timeout
        self.max_length = max_length

        # 确保流存在并设置最大长度
        self.redis.xtrim(self.stream_name, maxlen=self.max_length, approximate=False)

    def _connect_to_redis(self):
        try:
            client = redis.StrictRedis.from_url(self.redis_url, decode_responses=True)
            client.ping()  # 测试连接是否成功
            return client
        except redis.ConnectionError as e:
            raise e

    def add(self, message_data: Dict[str, str]) -> str:
        """添加消息到流中"""
        return self.redis.xadd(self.stream_name, message_data, maxlen=self.max_length)

    def ensure_group(self, group_name: str):
        """确保消费者组存在，如果不存在则创建"""
        try:
            self.redis.xgroup_create(self.stream_name, group_name, id="0", mkstream=True)
        except redis.exceptions.ResponseError as e:
            # 消费者组可能已经存在，忽略异常
            if "BUSYGROUP Consumer Group name already exists" not in str(e):
                raise e

    def consume(self, group_name: str, consumer_name: str, count: int = 10) -> StreamMessages:
        """从流中消费消息"""
        self.ensure_group(group_name)  # 确保消费者组存在
        raw_messages = self.redis.xreadgroup(
            group_name, consumer_name, {self.stream_name: ">"}, count=count, block=5000
        )

        # 解析 Redis 流消息到 Pydantic 模型
        stream_messages = []
        for _, message_list in raw_messages:
            for message_id, message_data in message_list:
                stream_message = StreamMessage(message_id=message_id, message_data=dict(message_data))
                stream_messages.append(stream_message)

        return StreamMessages(stream_name=self.stream_name, messages=stream_messages)

    def ack(self, group_name: str, message_id: str):
        """确认消息已处理"""
        self.redis.xack(self.stream_name, group_name, message_id)

    def reassign(self, group_name: str, consumer_name: str, message_id: str):
        """根据消息ID使消息重新可分配"""
        self.redis.xclaim(self.stream_name, group_name, consumer_name, min_idle_time=self.timeout * 1000, id=message_id)

    def query_unconfirmed(self, group_name: str) -> StreamMessages:
        """查询未确认的消息"""
        # 查询未确认的消息
        pending = self.redis.xpending(self.stream_name, group_name)

        # 处理返回的未确认消息
        stream_messages = []
        min_id = pending.get("min")
        max_id = pending.get("max")

        # 通过 XRANGE 查询未确认消息的具体内容
        if min_id and max_id:
            message_list = self.redis.xrange(self.stream_name, min_id, max_id)

            for message_id, message_data in message_list:
                stream_message = StreamMessage(message_id=message_id, message_data=dict(message_data))
                stream_messages.append(stream_message)

        return StreamMessages(stream_name=self.stream_name, messages=stream_messages)

    def check_timeout(self, message_id: str) -> bool:
        """检查消息是否超时"""
        message_time = int(message_id.split("-")[0])
        return (time.time() - message_time / 1000) > self.timeout

    def get_all_timeout_messages(self, group_name: str) -> List[str]:
        """获取所有超时的未确认消息ID"""
        pending_messages = self.query_unconfirmed(group_name)
        timeout_messages = []
        for message in pending_messages.get("messages", []):
            message_id = message[0]
            if self.check_timeout(message_id):
                timeout_messages.append(message_id)
        return timeout_messages

    def handle_timeout(self, group_name: str, consumer_name: str):
        """处理超时的未确认消息"""
        timeout_messages = self.get_all_timeout_messages(group_name)
        for message_id in timeout_messages:
            # 重新分配超时的未确认消息
            self.reassign(group_name, consumer_name, message_id)

    def reassign_all_unconfirmed(self, group_name: str, consumer_name: str):
        """将所有未确认的消息恢复到可分配状态"""
        unconfirmed_messages = self.query_unconfirmed(group_name)
        for message in unconfirmed_messages.get("messages", []):
            message_id = message[0]  # 直接使用，不需要解码
            self.reassign(group_name, consumer_name, message_id)

    def query_all(self, start: str = "-", end: str = "+"):
        """查询流中的所有消息"""
        try:
            messages = self.redis.xrange(self.stream_name, min=start, max=end)
            return messages
        except redis.exceptions.ResponseError as e:
            print(f"Error querying all messages: {e}")
            return None

    def weixiaofei(self, group_name):
        return self.redis.xpending(self.stream_name, group_name)


def test4():
    stream_name = "my_stream"
    group_name = "my_group"
    consumer_name = "my_consumer"
    manager = RedisStreamManager(stream_name=stream_name)

    # 确保消费者组存在
    manager.ensure_group(group_name)

    # 示例：添加消息
    print("Adding messages...")
    for i in range(5):
        message_id = manager.add({"key": f"value_{i}"})
        print(f"Added message with ID: {message_id}")

    data = manager.query_all()
    print(f"\ndata:{data}")

    data = manager.consume(group_name, consumer_name, count=1)
    print(f"\n未消费的:{data.messages[0].message_id}")
    manager.ack(group_name, data.messages[0].message_id)

    data = manager.consume(group_name, consumer_name, count=1)
    print(f"\n消费未确认:{data}")

    data = manager.query_unconfirmed(group_name)
    print(f"\n未消费列表:{data}")


def test3():
    stream_name = "my_stream"
    group_name = "my_group"
    consumer_name = "my_consumer"
    manager = RedisStreamManager(stream_name=stream_name)

    # 确保消费者组存在
    manager.ensure_group(group_name)

    data = manager.query_unconfirmed(group_name)
    print(f"\n未消费的:{data}")


def test1():
    stream_name = "my_stream"
    group_name = "my_group"
    consumer_name = "my_consumer"
    manager = RedisStreamManager(stream_name=stream_name)

    # 确保消费者组存在
    manager.ensure_group(group_name)

    # 示例：添加消息
    print("Adding messages...")
    for i in range(5):
        message_id = manager.add({"key": f"value_{i}"})
        print(f"Added message with ID: {message_id}")

    data = manager.query_all()
    print(f"\ndata:{data}")

    # 示例：消费消息
    test_cus = False
    test_timeout = False
    print("\nConsuming messages...")
    messages = manager.consume(group_name, consumer_name)
    for stream, message_list in messages:
        for message_id, message_data in message_list:
            message_id = message_id  # 直接使用，不需要解码
            message_data = dict(message_data)
            print(f"Received message {message_id}: {message_data}")

            if not test_cus:
                test_cus = True
                continue
            elif not test_timeout:
                test_timeout = True
                # 使消息超时以模拟未确认
                time.sleep(2 * manager.timeout)

            # 确认消息
            manager.ack(group_name, message_id)
            print(f"Acknowledged message {message_id}")

    data = manager.query_all()
    print(f"\n所有数据:{data}")

    data = manager.consume(group_name, consumer_name)
    print(f"\n未消费的:{data}")

    data = manager.handle_timeout(group_name, consumer_name)
    print(f"\n超时的:{data}")

    data = manager.query_unconfirmed(group_name)
    print(f"\n待确认的:{data}")


def test2():
    stream_name = "my_stream"
    group_name = "my_group"
    consumer_name = "my_consumer"
    manager = RedisStreamManager(stream_name=stream_name)

    # 确保消费者组存在
    manager.ensure_group(group_name)

    # 示例：消费消息
    print("\nConsuming messages...")
    messages = manager.consume(group_name, consumer_name)
    for stream, message_list in messages:
        for message_id, message_data in message_list:
            message_id = message_id  # 直接使用，不需要解码
            message_data = dict(message_data)
            print(f"Received message {message_id}: {message_data}")

            # 使消息超时以模拟未确认
            time.sleep(2 * manager.timeout)  # 等待超时

            # 确认消息
            manager.ack(group_name, message_id)
            print(f"Acknowledged message {message_id}")

    # 处理超时消息
    print("\nHandling timeout messages...")
    manager.handle_timeout(group_name, consumer_name)
    print("Handled timeout messages")

    # 获取所有超时消息
    print("\nGetting all timeout messages...")
    timeout_messages = manager.get_all_timeout_messages(group_name)
    print(f"All timeout messages: {timeout_messages}")

    # 将所有未确认的消息恢复到可分配状态
    print("\nReassigning all unconfirmed messages...")
    manager.reassign_all_unconfirmed(group_name, consumer_name)
    print("Reassigned all unconfirmed messages")