Coverage for faststream / nats / publisher / producer.py: 94%

1import asyncio (empty)

2from abc import abstractmethod (empty)

3from typing import TYPE_CHECKING, Any, Optional (empty)

4 

5import anyio (empty)

6import nats (empty)

7from nats.aio.client import NO_RESPONDERS_STATUS (empty)

8from nats.js.api import Header (empty)

9from typing_extensions import override (empty)

10 

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

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

13from faststream.exceptions import FeatureNotSupportedException (empty)

14from faststream.message import encode_message (empty)

15from faststream.nats.helpers.state import ( (empty)

16 ConnectedState, 

17 ConnectionState, 

18 EmptyConnectionState, 

19) 

20from faststream.nats.parser import NatsParser (empty)

21from faststream.nats.response import NatsPublishCommand (empty)

22 

23if TYPE_CHECKING: (empty)

24 from fast_depends.library.serializer import SerializerProto 

25 from nats.aio.client import Client 

26 from nats.aio.msg import Msg 

27 from nats.js import JetStreamContext 

28 

29 from faststream._internal.types import ( 

30 AsyncCallable, 

31 CustomCallable, 

32 ) 

33 from faststream.nats.schemas import PubAck 

34 

35 

36class NatsFastProducer(ProducerProto[NatsPublishCommand]): (empty)

37 def connect( (empty)

38 self, 

39 connection: Any, (empty)

40 serializer: Optional["SerializerProto"], (empty)

41 ) -> None: ... (empty)

42 

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

44 

45 @abstractmethod (empty)

46 async def publish(self, cmd: "NatsPublishCommand") -> Optional["PubAck"]: ... (empty)

47 

48 @abstractmethod (empty)

49 async def request(self, cmd: "NatsPublishCommand") -> "Msg": ... (empty)

50 

51 async def publish_batch(self, cmd: "NatsPublishCommand") -> None: (empty)

52 msg = "NATS doesn't support publishing in batches." 

53 raise FeatureNotSupportedException(msg) 

54 

55 

56class NatsFastProducerImpl(NatsFastProducer): (empty)

57 """A class to represent a NATS producer.""" 

58 

59 _decoder: "AsyncCallable" (empty)

60 _parser: "AsyncCallable" (empty)

61 

62 def __init__( (empty)

63 self, 

64 parser: Optional["CustomCallable"], 

65 decoder: Optional["CustomCallable"], 

66 ) -> None: 

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

68 

69 default = NatsParser(pattern="", is_ack_disabled=True) (empty)

70 self._parser = ParserComposition(parser, default.parse_message) (empty)

71 self._decoder = ParserComposition(decoder, default.decode_message) (empty)

72 

73 self.__state: ConnectionState[Client] = EmptyConnectionState() (empty)

74 

75 def connect( (empty)

76 self, 

77 connection: "Client", 

78 serializer: Optional["SerializerProto"], 

79 ) -> None: 

80 self.serializer = serializer (empty)

81 self.__state = ConnectedState(connection) (empty)

82 

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

84 self.__state = EmptyConnectionState() (empty)

85 

86 @override (empty)

87 async def publish(self, cmd: "NatsPublishCommand") -> None: (empty)

88 payload, content_type = encode_message(cmd.body, self.serializer) (empty)

89 

90 headers_to_send = { (empty)

91 "content-type": content_type or "", 

92 **cmd.headers_to_publish(), 

93 } 

94 

95 return await self.__state.connection.publish( (empty)

96 subject=cmd.destination, 

97 payload=payload, 

98 reply=cmd.reply_to, 

99 headers=headers_to_send, 

100 ) 

101 

102 @override (empty)

103 async def request(self, cmd: "NatsPublishCommand") -> "Msg": (empty)

104 payload, content_type = encode_message(cmd.body, self.serializer) (empty)

105 

106 headers_to_send = { (empty)

107 "content-type": content_type or "", 

108 **cmd.headers_to_publish(), 

109 } 

110 

111 return await self.__state.connection.request( (empty)

112 subject=cmd.destination, 

113 payload=payload, 

114 headers=headers_to_send, 

115 timeout=cmd.timeout, 

116 ) 

117 

118 

119class NatsJSFastProducer(NatsFastProducer): (empty)

120 """A class to represent a NATS JetStream producer.""" 

121 

122 _decoder: "AsyncCallable" (empty)

123 _parser: "AsyncCallable" (empty)

124 

125 def __init__( (empty)

126 self, 

127 *, 

128 parser: Optional["CustomCallable"], 

129 decoder: Optional["CustomCallable"], 

130 ) -> None: 

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

132 

133 default = NatsParser(pattern="", is_ack_disabled=True) (empty)

134 self._parser = ParserComposition(parser, default.parse_message) (empty)

135 self._decoder = ParserComposition(decoder, default.decode_message) (empty)

136 

137 self.__state: ConnectionState[JetStreamContext] = EmptyConnectionState() (empty)

138 

139 def connect( (empty)

140 self, 

141 connection: "JetStreamContext", 

142 serializer: Optional["SerializerProto"], 

143 ) -> None: 

144 self.serializer = serializer (empty)

145 self.__state = ConnectedState(connection) (empty)

146 

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

148 self.__state = EmptyConnectionState() (empty)

149 

150 @override (empty)

151 async def publish(self, cmd: "NatsPublishCommand") -> "PubAck": (empty)

152 payload, content_type = encode_message(cmd.body, self.serializer) (empty)

153 

154 headers_to_send = { (empty)

155 "content-type": content_type or "", 

156 **cmd.headers_to_publish(js=True), 

157 } 

158 

159 return await self.__state.connection.publish( (empty)

160 subject=cmd.destination, 

161 payload=payload, 

162 headers=headers_to_send, 

163 stream=cmd.stream, 

164 timeout=cmd.timeout, 

165 ) 

166 

167 @override (empty)

168 async def request(self, cmd: "NatsPublishCommand") -> "Msg": (empty)

169 payload, content_type = encode_message(cmd.body, self.serializer) (empty)

170 

171 reply_to = self.__state.connection._nc.new_inbox() (empty)

172 future: asyncio.Future[Msg] = asyncio.Future() (empty)

173 sub = await self.__state.connection._nc.subscribe( (empty)

174 reply_to, 

175 future=future, 

176 max_msgs=1, 

177 ) 

178 await sub.unsubscribe(limit=1) (empty)

179 

180 headers_to_send = { (empty)

181 "content-type": content_type or "", 

182 "reply_to": reply_to, 

183 **cmd.headers_to_publish(js=False), 

184 } 

185 

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

187 await self.__state.connection.publish( (empty)

188 subject=cmd.destination, 

189 payload=payload, 

190 headers=headers_to_send, 

191 stream=cmd.stream, 

192 timeout=cmd.timeout, 

193 ) 

194 

195 msg = await future (empty)

196 

197 if ( # pragma: no cover 

198 msg.headers and (msg.headers.get(Header.STATUS) == NO_RESPONDERS_STATUS) 

199 ): 

200 raise nats.errors.NoRespondersError 

201 

202 return msg (empty)

203 

204 

205class FakeNatsFastProducer(NatsFastProducer): (empty)

206 def connect(self, connection: Any, serializer: Optional["SerializerProto"]) -> None: (empty)

207 raise NotImplementedError 

208 

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

210 raise NotImplementedError 

211 

212 @override (empty)

213 async def publish(self, cmd: "NatsPublishCommand") -> None: (empty)

214 raise NotImplementedError 

215 

216 @override (empty)

217 async def request(self, cmd: "NatsPublishCommand") -> "Msg": (empty)

218 raise NotImplementedError 

219 

220 @override (empty)

221 async def publish_batch(self, cmd: "NatsPublishCommand") -> None: (empty)

222 msg = "NATS doesn't support publishing in batches." 

223 raise FeatureNotSupportedException(msg)