Coverage for faststream / kafka / subscriber / usecase.py: 92%

1import logging (empty)

2from abc import abstractmethod (empty)

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

4from itertools import chain (empty)

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

6 

7import anyio (empty)

8from aiokafka import ConsumerRecord, TopicPartition (empty)

9from aiokafka.errors import ConsumerStoppedError, KafkaError, UnsupportedCodecError (empty)

10from typing_extensions import override (empty)

11 

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

13from faststream._internal.endpoint.subscriber.usecase import SubscriberUsecase (empty)

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

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

16from faststream._internal.utils.path import compile_path (empty)

17from faststream.kafka.helpers import make_logging_listener (empty)

18from faststream.kafka.message import KafkaAckableMessage, KafkaMessage, KafkaRawMessage (empty)

19from faststream.kafka.parser import AioKafkaBatchParser, AioKafkaParser (empty)

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

21 

22if TYPE_CHECKING: (empty)

23 from aiokafka import AIOKafkaConsumer 

24 

25 from faststream._internal.endpoint.publisher import PublisherProto 

26 from faststream._internal.endpoint.subscriber import SubscriberSpecification 

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

28 from faststream.kafka.configs import KafkaBrokerConfig 

29 from faststream.message import StreamMessage 

30 

31 from .config import KafkaSubscriberConfig 

32 

33 

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

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

36 

37 consumer: Optional["AIOKafkaConsumer"] (empty)

38 

39 batch: bool (empty)

40 parser: AioKafkaParser (empty)

41 

42 _outer_config: "KafkaBrokerConfig" (empty)

43 

44 def __init__( (empty)

45 self, 

46 config: "KafkaSubscriberConfig", 

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

48 calls: "CallsCollection[MsgType]", 

49 ) -> None: 

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

51 

52 self._topics = config.topics (empty)

53 self._partitions = config.partitions (empty)

54 self.group_id = config.group_id (empty)

55 

56 self._pattern = config.pattern (empty)

57 self._listener = config.listener (empty)

58 self._connection_args = config.connection_args (empty)

59 

60 self.consumer = None (empty)

61 

62 @property (empty)

63 def pattern(self) -> str | None: (empty)

64 if not self._pattern: (empty)

65 return self._pattern (empty)

66 return f"{self._outer_config.prefix}{self._pattern}" (empty)

67 

68 @property (empty)

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

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

71 

72 @property (empty)

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

74 return [ (empty)

75 TopicPartition( 

76 topic=f"{self._outer_config.prefix}{p.topic}", 

77 partition=p.partition, 

78 ) 

79 for p in self._partitions 

80 ] 

81 

82 @property (empty)

83 def builder(self) -> Callable[..., "AIOKafkaConsumer"]: (empty)

84 return self._outer_config.builder (empty)

85 

86 @property (empty)

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

88 return self._outer_config.client_id (empty)

89 

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

91 """Start the consumer.""" 

92 await super().start() (empty)

93 

94 self.consumer = consumer = self.builder( (empty)

95 group_id=self.group_id, 

96 client_id=self.client_id, 

97 **self._connection_args, 

98 ) 

99 

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

101 

102 if self.topics or self.pattern: (empty)

103 consumer.subscribe( (empty)

104 topics=self.topics, 

105 pattern=self.pattern, 

106 listener=make_logging_listener( 

107 consumer=consumer, 

108 logger=self._outer_config.logger.logger.logger, 

109 log_extra=self.get_log_context(None), 

110 listener=self._listener, 

111 ), 

112 ) 

113 

114 elif self.partitions: 114 ↛ 117line 114 didn't jump to line 117 because the condition on line 114 was always true(empty)

115 consumer.assign(partitions=self.partitions) (empty)

116 

117 await consumer.start() (empty)

118 

119 self._post_start() (empty)

120 

121 if self.calls: (empty)

122 self.add_task(self._run_consume_loop, (self.consumer,)) (empty)

123 

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

125 await super().stop() (empty)

126 

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

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

129 self.consumer = None (empty)

130 

131 @override (empty)

132 async def get_one( (empty)

133 self, 

134 *, 

135 timeout: float = 5.0, 

136 ) -> "KafkaMessage | None": 

137 assert not self.calls, ( (empty)

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

139 ) 

140 

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

142 

143 raw_messages = await self.consumer.getmany( (empty)

144 timeout_ms=timeout * 1000, 

145 max_records=1, 

146 ) 

147 

148 if not raw_messages: (empty)

149 return None (empty)

150 

151 ((raw_message,),) = raw_messages.values() (empty)

152 

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

154 

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

156 

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

158 msg=raw_message, 

159 middlewares=( 

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

161 ), 

