Coverage for faststream / confluent / helpers / client.py: 82%

1import asyncio (empty)

2import logging (empty)

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

4from concurrent.futures import ThreadPoolExecutor (empty)

5from contextlib import suppress (empty)

6from time import time (empty)

7from typing import TYPE_CHECKING, Any, cast (empty)

8 

9import anyio (empty)

10from confluent_kafka import Consumer, KafkaError, KafkaException, Message, Producer (empty)

11 

12from faststream._internal.utils.functions import call_or_await, run_in_executor (empty)

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

14from faststream.exceptions import SetupError (empty)

15 

16from . import config as config_module (empty)

17 

18if TYPE_CHECKING: (empty)

19 from typing_extensions import NotRequired, TypedDict 

20 

21 from faststream._internal.logger import LoggerState 

22 

23 from .admin import AdminService 

24 

25 class _SendKwargs(TypedDict): 

26 value: bytes | str | None 

27 key: bytes | str | None 

28 headers: list[tuple[str, str | bytes]] | None 

29 partition: NotRequired[int] 

30 timestamp: NotRequired[int] 

31 on_delivery: NotRequired[Callable[..., None]] 

32 

33 

34class AsyncConfluentProducer: (empty)

35 """An asynchronous Python Kafka client using the "confluent-kafka" package.""" 

36 

37 def __init__( (empty)

38 self, 

39 *, 

40 logger: "LoggerState", 

41 config: config_module.ConfluentFastConfig, 

42 ) -> None: 

43 self.logger_state = logger (empty)

44 

45 self.config = config.producer_config (empty)

46 self.producer = Producer( (empty)

47 self.config, 

48 logger=self.logger_state.logger.logger, 

49 ) 

50 

51 self.__running = True (empty)

52 self._poll_task = asyncio.create_task(self._poll_loop()) (empty)

53 

54 async def _poll_loop(self) -> None: (empty)

55 while self.__running: 55 ↛ exitline 55 didn't return from function '_poll_loop' because the condition on line 55 was always true(empty)

56 with suppress(Exception): (empty)

57 await call_or_await(self.producer.poll, 0.1) (empty)

58 

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

60 """Stop the Kafka producer and flush remaining messages.""" 

61 if self.__running: 61 ↛ exitline 61 didn't return from function 'stop' because the condition on line 61 was always true(empty)

62 self.__running = False (empty)

63 if not self._poll_task.done(): 63 ↛ 65line 63 didn't jump to line 65 because the condition on line 63 was always true(empty)

64 self._poll_task.cancel() (empty)

65 await call_or_await(self.producer.flush) (empty)

66 

67 async def flush(self) -> None: (empty)

68 await call_or_await(self.producer.flush) 

69 

70 async def send( (empty)

71 self, 

72 topic: str, 

73 value: bytes | str | None = None, 

74 key: bytes | str | None = None, 

75 partition: int | None = None, 

76 timestamp_ms: int | None = None, 

77 headers: list[tuple[str, str | bytes]] | None = None, 

78 no_confirm: bool = False, 

79 ) -> "asyncio.Future[Message | None] | Message | None": 

80 """Sends a single message to a Kafka topic.""" 

81 kwargs: _SendKwargs = { (empty)

82 "value": value, 

83 "key": key, 

84 "headers": headers, 

85 } 

86 

87 if partition is not None: 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true(empty)

88 kwargs["partition"] = partition 

89 

90 if timestamp_ms is not None: (empty)

91 kwargs["timestamp"] = timestamp_ms (empty)

92 

93 loop = asyncio.get_running_loop() (empty)

94 result_future: asyncio.Future[Message | None] = loop.create_future() (empty)

95 

96 def ack_callback(err: Any, msg: Message | None) -> None: (empty)

97 if err or (msg is not None and (err := msg.error())): 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true(empty)

98 loop.call_soon_threadsafe( 

99 result_future.set_exception, 

100 KafkaException(err), 

101 ) 

102 else: 

103 loop.call_soon_threadsafe(result_future.set_result, msg) (empty)

104 

105 kwargs["on_delivery"] = ack_callback (empty)

106 

107 # should be sync to prevent segfault 

108 # confluent stub expects bytes|None for value/key; we accept str and encode 

109 produce_value: bytes | None = ( (empty)

110 kwargs["value"] 

111 if isinstance(kwargs["value"], (bytes, type(None))) 

112 else kwargs["value"].encode() 

113 ) 

114 produce_key: bytes | None = ( (empty)

115 kwargs["key"] 

116 if isinstance(kwargs["key"], (bytes, type(None))) 

117 else (kwargs["key"].encode() if kwargs["key"] is not None else None) 

118 ) 

