redis stream broker support custom id

[vvanglro]

I have multiple streams and use middleware to record the processing status of tasks. I use the ID returned by publish, which is xadd, to distinguish each task. Since it is currently impossible to pass in a custom ID, different streams may generate the same ID when concurrent.

import json
import traceback
from datetime import datetime

import redis.asyncio as redis
from types import TracebackType
from typing import Any, Awaitable, Callable

from faststream import BaseMiddleware, PublishCommand, StreamMessage

from config.settings import settings


def _convert_bytes_to_str(obj):
    if isinstance(obj, bytes):
        try:
            return obj.decode('utf-8')
        except UnicodeDecodeError:
            return obj.decode('latin-1')
    elif isinstance(obj, dict):
        return {k.decode('utf-8') if isinstance(k, bytes) else k:
                _convert_bytes_to_str(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_convert_bytes_to_str(item) for item in obj]
    elif isinstance(obj, tuple):
        return tuple(_convert_bytes_to_str(item) for item in obj)
    return obj

class TaskStatusMiddleware(BaseMiddleware):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.redis_client = redis.Redis(
            host=settings.redis_host,
            port=settings.redis_port,
            db=settings.redis_db,
            password=settings.redis_password,
            decode_responses=True
        )
        self.status_key_prefix = f"task:status:"

    # Use this if you want to add logic when a message is received for the first time,
    # such as logging incoming messages, validating headers, or setting up the context.
    async def on_receive(self) -> Any:
        message_ids = self.msg.get("message_ids", [])
        msg_id = message_ids[0].decode() if message_ids else "unknown"  # '1767854208631-0'

        await self._update_status(msg_id, "processing", {
                "raw_data": self.msg
            })
        return await super().on_receive()

    # Use this if you want to wrap the entire message processing process,
    # such as implementing retry logic, circuit breakers, rate limiting, or authentication.
    async def consume_scope(
        self,
        call_next: Callable[[StreamMessage[Any]], Awaitable[Any]],
        msg: StreamMessage[Any],
    ) -> Any:
        raw_message = msg.raw_message
        message_ids = raw_message.get("message_ids", [])
        msg_id = message_ids[0].decode() if message_ids else "unknown"  # '1767854208631-0'
        try:
            raw_resp = await call_next(msg)
            await self._update_status(msg_id, "completed", {
                "raw_resp": raw_resp
            })
            return raw_resp
        except Exception as e:
            raise e

    # Use this if you want to customize outgoing messages before they are sent,
    # such as adding encryption, compression, or custom headers.
    async def publish_scope(
        self,
        call_next: Callable[[PublishCommand], Awaitable[Any]],
        cmd: PublishCommand,
    ) -> Any:
        msg_id = await call_next(cmd)
        await self._update_status(msg_id.decode(), "pending", {})
        return msg_id

    # Use this if you want to perform post-processing tasks after message handling has completed,
    # such as cleaning up, logging errors, collecting metrics, or committing transactions.
    async def after_processed(
        self,
        exc_type: type[BaseException] | None = None,
        exc_val: BaseException | None = None,
        exc_tb: TracebackType | None = None,
    ) -> bool | None:
        if exc_val:
            message_ids = self.msg.get("message_ids", [])
            msg_id = message_ids[0].decode() if message_ids else "unknown"  # '1767854208631-0'
            error_details = {
                "error_type": exc_type.__name__ if exc_type else None,
                "error_message": str(exc_val),
                "traceback": traceback.format_exc()
            }
            await self._update_status(msg_id, "failed", {
                "raw_error": error_details
            })

        return await super().after_processed(exc_type, exc_val, exc_tb)

    async def _update_status(self, task_id: str, status: str, details: dict):
        status_data = {
            "status": status,
            "updated_at": datetime.now().isoformat(),
            **details
        }
        await self.redis_client.setex(
            f"{self.status_key_prefix}{task_id}",
            3 * 24 * 60 * 60, 
            json.dumps(_convert_bytes_to_str(status_data), ensure_ascii=False)
        )
        return True