162 parser=async_parser, 

163 decoder=async_decoder, 

164 ) 

165 return msg (empty)

166 

167 @override (empty)

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

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

170 assert not self.calls, ( (empty)

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

172 ) 

173 

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

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

176 

177 async for raw_message in self.consumer: 177 ↛ exitline 177 didn't return from function '__aiter__' because the loop on line 177 didn't complete(empty)

178 msg: KafkaMessage = 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 return ( (empty)

193 KafkaFakePublisher( 

194 self._outer_config.producer, 

195 topic=message.reply_to, 

196 ), 

197 ) 

198 

199 @abstractmethod (empty)

200 async def get_msg(self, consumer: "AIOKafkaConsumer") -> MsgType: (empty)

201 raise NotImplementedError 

202 

203 async def _run_consume_loop(self, consumer: "AIOKafkaConsumer") -> None: (empty)

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

205 

206 connected = True (empty)

207 while self.running: (empty)

208 try: (empty)

209 msg = await self.get_msg(consumer) (empty)

210 

211 except UnsupportedCodecError as e: # noqa: PERF203 (empty)

212 self._log( 

213 logging.ERROR, 

214 "There is no suitable compression library available. Please refer to the Kafka " 

215 "documentation for more information - " 

216 "https://aiokafka.readthedocs.io/en/stable/#installation", 

217 exc_info=e, 

218 ) 

219 await anyio.sleep(15) 

220 

221 except KafkaError as e: (empty)

222 self._log(logging.ERROR, "Kafka error occurred", exc_info=e) 

223 

224 if connected: 

225 connected = False 

226 

227 await anyio.sleep(5) 

228 

229 except ConsumerStoppedError: (empty)

230 return (empty)

231 

232 else: 

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

234 connected = True 

235 

236 if msg: 236 ↛ 207line 236 didn't jump to line 207 because the condition on line 236 was always true(empty)

237 await self.consume_one(msg) (empty)

238 

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

240 await self.consume(msg) (empty)

241 

242 @property (empty)

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

244 if self.pattern: (empty)

245 topics = [self.pattern] (empty)

246 

247 elif self.topics: (empty)

248 topics = self.topics (empty)

249 

250 else: 

251 topics = [f"{p.topic}-{p.partition}" for p in self.partitions] (empty)

252 

253 return topics (empty)

254 

255 @staticmethod (empty)

256 def build_log_context( (empty)

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

258 topic: str, 

259 group_id: str | None = None, 

260 ) -> dict[str, str]: 

261 return { (empty)

262 "topic": topic, 

263 "group_id": group_id or "", 

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

265 } 

266 

267 

268class DefaultSubscriber(LogicSubscriber["ConsumerRecord"]): (empty)

269 def __init__( (empty)

270 self, 

271 config: "KafkaSubscriberConfig", 

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

273 calls: "CallsCollection[ConsumerRecord]", 

274 ) -> None: 

275 if config.pattern: (empty)

276 reg, pattern = compile_path( (empty)

277 config.pattern, 

278 replace_symbol=".*", 

279 patch_regex=lambda x: x.replace(r"\*", ".*"), 

280 ) 

281 config.pattern = pattern (empty)

282 

283 else: 

284 reg = None (empty)

285 

286 self.parser = AioKafkaParser( (empty)

287 msg_class=KafkaMessage if config.ack_first else KafkaAckableMessage, 

288 regex=reg, 

289 ) 

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

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

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

293 

294 async def get_msg(self, consumer: "AIOKafkaConsumer") -> "ConsumerRecord": (empty)

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

296 return await consumer.getone() (empty)

297 

298 def get_log_context( (empty)

299 self, 

300 message: Optional["StreamMessage[ConsumerRecord]"], 

301 ) -> dict[str, str]: 

302 if message is None: (empty)

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

304 else: 

305 topic = message.raw_message.topic (empty)

306 

307 return self.build_log_context( (empty)

308 message=message, 

309 topic=topic, 

310 group_id=self.group_id, 

311 ) 

312 

313 

314class BatchSubscriber(LogicSubscriber[tuple["ConsumerRecord", ...]]): (empty)

315 def __init__( (empty)

316 self, 

317 config: "KafkaSubscriberConfig", (empty)

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

319 calls: "CallsCollection[tuple[ConsumerRecord, ...]]", (empty)

320 batch_timeout_ms: int, (empty)

321 max_records: int | None, (empty)

322 ) -> None: (empty)

323 if config.pattern: (empty)

324 reg, pattern = compile_path( 

325 config.pattern, 

326 replace_symbol=".*", 

327 patch_regex=lambda x: x.replace(r"\*", ".*"), 

328 ) 