119 produce_headers: ( (empty)

120 dict[str, str | bytes | None] | list[tuple[str, str | bytes | None]] | None 

121 ) = cast("Any", kwargs["headers"]) if kwargs.get("headers") is not None else None 

122 produce_kwargs: dict[str, Any] = { (empty)

123 "value": produce_value, 

124 "key": produce_key, 

125 "headers": produce_headers, 

126 "on_delivery": kwargs["on_delivery"], 

127 } 

128 if kwargs.get("partition") is not None: 128 ↛ 129line 128 didn't jump to line 129 because the condition on line 128 was never true(empty)

129 produce_kwargs["partition"] = kwargs["partition"] 

130 if kwargs.get("timestamp") is not None: (empty)

131 produce_kwargs["timestamp"] = kwargs["timestamp"] (empty)

132 self.producer.produce(topic, **produce_kwargs) (empty)

133 

134 if no_confirm: 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true(empty)

135 return result_future 

136 return await result_future (empty)

137 

138 def create_batch(self) -> "BatchBuilder": (empty)

139 """Creates a batch for sending multiple messages.""" 

140 return BatchBuilder() (empty)

141 

142 async def send_batch( (empty)

143 self, 

144 batch: "BatchBuilder", 

145 topic: str, 

146 *, 

147 partition: int | None, 

148 no_confirm: bool = False, 

149 ) -> None: 

150 """Sends a batch of messages to a Kafka topic.""" 

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

152 for msg in batch._builder: (empty)

153 tg.start_soon( (empty)

154 self.send, 

155 topic, 

156 msg["value"], 

157 msg["key"], 

158 partition, 

159 msg["timestamp_ms"], 

160 msg["headers"], 

161 no_confirm, 

162 ) 

163 

164 async def ping( (empty)

165 self, 

166 timeout: float | None = 5.0, 

167 ) -> bool: 

168 """Implement ping using `list_topics` information request.""" 

169 if timeout is None: 169 ↛ 170line 169 didn't jump to line 170 because the condition on line 169 was never true(empty)

170 timeout = -1 

171 

172 try: (empty)

173 cluster_metadata = await call_or_await( (empty)

174 self.producer.list_topics, 

175 timeout=timeout, 

176 ) 

177 

178 return bool(cluster_metadata) (empty)

179 

180 except Exception: 

181 return False 

182 

183 

184class AsyncConfluentConsumer: (empty)

185 """An asynchronous Python Kafka client for consuming messages using the "confluent-kafka" package.""" 

186 

187 def __init__( (empty)

188 self, 

189 *topics: str, (empty)

190 config: config_module.ConfluentFastConfig, (empty)

191 logger: "LoggerState", (empty)

192 admin_service: "AdminService", (empty)

193 # kwargs options 

194 partitions: Sequence["TopicPartition"], (empty)

195 bootstrap_servers: str | list[str] = "localhost", (empty)

196 # consumer options 

197 client_id: str | None = "confluent-kafka-consumer", (empty)

198 group_id: str | None = None, (empty)

199 group_instance_id: str | None = None, (empty)

200 fetch_max_wait_ms: int = 500, (empty)

201 fetch_max_bytes: int = 52428800, (empty)

202 fetch_min_bytes: int = 1, (empty)

203 max_partition_fetch_bytes: int = 1 * 1024 * 1024, (empty)

204 retry_backoff_ms: int = 100, (empty)

205 auto_offset_reset: str = "latest", (empty)

206 enable_auto_commit: bool = True, (empty)

207 auto_commit_interval_ms: int = 5000, (empty)

208 check_crcs: bool = True, (empty)

209 metadata_max_age_ms: int = 5 * 60 * 1000, (empty)

210 partition_assignment_strategy: str | list[Any] = "roundrobin", (empty)

211 max_poll_interval_ms: int = 300000, (empty)

212 session_timeout_ms: int = 10000, (empty)

213 heartbeat_interval_ms: int = 3000, (empty)

214 security_protocol: str = "PLAINTEXT", (empty)

215 connections_max_idle_ms: int = 540000, (empty)

216 isolation_level: str = "read_uncommitted", (empty)

217 allow_auto_create_topics: bool = True, (empty)

218 # rebalance callbacks 

219 on_assign: Callable[..., None] | None = None, (empty)

220 on_revoke: Callable[..., None] | None = None, (empty)

221 on_lost: Callable[..., None] | None = None, (empty)

222 ) -> None: (empty)

223 self.admin_client = admin_service (empty)

224 self.logger_state = logger (empty)

225 

226 self._on_assign = on_assign (empty)

227 self._on_revoke = on_revoke (empty)

228 self._on_lost = on_lost (empty)

229 

230 self.topics = list(topics) (empty)

231 self.partitions = partitions (empty)

232 

233 if not isinstance(partition_assignment_strategy, str): (empty)

234 partition_assignment_strategy = ",".join( (empty)

235 [ (empty)

236 x if isinstance(x, str) else x().name (empty)

237 for x in partition_assignment_strategy (empty)

238 ], 

239 ) 

240 

241 config_from_params = { (empty)

242 "allow.auto.create.topics": allow_auto_create_topics, (empty)

243 "topic.metadata.refresh.interval.ms": 1000, (empty)

244 "bootstrap.servers": bootstrap_servers, (empty)

245 "client.id": client_id, (empty)

246 "group.id": group_id or "faststream-consumer-group", (empty)

247 "group.instance.id": group_instance_id, (empty)

248 "fetch.wait.max.ms": fetch_max_wait_ms, (empty)

249 "fetch.max.bytes": fetch_max_bytes, (empty)

250 "fetch.min.bytes": fetch_min_bytes, (empty)

251 "max.partition.fetch.bytes": max_partition_fetch_bytes, (empty)

252 "fetch.error.backoff.ms": retry_backoff_ms, (empty)

253 "auto.offset.reset": auto_offset_reset, (empty)

254 "enable.auto.commit": enable_auto_commit, (empty)

255 "auto.commit.interval.ms": auto_commit_interval_ms, (empty)

256 "check.crcs": check_crcs, (empty)

257 "metadata.max.age.ms": metadata_max_age_ms, (empty)

258 "partition.assignment.strategy": partition_assignment_strategy, (empty)

259 "max.poll.interval.ms": max_poll_interval_ms, (empty)

260 "session.timeout.ms": session_timeout_ms, (empty)

261 "heartbeat.interval.ms": heartbeat_interval_ms, (empty)

262 "security.protocol": security_protocol.lower(), (empty)

263 "connections.max.idle.ms": connections_max_idle_ms, (empty)

264 "isolation.level": isolation_level, (empty)

265 } | config.consumer_config (empty)

266 

267 self.config = config_from_params (empty)

268 self.consumer = Consumer(self.config, logger=self.logger_state.logger.logger) (empty)

269 

270 # A pool with single thread is used in order to execute the commands of the consumer sequentially: 

271 # https://github.com/ag2ai/faststream/issues/1904#issuecomment-2506990895 

272 self._thread_pool = ThreadPoolExecutor(max_workers=1) (empty)

273 

274 @property (empty)

275 def topics_to_create(self) -> list[str]: (empty)

276 return list({*self.topics, *(p.topic for p in self.partitions)}) (empty)

277 

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

279 """Starts the Kafka consumer and subscribes to the specified topics.""" 

280 if self.config.get("allow.auto.create.topics", True): (empty)

281 topics_creation_result = await run_in_executor( (empty)

282 self._thread_pool, 

283 self.admin_client.create_topics, 

284 self.topics_to_create, 

285 ) 

286 

287 for create_result in topics_creation_result: (empty)

288 if create_result.error: (empty)

289 self.logger_state.log( 

290 log_level=logging.WARNING, 

291 message=f"Failed to create topic {create_result.topic}: {create_result.error}", 

292 ) 

293 

294 else: 

295 self.logger_state.log( 

296 log_level=logging.WARNING, 

297 message="Auto create topics is disabled. Make sure the topics exist.", 

298 ) 

299 

300 if self.topics: (empty)

301 subscribe_kwargs: dict[str, Any] = {"topics": self.topics} (empty)

302 if self._on_assign is not None: 302 ↛ 303line 302 didn't jump to line 303 because the condition on line 302 was never true(empty)

303 subscribe_kwargs["on_assign"] = self._on_assign 

304 if self._on_revoke is not None: 304 ↛ 305line 304 didn't jump to line 305 because the condition on line 304 was never true(empty)

305 subscribe_kwargs["on_revoke"] = self._on_revoke 

306 if self._on_lost is not None: 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true(empty)

307 subscribe_kwargs["on_lost"] = self._on_lost 

308 await run_in_executor( (empty)

309 self._thread_pool, 

310 self.consumer.subscribe, 

311 **subscribe_kwargs, 

312 ) 

313 

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

315 await run_in_executor( (empty)

316 self._thread_pool, 

317 self.consumer.assign, 

318 [p.to_confluent() for p in self.partitions], 

319 ) 

320 

321 else: 

