Coverage for faststream / rabbit / subscriber / usecase.py: 99%

1import asyncio (empty)

2import contextlib (empty)

3from collections.abc import AsyncIterator, Sequence (empty)

4from typing import TYPE_CHECKING, Any, Optional, cast (empty)

5 

6import anyio (empty)

7from typing_extensions import override (empty)

8 

9from faststream._internal.endpoint.subscriber import SubscriberUsecase (empty)

10from faststream._internal.endpoint.utils import process_msg (empty)

11from faststream.rabbit.parser import AioPikaParser (empty)

12from faststream.rabbit.publisher.fake import RabbitFakePublisher (empty)

13from faststream.rabbit.schemas import RabbitExchange (empty)

14from faststream.rabbit.schemas.constants import REPLY_TO_QUEUE_EXCHANGE_DELIMITER (empty)

15 

16if TYPE_CHECKING: (empty)

17 from aio_pika import IncomingMessage, RobustQueue 

18 

19 from faststream._internal.endpoint.publisher import PublisherProto 

20 from faststream._internal.endpoint.subscriber.call_item import CallsCollection 

21 from faststream._internal.endpoint.subscriber.specification import ( 

22 SubscriberSpecification, 

23 ) 

24 from faststream.message import StreamMessage 

25 from faststream.rabbit.configs import RabbitBrokerConfig 

26 from faststream.rabbit.message import RabbitMessage 

27 from faststream.rabbit.schemas import RabbitQueue 

28 

29 from .config import RabbitSubscriberConfig 

30 

31 

32class RabbitSubscriber(SubscriberUsecase["IncomingMessage"]): (empty)

33 """A class to handle logic for RabbitMQ message consumption.""" 

34 

35 _outer_config: "RabbitBrokerConfig" (empty)

36 

37 def __init__( (empty)

38 self, 

39 config: "RabbitSubscriberConfig", 

40 specification: "SubscriberSpecification[Any, Any]", 

41 calls: "CallsCollection[IncomingMessage]", 

42 ) -> None: 

43 parser = AioPikaParser(pattern=config.queue.path_regex) (empty)

44 config.decoder = parser.decode_message (empty)

45 config.parser = parser.parse_message (empty)

46 super().__init__( (empty)

47 config, 

48 specification=specification, 

49 calls=calls, 

50 ) 

51 

52 self.queue = config.queue (empty)

53 self.exchange = config.exchange (empty)

54 

55 self.consume_args = config.consume_args or {} (empty)

56 

57 self.__no_ack = config.ack_first (empty)

58 

59 self._consumer_tag: str | None = None (empty)

60 self._queue_obj: RobustQueue | None = None (empty)

61 self.channel = config.channel (empty)

62 

63 @property (empty)

64 def app_id(self) -> str | None: (empty)

65 return self._outer_config.app_id (empty)

66 

67 def routing(self) -> str: (empty)

68 return f"{self._outer_config.prefix}{self.queue.routing()}" (empty)

69 

70 @override (empty)

71 async def start(self) -> None: (empty)

72 """Starts the consumer for the RabbitMQ queue.""" 

73 await super().start() (empty)

74 

75 queue_to_bind = self.queue.add_prefix(self._outer_config.prefix) (empty)

76 

77 declarer = self._outer_config.declarer (empty)

78 

79 self._queue_obj = queue = await declarer.declare_queue( (empty)

80 queue_to_bind, 

81 channel=self.channel, 

82 ) 

83 

84 if ( 

85 self.exchange is not None 

86 and queue_to_bind.declare # queue just getted from RMQ 

87 and self.exchange.name # check Exchange is not default 

88 ): 

89 exchange = await declarer.declare_exchange( (empty)

90 self.exchange, 

91 channel=self.channel, 

92 ) 

93 

94 await queue.bind( (empty)

95 exchange, 

96 routing_key=queue_to_bind.routing(), 

97 arguments=queue_to_bind.bind_arguments, 

98 timeout=queue_to_bind.timeout, 

99 robust=self.queue.robust, 

100 ) 

101 

102 if self.calls: (empty)

103 self._consumer_tag = await self._queue_obj.consume( (empty)

104 # NOTE: aio-pika expects AbstractIncomingMessage, not IncomingMessage 

105 self.consume, # type: ignore[arg-type] 

106 no_ack=self.__no_ack, 

107 arguments=self.consume_args, 

108 ) 

109 

110 self._post_start() (empty)

111 

112 async def stop(self) -> None: (empty)

113 await super().stop() (empty)

114 

115 if self._queue_obj is not None: (empty)