329 config.pattern = pattern 

330 

331 else: 

332 reg = None (empty)

333 

334 self.parser = AioKafkaBatchParser( (empty)

335 msg_class=KafkaMessage if config.ack_first else KafkaAckableMessage, (empty)

336 regex=reg, (empty)

337 ) 

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

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

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

341 

342 self.batch_timeout_ms = batch_timeout_ms (empty)

343 self.max_records = max_records (empty)

344 

345 async def get_msg( (empty)

346 self, 

347 consumer: "AIOKafkaConsumer", (empty)

348 ) -> tuple["ConsumerRecord", ...]: (empty)

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

350 

351 messages = await consumer.getmany( (empty)

352 timeout_ms=self.batch_timeout_ms, (empty)

353 max_records=self.max_records, (empty)

354 ) 

355 

356 if not messages: # pragma: no cover (empty)

357 await anyio.sleep(self.batch_timeout_ms / 1000) 

358 return () 

359 

360 return tuple(chain(*messages.values())) (empty)

361 

362 def get_log_context( (empty)

363 self, 

364 message: Optional["StreamMessage[tuple[ConsumerRecord, ...]]"], (empty)

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

366 if message is None: (empty)

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

368 else: 

369 topic = message.raw_message[0].topic (empty)

370 

371 return self.build_log_context( (empty)

372 message=message, (empty)

373 topic=topic, (empty)

374 group_id=self.group_id, (empty)

375 ) 

376 

377 

378class ConcurrentDefaultSubscriber(ConcurrentMixin["ConsumerRecord"], DefaultSubscriber): (empty)

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

380 await super().start() (empty)

381 self.start_consume_task() (empty)

382 

383 async def consume_one(self, msg: "ConsumerRecord") -> None: (empty)

384 await self._put_msg(msg) (empty)

385 

386 

387class ConcurrentBetweenPartitionsSubscriber(DefaultSubscriber): (empty)

388 consumer_subgroup: list["AIOKafkaConsumer"] (empty)

389 

390 def __init__( (empty)

391 self, 

392 config: "KafkaSubscriberConfig", 

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

394 calls: "CallsCollection[ConsumerRecord]", 

395 max_workers: int, 

396 ) -> None: 

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

398 

399 self.max_workers = max_workers (empty)

400 self.consumer_subgroup = [] (empty)

401 

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

403 """Start the consumer subgroup.""" 

404 await super(LogicSubscriber, self).start() (empty)

405 

406 if self.calls: 406 ↛ 419line 406 didn't jump to line 419 because the condition on line 406 was always true(empty)

407 self.consumer_subgroup = [ (empty)

408 self.builder( 

409 group_id=self.group_id, 

410 client_id=self.client_id, 

411 **self._connection_args, 

412 ) 

413 for _ in range(self.max_workers) 

414 ] 

415 

416 else: 

417 # We should create single consumer to support 

418 # `get_one()` and `__aiter__` methods 

419 self.consumer = self.builder( 

420 group_id=self.group_id, 

421 client_id=self.client_id, 

422 **self._connection_args, 

423 ) 

424 self.consumer_subgroup = [self.consumer] 

425 

426 # Subscribers starting should be called concurrently 

427 # to balance them correctly 

428 async with anyio.create_task_group() as tg: (empty)

429 for c in self.consumer_subgroup: (empty)

430 c.subscribe( (empty)

431 topics=self.topics, 

432 listener=make_logging_listener( 

433 consumer=c, 

434 logger=self._outer_config.logger.logger.logger, 

435 log_extra=self.get_log_context(None), 

436 listener=self._listener, 

437 ), 

438 ) 

439 

440 tg.start_soon(c.start) (empty)

441 

442 self._post_start() (empty)

443 

444 if self.calls: 444 ↛ exitline 444 didn't return from function 'start' because the condition on line 444 was always true(empty)

445 for c in self.consumer_subgroup: (empty)

446 self.add_task(self._run_consume_loop, (c,)) (empty)

447 

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

449 if self.consumer_subgroup: 449 ↛ 456line 449 didn't jump to line 456 because the condition on line 449 was always true(empty)

450 async with anyio.create_task_group() as tg: (empty)

451 for consumer in self.consumer_subgroup: (empty)

452 tg.start_soon(consumer.stop) (empty)

453 

454 self.consumer_subgroup = [] (empty)

455 

456 await super().stop() (empty)

457 

458 async def get_msg(self, consumer: "AIOKafkaConsumer") -> "KafkaRawMessage": (empty)

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

460 message = await consumer.getone() (empty)

461 message.consumer = consumer (empty)

462 return cast("KafkaRawMessage", message) (empty)