Coverage for faststream / rabbit / publisher / producer.py: 92%

1from abc import abstractmethod (empty)

2from typing import ( (empty)

3 TYPE_CHECKING, 

4 Optional, 

5 Protocol, 

6 cast, 

7) 

8 

9import anyio (empty)

10from typing_extensions import Unpack, override (empty)

11 

12from faststream._internal.endpoint.utils import ParserComposition (empty)

13from faststream._internal.producer import ProducerProto (empty)

14from faststream.exceptions import FeatureNotSupportedException, IncorrectState (empty)

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

16from faststream.rabbit.response import RabbitPublishCommand (empty)

17from faststream.rabbit.schemas import RABBIT_REPLY, RabbitExchange (empty)

18 

19if TYPE_CHECKING: (empty)

20 from types import TracebackType 

21 

22 import aiormq 

23 from aio_pika import IncomingMessage, RobustQueue 

24 from aio_pika.abc import AbstractIncomingMessage, TimeoutType 

25 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream 

26 from fast_depends.library.serializer import SerializerProto 

27 

28 from faststream._internal.types import ( 

29 AsyncCallable, 

30 CustomCallable, 

31 ) 

32 from faststream.rabbit.helpers import RabbitDeclarer 

33 from faststream.rabbit.types import AioPikaSendableMessage 

34 

35 from .options import MessageOptions 

36 

37 

38class LockState(Protocol): (empty)

39 @property (empty)

40 def lock(self) -> "anyio.Lock": ... (empty)

41 

42 

43class LockUnset: (empty)

44 __slots__ = () (empty)

45 

46 @property (empty)

47 def lock(self) -> "anyio.Lock": (empty)

48 msg = "You should call `producer.connect()` method at first." 

49 raise IncorrectState(msg) 

50 

51 

52class RealLock: (empty)

53 __slots__ = ("lock",) (empty)

54 

55 def __init__(self) -> None: (empty)

56 self.lock = anyio.Lock() (empty)

57 

58 

59class AioPikaFastProducer(ProducerProto[RabbitPublishCommand]): (empty)

60 def connect(self, serializer: Optional["SerializerProto"] = None) -> None: ... (empty)

61 

62 def disconnect(self) -> None: ... (empty)

63 

64 @abstractmethod (empty)

65 async def publish( (empty)

66 self, 

67 cmd: "RabbitPublishCommand", (empty)

68 ) -> Optional["aiormq.abc.ConfirmationFrameType"]: ... (empty)

69 

70 @abstractmethod (empty)

71 async def request(self, cmd: "RabbitPublishCommand") -> "IncomingMessage": ... (empty)

72 

73 @override (empty)

74 async def publish_batch(self, cmd: "RabbitPublishCommand") -> None: (empty)

75 msg = "RabbitMQ doesn't support publishing in batches." 

76 raise FeatureNotSupportedException(msg) 

77 

78 

79class FakeAioPikaFastProducer(AioPikaFastProducer): (empty)

80 def __bool__(self) -> bool: (empty)

81 return False 

82 

83 def connect(self, serializer: Optional["SerializerProto"] = None) -> None: (empty)

84 raise NotImplementedError 

85 

86 def disconnect(self) -> None: (empty)

87 raise NotImplementedError 

88 

89 @override (empty)

90 async def publish( (empty)

91 self, 

92 cmd: "RabbitPublishCommand", 

93 ) -> Optional["aiormq.abc.ConfirmationFrameType"]: 

94 raise NotImplementedError 

95 

96 @override (empty)

97 async def request(self, cmd: "RabbitPublishCommand") -> "IncomingMessage": (empty)

98 raise NotImplementedError 

99 

100 

101class AioPikaFastProducerImpl(AioPikaFastProducer): (empty)

102 """A class for fast producing messages using aio-pika.""" 

103 

104 _decoder: "AsyncCallable" (empty)

105 _parser: "AsyncCallable" (empty)

106 

107 def __init__( (empty)

108 self, 

109 *, 

110 declarer: "RabbitDeclarer", 

111 parser: Optional["CustomCallable"], 

112 decoder: Optional["CustomCallable"], 

113 ) -> None: 

114 self.declarer = declarer (empty)

115 

116 self.__lock: LockState = LockUnset() (empty)

117 self.serializer: SerializerProto | None = None (empty)

118 

