Add broker.subscriber.get_one()

faststream/broker/subscriber/proto.py

@@ -85,6 +85,11 @@ async def consume(self, msg: MsgType) -> Any: ...
     @abstractmethod
     async def process_message(self, msg: MsgType) -> Any: ...
 
+    @abstractmethod
+    async def get_one(
+        self, *, timeout: float = 5.0
+    ) -> "Optional[StreamMessage[MsgType]]": ...
+
     @abstractmethod
     def add_call(
         self,

faststream/broker/subscriber/usecase.py

@@ -107,8 +107,8 @@ def __init__(
         """Initialize a new instance of the class."""
         self.calls = []
 
-        self._default_parser = default_parser
-        self._default_decoder = default_decoder
+        self._parser = default_parser
+        self._decoder = default_decoder
         self._no_reply = no_reply
         # Watcher args
         self._no_ack = no_ack
@@ -163,18 +163,17 @@ def setup(  # type: ignore[override]
 
         for call in self.calls:
             if parser := call.item_parser or broker_parser:
-                async_parser = resolve_custom_func(
-                    to_async(parser), self._default_parser
-                )
+                async_parser = resolve_custom_func(to_async(parser), self._parser)
             else:
-                async_parser = self._default_parser
+                async_parser = self._parser
 
             if decoder := call.item_decoder or broker_decoder:
-                async_decoder = resolve_custom_func(
-                    to_async(decoder), self._default_decoder
-                )
+                async_decoder = resolve_custom_func(to_async(decoder), self._decoder)
             else:
-                async_decoder = self._default_decoder
+                async_decoder = self._decoder
+
+            self._parser = async_parser
+            self._decoder = async_decoder
 
             call.setup(
                 parser=async_parser,

faststream/confluent/subscriber/usecase.py

@@ -1,5 +1,7 @@
 import asyncio
 from abc import ABC, abstractmethod
+from contextlib import AsyncExitStack
+from functools import partial
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -9,7 +11,7 @@
     List,
     Optional,
     Sequence,
-    Tuple,
+    Tuple, Awaitable,
 )
 
 import anyio
@@ -19,8 +21,10 @@
 from faststream.broker.publisher.fake import FakePublisher
 from faststream.broker.subscriber.usecase import SubscriberUsecase
 from faststream.broker.types import MsgType
+from faststream.confluent.message import KafkaMessage
 from faststream.confluent.parser import AsyncConfluentParser
 from faststream.confluent.schemas import TopicPartition
+from faststream.utils.functions import return_input
 
 if TYPE_CHECKING:
     from fast_depends.dependencies import Depends
@@ -152,6 +156,9 @@ async def start(self) -> None:
 
         await super().start()
 
+        if not self.calls:
+            return
+
         self.task = asyncio.create_task(self._consume())
 
     async def close(self) -> None:
@@ -166,6 +173,28 @@ async def close(self) -> None:
 
         self.task = None
 
+    async def get_one(self, timeout: float = 5.0) -> "Optional[KafkaMessage]":
+        assert self.consumer, "You should start subscriber at first."
+        assert (  # nosec B101
+            not self.calls
+        ), "You can't use `get_one` method if subscriber has registered handlers."
+
+        raw_message = await self.consumer.getone(timeout=timeout)
+
+        async with AsyncExitStack() as stack:
+            return_msg: Callable[[KafkaMessage], Awaitable[KafkaMessage]] = (
+                return_input
+            )
+
+            for m in self._broker_middlewares:
+                mid = m(raw_message)
+                await stack.enter_async_context(mid)
+                return_msg = partial(mid.consume_scope, return_msg)
+
+            parsed_msg = await self._parser(raw_message)
+            parsed_msg._decoded_body = await self._decoder(parsed_msg)
+            return await return_msg(parsed_msg)
+
     def _make_response_publisher(
         self,
         message: "StreamMessage[Any]",

faststream/kafka/subscriber/usecase.py

@@ -1,5 +1,7 @@
 import asyncio
 from abc import ABC, abstractmethod
+from contextlib import AsyncExitStack
+from functools import partial
 from itertools import chain
 from typing import (
     TYPE_CHECKING,
@@ -10,7 +12,7 @@
     List,
     Optional,
     Sequence,
-    Tuple,
+    Tuple, Awaitable,
 )
 
 import anyio
@@ -28,6 +30,7 @@
 )
 from faststream.kafka.message import KafkaAckableMessage, KafkaMessage
 from faststream.kafka.parser import AioKafkaBatchParser, AioKafkaParser
+from faststream.utils.functions import return_input
 from faststream.utils.path import compile_path
 
 if TYPE_CHECKING:
@@ -164,6 +167,9 @@ async def start(self) -> None:
         await consumer.start()
         await super().start()
 
+        if not self.calls:
+            return
+
         self.task = asyncio.create_task(self._consume())
 
     async def close(self) -> None:
@@ -178,6 +184,33 @@ async def close(self) -> None:
 
         self.task = None
 
+    async def get_one(self, timeout: float = 5) -> "Optional[KafkaMessage]":
+        assert self.consumer, "You should start subscriber at first."
+        assert (  # nosec B101
+            not self.calls
+        ), "You can't use `get_one` method if subscriber has registered handlers."
+
+        raw_messages = await self.consumer.getmany(timeout_ms=timeout * 1000, max_records=1)
+
+        if not raw_messages:
+            return None
+
+        (raw_message,) ,= raw_messages.values()
+
+        async with AsyncExitStack() as stack:
+            return_msg: Callable[[KafkaMessage], Awaitable[KafkaMessage]] = (
+                return_input
+            )
+
+            for m in self._broker_middlewares:
+                mid = m(raw_message)
+                await stack.enter_async_context(mid)
+                return_msg = partial(mid.consume_scope, return_msg)
+
+            parsed_msg = await self._parser(raw_message)
+            parsed_msg._decoded_body = await self._decoder(parsed_msg)
+            return await return_msg(parsed_msg)
+
     def _make_response_publisher(
         self,
         message: "StreamMessage[Any]",

faststream/rabbit/subscriber/usecase.py

@@ -1,6 +1,9 @@
+from contextlib import AsyncExitStack
+from functools import partial
 from typing import (
     TYPE_CHECKING,
     Any,
+    Awaitable,
     Callable,
     Dict,
     Iterable,
@@ -9,13 +12,16 @@
     Union,
 )
 
+import anyio
 from typing_extensions import override
 
 from faststream.broker.publisher.fake import FakePublisher
 from faststream.broker.subscriber.usecase import SubscriberUsecase
 from faststream.exceptions import SetupError
+from faststream.rabbit.helpers.declarer import RabbitDeclarer
 from faststream.rabbit.parser import AioPikaParser
 from faststream.rabbit.schemas import BaseRMQInformation
+from faststream.utils.functions import return_input
 
 if TYPE_CHECKING:
     from aio_pika import IncomingMessage, RobustQueue
@@ -24,6 +30,7 @@
     from faststream.broker.message import StreamMessage
     from faststream.broker.types import BrokerMiddleware, CustomCallable
     from faststream.rabbit.helpers.declarer import RabbitDeclarer
+    from faststream.rabbit.message import RabbitMessage
     from faststream.rabbit.publisher.producer import AioPikaFastProducer
     from faststream.rabbit.schemas import (
         RabbitExchange,
@@ -156,11 +163,12 @@ async def start(self) -> None:
                 robust=self.queue.robust,
             )
 
-        self._consumer_tag = await queue.consume(
-            # NOTE: aio-pika expects AbstractIncomingMessage, not IncomingMessage
-            self.consume,  # type: ignore[arg-type]
-            arguments=self.consume_args,
-        )
+        if self.calls:
+            self._consumer_tag = await self._queue_obj.consume(
+                # NOTE: aio-pika expects AbstractIncomingMessage, not IncomingMessage
+                self.consume,  # type: ignore[arg-type]
+                arguments=self.consume_args,
+            )
 
         await super().start()
 
@@ -175,6 +183,47 @@ async def close(self) -> None:
 
             self._queue_obj = None
 
+    async def get_one(
+        self,
+        *,
+        timeout: float = 5.0,
+        no_ack: bool = True,
+    ) -> "Optional[RabbitMessage]":
+        assert self._queue_obj, "You should start subscriber at first."  # nosec B101
+        assert (  # nosec B101
+            not self.calls
+        ), "You can't use `get_one` method if subscriber has registered handlers."
+
+        sleep_interval = timeout / 10
+
+        raw_message: Optional[IncomingMessage] = None
+        with anyio.move_on_after(timeout):
+            while (  # noqa: ASYNC110
+                raw_message := await self._queue_obj.get(
+                    fail=False,
+                    no_ack=no_ack,
+                    timeout=timeout,
+                )
+            ) is None:
+                await anyio.sleep(sleep_interval)
+
+        if raw_message is None:
+            return None
+
+        async with AsyncExitStack() as stack:
+            return_msg: Callable[[RabbitMessage], Awaitable[RabbitMessage]] = (
+                return_input
+            )
+
+            for m in self._broker_middlewares:
+                mid = m(raw_message)
+                await stack.enter_async_context(mid)
+                return_msg = partial(mid.consume_scope, return_msg)
+
+            parsed_msg = await self._parser(raw_message)
+            parsed_msg._decoded_body = await self._decoder(parsed_msg)
+            return await return_msg(parsed_msg)
+
     def _make_response_publisher(
         self,
         message: "StreamMessage[Any]",