feat: [redis] add optional redis PEL reading support

faststream/redis/subscriber/usecases/stream_subscriber.py

@@ -1,6 +1,7 @@
 import math
-from collections.abc import AsyncIterator, Awaitable, Callable
-from typing import TYPE_CHECKING, Any, Optional, TypeAlias
+from collections.abc import AsyncIterator, Awaitable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Literal, Optional, TypeAlias
 
 from redis.exceptions import ResponseError
 from typing_extensions import override
@@ -21,6 +22,7 @@
 
 if TYPE_CHECKING:
     from anyio import Event
+    from redis.asyncio import Redis
 
     from faststream._internal.endpoint.subscriber import SubscriberSpecification
     from faststream._internal.endpoint.subscriber.call_item import (
@@ -34,8 +36,155 @@
 TopicName: TypeAlias = bytes
 Offset: TypeAlias = bytes
 
+RedisStreamReadResponse: TypeAlias = tuple[
+    tuple[
+        TopicName,
+        tuple[
+            tuple[
+                Offset,
+                dict[bytes, bytes],
+            ],
+            ...,
+        ],
+    ],
+    ...,
+]
+
+
+@dataclass(kw_only=True, repr=False)
+class RedisStreamReader:
+    stream_sub: "StreamSub"
+    resumable: Literal[False] = False
+
+    _last_id: str | None = field(default=None, init=False)
+    _client: Optional["Redis[bytes]"] = field(default=None, init=False)
+
+    async def __call__(self, last_id: str | None = None) -> RedisStreamReadResponse:
+        last_id = last_id or self._last_id or self.stream_sub.last_id
+
+        response = await self.read(last_id)
+        if response and response[0][1]:
+            self._last_id = response[0][1][-1][0].decode()
+            return response
+
+        return response
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}[{self.stream_sub.name}]"
+
+    @property
+    def client(self) -> "Redis[bytes]":
+        assert self._client, f"{self} has not been configured"
+        return self._client
+
+    async def configure(self, client: "Redis[bytes]", prefix: str) -> None:
+        self._client = client
+        self._last_id = self.stream_sub.last_id
+
+        self.stream_sub = self.stream_sub.add_prefix(prefix)
+
+    def read(self, last_id: str) -> Awaitable[RedisStreamReadResponse]:
+        stream = self.stream_sub
+
+        return self.client.xread(
+            {stream.name: last_id},
+            block=stream.polling_interval,
+            count=stream.max_records or 1,
+        )
+
+
+@dataclass(kw_only=True, repr=False)
+class RedisGroupReader(RedisStreamReader):
+    async def __call__(self, last_id: str | None = None) -> RedisStreamReadResponse:
+        last_id = last_id or self.stream_sub.last_id
+
+        response = await self.read(last_id)
+        if response and response[0][1]:
+            self._last_id = response[0][1][-1][0].decode()
+            return response
+
+        return response
+
+    @override
+    async def configure(self, client: "Redis[bytes]", prefix: str) -> None:
+        await super().configure(client, prefix)
+
+        stream = self.stream_sub
+
+        assert stream.group, f"{self} requires a stream with group"
+        assert stream.consumer, f"{self} requires a stream with consumer"
+
+        group_create_id = "$" if stream.last_id == ">" else stream.last_id
+        try:
+            await self.client.xgroup_create(
+                name=stream.name,
+                id=group_create_id,
+                groupname=stream.group,
+                mkstream=True,
+            )
+        except ResponseError as e:
+            if "already exists" not in str(e):
+                raise
+
+    def read(self, last_id: str) -> Awaitable[RedisStreamReadResponse]:
+        stream = self.stream_sub
+
+        assert stream.group, f"{self} requires a stream with group"
+        assert stream.consumer, f"{self} requires a stream with consumer"
+
+        return self.client.xreadgroup(
+            groupname=stream.group,
+            consumername=stream.consumer,
+            streams={stream.name: last_id},
+            count=stream.max_records or 1,
+            block=stream.polling_interval,
+            noack=stream.no_ack,
+        )
+
+
+@dataclass(kw_only=True, repr=False)
+class RedisConsumerReader(RedisGroupReader):
+    resumable: bool = True  # type: ignore[assignment] # redefinition
+    _resume_from: str | None = field(default=None, init=False)
+
+    @override
+    async def configure(self, client: "Redis[bytes]", prefix: str) -> None:
+        await super().configure(client, prefix)
+
+        if self.stream_sub.last_id != ">":
+            self.resume_from(self.stream_sub.last_id)
+            self.stream_sub.last_id = ">"
+
+    def resume_from(self, last_id: str, *, resumable: bool | None = None) -> None:
+        self._resume_from = last_id
+        if resumable is not None:
+            self.resumable = resumable
+
+    @override
+    async def __call__(self, last_id: str | None = None) -> RedisStreamReadResponse:
+        if last_id:
+            self._resume_from = last_id
+
+        resume_from = self._resume_from
+        resumable = self.resumable
+
+        if resume_from in {">", None}:
+            resumable = False
+
+        response = await super().__call__(resume_from)
+        if response and response[0][1]:
+            return response
+
+        if resumable:
+            self._resume_from = None
+            return await super().__call__(">")
+
+        return response
+
 
 class _StreamHandlerMixin(LogicSubscriber):