119 default_parser = AioPikaParser() (empty)

120 self._parser = ParserComposition(parser, default_parser.parse_message) (empty)

121 self._decoder = ParserComposition(decoder, default_parser.decode_message) (empty)

122 

123 def connect(self, serializer: Optional["SerializerProto"] = None) -> None: (empty)

124 """Lock initialization. 

125 

126 Should be called in async context due `anyio.Lock` object can't be created outside event loop. 

127 """ 

128 self.serializer = serializer (empty)

129 self.__lock = RealLock() (empty)

130 

131 def disconnect(self) -> None: (empty)

132 self.__lock = LockUnset() (empty)

133 

134 @override (empty)

135 async def publish( (empty)

136 self, 

137 cmd: "RabbitPublishCommand", 

138 ) -> Optional["aiormq.abc.ConfirmationFrameType"]: 

139 return await self._publish( (empty)

140 message=cmd.body, 

141 exchange=cmd.exchange, 

142 routing_key=cmd.destination, 

143 reply_to=cmd.reply_to, 

144 headers=cmd.headers, 

145 correlation_id=cmd.correlation_id, 

146 **cmd.publish_options, 

147 **cmd.message_options, 

148 ) 

149 

150 @override (empty)

151 async def request(self, cmd: "RabbitPublishCommand") -> "IncomingMessage": (empty)

152 async with _RPCCallback( (empty)

153 self.__lock.lock, 

154 await self.declarer.declare_queue(RABBIT_REPLY), 

155 ) as response_queue: 

156 with anyio.fail_after(cmd.timeout): (empty)

157 await self._publish( (empty)

158 message=cmd.body, 

159 exchange=cmd.exchange, 

160 routing_key=cmd.destination, 

161 reply_to=RABBIT_REPLY.name, 

162 headers=cmd.headers, 

163 correlation_id=cmd.correlation_id, 

164 **cmd.publish_options, 

165 **cmd.message_options, 

166 ) 

167 return await response_queue.receive() (empty)

168 

169 async def _publish( (empty)

170 self, 

171 message: "AioPikaSendableMessage", 

172 *, 

173 exchange: "RabbitExchange", 

174 routing_key: str, 

175 mandatory: bool = True, 

176 immediate: bool = False, 

177 timeout: "TimeoutType" = None, 

178 **message_options: Unpack["MessageOptions"], 

179 ) -> Optional["aiormq.abc.ConfirmationFrameType"]: 

180 message = AioPikaParser.encode_message( (empty)

181 message=message, serializer=self.serializer, **message_options 

182 ) 

183 

184 exchange_obj = await self.declarer.declare_exchange( (empty)

185 exchange=exchange, 

186 declare=False, 

187 ) 

188 

189 return await exchange_obj.publish( (empty)

190 message=message, 

191 routing_key=routing_key, 

192 mandatory=mandatory, 

193 immediate=immediate, 

194 timeout=timeout, 

195 ) 

196 

197 

198class _RPCCallback: (empty)

199 """A class provides an RPC lock.""" 

200 

201 def __init__(self, lock: "anyio.Lock", callback_queue: "RobustQueue") -> None: (empty)

202 self.lock = lock (empty)

203 self.queue = callback_queue (empty)

204 

205 async def __aenter__(self) -> "MemoryObjectReceiveStream[IncomingMessage]": (empty)

206 send_response_stream: MemoryObjectSendStream[AbstractIncomingMessage] 

207 receive_response_stream: MemoryObjectReceiveStream[AbstractIncomingMessage] 

208 

209 ( (empty)

210 send_response_stream, 

211 receive_response_stream, 

212 ) = anyio.create_memory_object_stream(max_buffer_size=1) 

213 await self.lock.acquire() (empty)

214 

215 self.consumer_tag = await self.queue.consume( (empty)

216 callback=send_response_stream.send, 

217 no_ack=True, 

218 ) 

219 

220 return cast( (empty)

221 "MemoryObjectReceiveStream[IncomingMessage]", 

222 receive_response_stream, 

223 ) 

224 

225 async def __aexit__( (empty)

226 self, 

227 exc_type: type[BaseException] | None = None, 

228 exc_val: BaseException | None = None, 

229 exc_tb: Optional["TracebackType"] = None, 

230 ) -> None: 

231 self.lock.release() (empty)

232 await self.queue.cancel(self.consumer_tag) (empty)