116 if self._consumer_tag is not None: # pragma: no branch (empty)

117 if not self._queue_obj.channel.is_closed: (empty)

118 await self._queue_obj.cancel(self._consumer_tag) (empty)

119 self._consumer_tag = None (empty)

120 

121 self._queue_obj = None (empty)

122 

123 @override (empty)

124 async def get_one( (empty)

125 self, 

126 *, 

127 timeout: float = 5.0, 

128 no_ack: bool = True, 

129 ) -> "RabbitMessage | None": 

130 assert self._queue_obj, "You should start subscriber at first." (empty)

131 assert not self.calls, ( (empty)

132 "You can't use `get_one` method if subscriber has registered handlers." 

133 ) 

134 

135 sleep_interval = timeout / 10 (empty)

136 

137 raw_message: IncomingMessage | None = None (empty)

138 with ( 

139 contextlib.suppress(asyncio.exceptions.CancelledError), 

140 anyio.move_on_after(timeout), 

141 ): 

142 while ( # noqa: ASYNC110 (empty)

143 raw_message := await self._queue_obj.get( 

144 fail=False, 

145 no_ack=no_ack, 

146 timeout=timeout, 

147 ) 

148 ) is None: 

149 await anyio.sleep(sleep_interval) (empty)

150 

151 context = self._outer_config.fd_config.context (empty)

152 async_parser, async_decoder = self._get_parser_and_decoder() (empty)

153 

154 msg: RabbitMessage | None = await process_msg( # type: ignore[assignment] (empty)

155 msg=raw_message, 

156 middlewares=( 

157 m(raw_message, context=context) for m in self._broker_middlewares 

158 ), 

159 parser=async_parser, 

160 decoder=async_decoder, 

161 ) 

162 return msg (empty)

163 

164 @override (empty)

165 async def __aiter__(self) -> AsyncIterator["RabbitMessage"]: # type: ignore[override] (empty)

166 assert self._queue_obj, "You should start subscriber at first." (empty)

167 assert not self.calls, ( (empty)

168 "You can't use iterator method if subscriber has registered handlers." 

169 ) 

170 

171 context = self._outer_config.fd_config.context (empty)

172 async_parser, async_decoder = self._get_parser_and_decoder() (empty)

173 

174 async with self._queue_obj.iterator() as queue_iter: (empty)

175 async for raw_message in queue_iter: 175 ↛ exitline 175 didn't jump to the function exit(empty)

176 raw_message = cast("IncomingMessage", raw_message) (empty)

177 

178 msg: RabbitMessage = await process_msg( # type: ignore[assignment] (empty)

179 msg=raw_message, 

180 middlewares=( 

181 m(raw_message, context=context) for m in self._broker_middlewares 

182 ), 

183 parser=async_parser, 

184 decoder=async_decoder, 

185 ) 

186 yield msg (empty)

187 

188 def _make_response_publisher( (empty)

189 self, 

190 message: "StreamMessage[Any]", 

191 ) -> Sequence["PublisherProto"]: 

192 if REPLY_TO_QUEUE_EXCHANGE_DELIMITER in message.reply_to: (empty)

193 queue_name, exchange_name = message.reply_to.split( (empty)

194 REPLY_TO_QUEUE_EXCHANGE_DELIMITER, 2 

195 ) 

196 publisher = RabbitFakePublisher( (empty)

197 self._outer_config.producer, 

198 app_id=self.app_id, 

199 routing_key=queue_name, 

200 exchange=RabbitExchange.validate(exchange_name), 

201 ) 

202 else: 

203 publisher = RabbitFakePublisher( (empty)

204 self._outer_config.producer, 

205 app_id=self.app_id, 

206 routing_key=message.reply_to, 

207 exchange=RabbitExchange(), 

208 ) 

209 

210 return (publisher,) (empty)

211 

212 @staticmethod (empty)

213 def build_log_context( (empty)

214 message: Optional["StreamMessage[Any]"], 

215 queue: "RabbitQueue", 

216 exchange: Optional["RabbitExchange"] = None, 

217 ) -> dict[str, str]: 

218 return { (empty)

219 "queue": queue.name, 

220 "exchange": getattr(exchange, "name", ""), 

221 "message_id": getattr(message, "message_id", ""), 

222 } 

223 

224 def get_log_context( (empty)

225 self, 

226 message: Optional["StreamMessage[Any]"], 

227 ) -> dict[str, str]: 

228 return self.build_log_context( (empty)

229 message=message, 

230 queue=self.queue, 

231 exchange=self.exchange, 

232 )