+    reader: RedisStreamReader | RedisGroupReader | RedisConsumerReader
+
     def __init__(
         self,
         config: "RedisSubscriberConfig",
@@ -45,7 +194,13 @@ def __init__(
         super().__init__(config, specification, calls)
 
         assert config.stream_sub
-        self._stream_sub = config.stream_sub
+        stream = self._stream_sub = config.stream_sub
+
+        if stream.consumer:
+            self.reader = RedisConsumerReader(stream_sub=stream)
+        else:
+            self.reader = RedisStreamReader(stream_sub=stream)
+
         self.last_id = config.stream_sub.last_id
 
     @property
@@ -79,92 +234,9 @@ async def start(self) -> None:
             group=self.stream_sub.group,
         )
 
-        stream = self.stream_sub
+        await self.reader.configure(client=client, prefix=self._outer_config.prefix)
 
-        read: Callable[
-            [str],
-            Awaitable[
-                tuple[
-                    tuple[
-                        TopicName,
-                        tuple[
-                            tuple[
-                                Offset,
-                                dict[bytes, bytes],
-                            ],
-                            ...,
-                        ],
-                    ],
-                    ...,
-                ],
-            ],
-        ]
-
-        if stream.group and stream.consumer:
-            group_create_id = "$" if self.last_id == ">" else self.last_id
-            try:
-                await client.xgroup_create(
-                    name=stream.name,
-                    id=group_create_id,
-                    groupname=stream.group,
-                    mkstream=True,
-                )
-            except ResponseError as e:
-                if "already exists" not in str(e):
-                    raise
-
-            def read(
-                _: str,
-            ) -> Awaitable[
-                tuple[
-                    tuple[
-                        TopicName,
-                        tuple[
-                            tuple[
-                                Offset,
-                                dict[bytes, bytes],
-                            ],
-                            ...,
-                        ],
-                    ],
-                    ...,
-                ],
-            ]:
-                return client.xreadgroup(
-                    groupname=stream.group,
-                    consumername=stream.consumer,
-                    streams={stream.name: stream.last_id},
-                    count=stream.max_records,
-                    block=stream.polling_interval,
-                    noack=stream.no_ack,
-                )
-
-        else:
-
-            def read(
-                last_id: str,
-            ) -> Awaitable[
-                tuple[
-                    tuple[
-                        TopicName,
-                        tuple[
-                            tuple[
-                                Offset,
-                                dict[bytes, bytes],
-                            ],
-                            ...,
-                        ],
-                    ],
-                    ...,
-                ],
-            ]:
-                return client.xread(
-                    {stream.name: last_id},
-                    block=stream.polling_interval,
-                    count=stream.max_records,
-                )
-
-        await super().start(read)
+        await super().start(self.reader)
 
     @override
     async def get_one(
@@ -262,31 +334,9 @@ def __init__(
         config.parser = parser.parse_message
         super().__init__(config, specification, calls)
 
-    async def _get_msgs(
-        self,
-        read: Callable[
-            [str],
-            Awaitable[
-                tuple[
-                    tuple[
-                        TopicName,
-                        tuple[
-                            tuple[
-                                Offset,
-                                dict[bytes, bytes],
-                            ],
-                            ...,
-                        ],
-                    ],
-                    ...,
-                ],
-            ],
-        ],
-    ) -> None:
-        for stream_name, msgs in await read(self.last_id):
+    async def _get_msgs(self, read: RedisStreamReader) -> None:
+        for stream_name, msgs in await read():
             if msgs:
-                self.last_id = msgs[-1][0].decode()
-
                 for message_id, raw_msg in msgs:
                     msg = DefaultStreamMessage(
                         type="stream",
@@ -310,19 +360,9 @@ def __init__(
         config.parser = parser.parse_message
         super().__init__(config, specification, calls)
 
-    async def _get_msgs(
-        self,
-        read: Callable[
-            [str],
-            Awaitable[
-                tuple[tuple[bytes, tuple[tuple[bytes, dict[bytes, bytes]], ...]], ...],
-            ],
-        ],
-    ) -> None:
-        for stream_name, msgs in await read(self.last_id):
+    async def _get_msgs(self, read: RedisStreamReader) -> None:
+        for stream_name, msgs in await read():
             if msgs:
-                self.last_id = msgs[-1][0].decode()
-
                 data: list[dict[bytes, bytes]] = []
                 ids: list[bytes] = []
                 for message_id, i in msgs: