Coverage for faststream / confluent / subscriber / usecase.py: 98%

1import logging (empty)

2from abc import abstractmethod (empty)

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

4from typing import ( (empty)

5 TYPE_CHECKING, 

6 Any, 

7 Optional, 

8 cast, 

9) 

10 

11import anyio (empty)

12from confluent_kafka import KafkaException, Message (empty)

13from typing_extensions import override (empty)

14 

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

16from faststream._internal.endpoint.subscriber.mixins import ConcurrentMixin, TasksMixin (empty)

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

18from faststream._internal.types import MsgType (empty)

19from faststream.confluent.parser import AsyncConfluentParser (empty)

20from faststream.confluent.publisher.fake import KafkaFakePublisher (empty)

21from faststream.confluent.schemas import TopicPartition (empty)

22 

23if TYPE_CHECKING: (empty)

24 from faststream._internal.endpoint.publisher import PublisherProto 

25 from faststream._internal.endpoint.subscriber import SubscriberSpecification 

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

27 from faststream.confluent.configs import KafkaBrokerConfig 

28 from faststream.confluent.helpers.client import AsyncConfluentConsumer 

29 from faststream.confluent.message import KafkaMessage 

30 from faststream.message import StreamMessage 

31 

32 from .config import KafkaSubscriberConfig 

33 

34 

35class LogicSubscriber(TasksMixin, SubscriberUsecase[MsgType]): (empty)

36 """A class to handle logic for consuming messages from Kafka.""" (empty)

37 

38 _outer_config: "KafkaBrokerConfig" (empty)

39 

40 group_id: str | None (empty)

41 

42 consumer: Optional["AsyncConfluentConsumer"] (empty)

43 parser: AsyncConfluentParser (empty)

44 

45 def __init__( (empty)

46 self, 

47 config: "KafkaSubscriberConfig", 

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

49 calls: "CallsCollection[MsgType]", 

50 ) -> None: 

51 super().__init__(config, specification, calls) (empty)

52 

53 self.__connection_data = config.connection_data (empty)

54 

55 self.group_id = config.group_id (empty)

56 

57 self._topics = config.topics (empty)

58 self._partitions = config.partitions (empty)

59 

60 self.consumer = None (empty)

61 self.polling_interval = config.polling_interval (empty)

62 

63 @property (empty)

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

65 return self._outer_config.client_id (empty)

66 

67 @property (empty)

68 def topics(self) -> list[str]: (empty)

69 return [f"{self._outer_config.prefix}{t}" for t in self._topics] (empty)

70 

71 @property (empty)

72 def partitions(self) -> list[TopicPartition]: (empty)

73 return [p.add_prefix(self._outer_config.prefix) for p in self._partitions] (empty)

74 

75 @override (empty)

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

77 """Start the consumer.""" 

78 await super().start() (empty)

79 self.consumer = consumer = self._outer_config.builder( (empty)

80 *self.topics, 

81 partitions=self.partitions, 

82 group_id=self.group_id, 

83 client_id=self.client_id, 

84 **self.__connection_data, 

85 ) 

86 self.parser._setup(consumer) (empty)

87 await consumer.start() (empty)

88 

89 self._post_start() (empty)

90 

91 if self.calls: (empty)

92 self.add_task(self._consume) (empty)

93 

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

95 await super().stop() (empty)

96 

97 if self.consumer is not None: (empty)

98 await self.consumer.stop() (empty)

99 self.consumer = None (empty)

100 

101 @override (empty)

102 async def get_one( (empty)

103 self, 

104 *, 

105 timeout: float = 5.0, 

106 ) -> "KafkaMessage | None": 

107 assert self.consumer, "You should start subscriber at first." (empty)

108 assert not self.calls, ( (empty)

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

110 ) 

111 

112 raw_message = await self.consumer.getone(timeout=timeout) (empty)

113 

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

115 

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

117 

118 return await process_msg( # type: ignore[return-value] (empty)

119 msg=raw_message, 

120 middlewares=( 

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

122 ), 

123 parser=async_parser, 

124 decoder=async_decoder, 

125 ) 

126 

127 @override (empty)

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

129 assert self.consumer, "You should start subscriber at first." (empty)

130 assert not self.calls, ( (empty)

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

132 ) 

133 

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

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

136 

137 timeout = 5.0 (empty)

138 while True: (empty)

139 raw_message = await self.consumer.getone(timeout=timeout) (empty)

140 

141 if raw_message is None: 141 ↛ 142line 141 didn't jump to line 142 because the condition on line 141 was never true(empty)

142 continue 

143 

144 yield cast( (empty)

145 "KafkaMessage", 

146 await process_msg( 

147 msg=raw_message, 

148 middlewares=( 

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

150 ), 

151 parser=async_parser, 

152 decoder=async_decoder, 

153 ), 

154 ) 

155 

156 def _make_response_publisher( (empty)

157 self, 

158 message: "StreamMessage[Any]", 

159 ) -> Sequence["PublisherProto"]: 