322 msg = "You must provide either `topics` or `partitions` option." 

323 raise SetupError(msg) 

324 

325 async def commit(self, asynchronous: bool = True) -> None: (empty)

326 """Commits the offsets of all messages returned by the last poll operation.""" 

327 await run_in_executor( (empty)

328 self._thread_pool, 

329 lambda: self.consumer.commit(asynchronous=asynchronous), # type: ignore[call-overload] 

330 ) 

331 

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

333 """Stops the Kafka consumer and releases all resources.""" 

334 # NOTE: If we don't explicitly call commit and then close the consumer, the confluent consumer gets stuck. 

335 # We are doing this to avoid the issue. 

336 enable_auto_commit = self.config.get("enable.auto.commit", True) (empty)

337 

338 try: (empty)

339 if enable_auto_commit: (empty)

340 await self.commit(asynchronous=False) (empty)

341 

342 except Exception as e: (empty)

343 # No offset stored issue is not a problem - https://github.com/confluentinc/confluent-kafka-python/issues/295#issuecomment-355907183 

344 if "No offset stored" in str(e): (empty)

345 pass (empty)

346 else: 

347 self.logger_state.log( (empty)

348 log_level=logging.ERROR, 

349 message="Consumer closing error occurred.", 

350 exc_info=e, 

351 ) 

352 

353 # Wrap calls to async to make method cancelable by timeout 

354 # We shouldn't read messages and close consumer concurrently 

355 # https://github.com/ag2ai/faststream/issues/1904#issuecomment-2506990895 

356 # Now it works without lock due `ThreadPoolExecutor(max_workers=1)` 

357 # that makes all calls to consumer sequential 

358 await run_in_executor(self._thread_pool, self.consumer.close) (empty)

359 

360 self._thread_pool.shutdown(wait=False) (empty)

361 

362 async def getone(self, timeout: float = 0.1) -> Message | None: (empty)

363 """Consumes a single message from Kafka.""" 

364 msg = await run_in_executor(self._thread_pool, self.consumer.poll, timeout) (empty)

365 return check_msg_error(msg) (empty)

366 

367 async def getmany( (empty)

368 self, 

369 timeout: float = 0.1, (empty)

370 max_records: int | None = 10, (empty)

371 ) -> tuple[Message, ...]: (empty)

372 """Consumes a batch of messages from Kafka and groups them by topic and partition.""" 

373 raw_messages: list[Message | None] = await run_in_executor( (empty)

374 self._thread_pool, (empty)

375 cast( (empty)

376 "Callable[..., list[Message | None]]", (empty)

377 lambda: self.consumer.consume( (empty)

378 num_messages=max_records or 10, (empty)

379 timeout=timeout, (empty)

380 ), 

381 ), 

382 ) 

383 return tuple(x for x in map(check_msg_error, raw_messages) if x is not None) (empty)

384 

385 async def seek(self, topic: str, partition: int, offset: int) -> None: (empty)

386 """Seeks to the specified offset in the specified topic and partition.""" 

387 topic_partition = TopicPartition( (empty)

388 topic=topic, 

389 partition=partition, 

390 offset=offset, 

391 ) 

392 await run_in_executor( (empty)

393 self._thread_pool, 

394 self.consumer.seek, 

395 topic_partition.to_confluent(), 

396 ) 

397 

398 

399def check_msg_error(msg: Message | None) -> Message | None: (empty)

400 """Checks for errors in the consumed message.""" 

401 if msg is None or msg.error(): (empty)

402 return None (empty)

403 

404 return msg (empty)

405 

406 

407class BatchBuilder: (empty)

408 """A helper class to build a batch of messages to send to Kafka.""" 

409 

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

411 """Initializes a new BatchBuilder instance.""" 

412 self._builder: list[dict[str, Any]] = [] (empty)

413 

414 def append( (empty)

415 self, 

416 *, 

417 value: bytes | str | None = None, 

418 key: bytes | str | None = None, 

419 timestamp: int | None = None, 

420 headers: list[tuple[str, bytes]] | None = None, 

421 ) -> None: 

422 """Appends a message to the batch with optional timestamp, key, value, and headers.""" 

423 if key is None and value is None: 423 ↛ 424line 423 didn't jump to line 424 because the condition on line 423 was never true(empty)

424 raise KafkaException( 

425 KafkaError(40, "Both key and value can't be None"), 

426 ) 

427 

428 self._builder.append( (empty)

429 { 

430 "timestamp_ms": timestamp or round(time() * 1000), 

431 "key": key, 

432 "value": value, 

433 "headers": headers or [], 

434 }, 

435 )