160 return ( (empty)

161 KafkaFakePublisher( 

162 self._outer_config.producer, 

163 topic=message.reply_to, 

164 ), 

165 ) 

166 

167 async def consume_one(self, msg: MsgType) -> None: (empty)

168 await self.consume(msg) (empty)

169 

170 @abstractmethod (empty)

171 async def get_msg(self) -> MsgType | None: (empty)

172 raise NotImplementedError 

173 

174 async def _consume(self) -> None: (empty)

175 assert self.consumer, "You should start subscriber at first." (empty)

176 

177 connected = True (empty)

178 while self.running: (empty)

179 try: (empty)

180 msg = await self.get_msg() (empty)

181 

182 except KafkaException as e: # pragma: no cover # noqa: PERF203 (empty)

183 self._log( 

184 logging.ERROR, 

185 message="Message fetch error", 

186 exc_info=e, 

187 ) 

188 

189 if connected: 

190 connected = False 

191 

192 await anyio.sleep(5) 

193 

194 else: 

195 if not connected: # pragma: no cover (empty)

196 connected = True 

197 

198 if msg is not None: (empty)

199 await self.consume_one(msg) (empty)

200 

201 @property (empty)

202 def topic_names(self) -> list[str]: (empty)

203 topics = self.topics or (f"{p.topic}-{p.partition}" for p in self.partitions) (empty)

204 return [f"{self._outer_config.prefix}{t}" for t in topics] (empty)

205 

206 @staticmethod (empty)

207 def build_log_context( (empty)

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

209 topic: str, 

210 group_id: str | None = None, 

211 ) -> dict[str, str]: 

212 return { (empty)

213 "topic": topic, 

214 "group_id": group_id or "", 

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

216 } 

217 

218 

219class DefaultSubscriber(LogicSubscriber[Message]): (empty)

220 def __init__( (empty)

221 self, 

222 config: "KafkaSubscriberConfig", 

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

224 calls: "CallsCollection[Message]", 

225 ) -> None: 

226 self.parser = AsyncConfluentParser(is_manual=not config.ack_first) (empty)

227 config.decoder = self.parser.decode_message (empty)

228 config.parser = self.parser.parse_message (empty)

229 super().__init__(config, specification, calls) (empty)

230 

231 async def get_msg(self) -> Optional["Message"]: (empty)

232 assert self.consumer, "You should setup subscriber at first." (empty)

233 return await self.consumer.getone(timeout=self.polling_interval) (empty)

234 

235 def get_log_context( (empty)

236 self, 

237 message: Optional["StreamMessage[Message]"], 

238 ) -> dict[str, str]: 

239 if message is None: (empty)

240 topic = ",".join(self.topic_names) (empty)

241 else: 

242 topic = message.raw_message.topic() or ",".join(self.topics) (empty)

243 

244 return self.build_log_context( (empty)

245 message=message, 

246 topic=topic, 

247 group_id=self.group_id, 

248 ) 

249 

250 

251class ConcurrentDefaultSubscriber(ConcurrentMixin["Message"], DefaultSubscriber): (empty)

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

253 await super().start() (empty)

254 self.start_consume_task() (empty)

255 

256 async def consume_one(self, msg: "Message") -> None: (empty)

257 await self._put_msg(msg) (empty)

258 

259 

260class BatchSubscriber(LogicSubscriber[tuple[Message, ...]]): (empty)

261 def __init__( (empty)

262 self, 

263 config: "KafkaSubscriberConfig", (empty)

264 specification: "SubscriberSpecification[Any, Any]", (empty)

265 calls: "CallsCollection[tuple[Message, ...]]", (empty)

266 max_records: int | None, (empty)

267 ) -> None: (empty)

268 self.parser = AsyncConfluentParser(is_manual=not config.ack_first) (empty)

269 config.decoder = self.parser.decode_batch (empty)

270 config.parser = self.parser.parse_batch (empty)

271 super().__init__(config, specification, calls) (empty)

272 

273 self.max_records = max_records (empty)

274 

275 async def get_msg(self) -> tuple["Message", ...] | None: (empty)

276 assert self.consumer, "You should setup subscriber at first." (empty)

277 return ( (empty)

278 await self.consumer.getmany( (empty)

279 timeout=self.polling_interval, (empty)

280 max_records=self.max_records, (empty)

281 ) 

282 or None (empty)

283 ) 

284 

285 def get_log_context( (empty)

286 self, 

287 message: Optional["StreamMessage[tuple[Message, ...]]"], (empty)

288 ) -> dict[str, str]: (empty)

289 if message is None: (empty)

290 topic = ",".join(self.topic_names) (empty)

291 else: 

292 topic = message.raw_message[0].topic() or ",".join(self.topic_names) (empty)

293 

294 return self.build_log_context( (empty)

295 message=message, (empty)

296 topic=topic, (empty)

297 group_id=self.group_id, (empty)

298 )