Coverage for faststream / redis / subscriber / usecases / stream_subscriber.py: 86%

1import asyncio (empty)

2import logging (empty)

3import math (empty)

4from collections.abc import AsyncIterator, Awaitable, Callable (empty)

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

6 

7from redis.exceptions import ResponseError (empty)

8from typing_extensions import override (empty)

9 

10from faststream._internal.endpoint.subscriber.mixins import ConcurrentMixin (empty)

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

12from faststream.redis.exceptions import StreamGroupNotFoundError (empty)

13from faststream.redis.message import ( (empty)

14 BatchStreamMessage, 

15 DefaultStreamMessage, 

16 RedisStreamMessage, 

17) 

18from faststream.redis.parser import ( (empty)

19 RedisBatchStreamParser, 

20 RedisStreamParser, 

21) 

22 

23from .basic import LogicSubscriber (empty)

24 

25if TYPE_CHECKING: (empty)

26 from anyio import Event 

27 

28 from faststream._internal.endpoint.subscriber import SubscriberSpecification 

29 from faststream._internal.endpoint.subscriber.call_item import ( 

30 CallsCollection, 

31 ) 

32 from faststream.message import StreamMessage as BrokerStreamMessage 

33 from faststream.redis.schemas import StreamSub 

34 from faststream.redis.subscriber.config import RedisSubscriberConfig 

35 

36 

37TopicName: TypeAlias = bytes (empty)

38Offset: TypeAlias = bytes (empty)

39 

40ReadResponse = tuple[ (empty)

41 tuple[ 

42 TopicName, 

43 tuple[ 

44 tuple[ 

45 Offset, 

46 dict[bytes, bytes], 

47 ], 

48 ..., 

49 ], 

50 ], 

51 ..., 

52] 

53ReadCallable = Callable[[str], Awaitable[ReadResponse]] (empty)

54 

55 

56class _StreamHandlerMixin(LogicSubscriber): (empty)

57 def __init__( (empty)

58 self, 

59 config: "RedisSubscriberConfig", 

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

61 calls: "CallsCollection[Any]", 

62 ) -> None: 

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

64 

65 assert config.stream_sub (empty)

66 self._stream_sub = config.stream_sub (empty)

67 self.last_id = config.stream_sub.last_id (empty)

68 self.min_idle_time = config.stream_sub.min_idle_time (empty)

69 self.autoclaim_start_id = b"0-0" (empty)

70 

71 @property (empty)

72 def stream_sub(self) -> "StreamSub": (empty)

73 return self._stream_sub.add_prefix(self._outer_config.prefix) (empty)

74 

75 def get_log_context( (empty)

76 self, 

77 message: Optional["BrokerStreamMessage[Any]"], 

78 ) -> dict[str, str]: 

79 return self.build_log_context( (empty)

80 message=message, 

81 channel=self.stream_sub.name, 

82 ) 

83 

84 @override (empty)

85 async def _consume(self, *args: Any, start_signal: "Event") -> None: (empty)

86 if await self._client.ping(): 86 ↛ 89line 86 didn't jump to line 89 because the condition on line 86 was always true(empty)

87 start_signal.set() (empty)

88 

89 while self.running: (empty)

90 try: (empty)

91 await self._get_msgs(*args) (empty)

92 

93 except ResponseError as e: # noqa: PERF203 (empty)

94 if "NOGROUP" in str(e): (empty)

95 msg = ( (empty)

96 f"Consumer group `{self.stream_sub.group}` for stream " 

97 f"`{self.stream_sub.name}` no longer exists. " 

98 "The stream was likely deleted or flushed. " 

99 "Stopping subscriber — restart the application to recreate the group." 

100 ) 

101 raise StreamGroupNotFoundError(msg) from e (empty)

102 raise 

103 

104 except Exception as e: (empty)

105 self._log( 

106 log_level=logging.ERROR, 

107 message="Message fetch error", 

108 exc_info=e, 

109 ) 

110 

111 finally: 

112 if not start_signal.is_set(): 112 ↛ 113line 112 didn't jump to line 113 because the condition on line 112 was never true(empty)

113 start_signal.set() 

114 

115 @override (empty)

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

117 client = self._client (empty)

118 

119 self.extra_watcher_options.update( (empty)

120 redis=client, 

121 group=self.stream_sub.group, 

122 ) 

123 

124 stream = self.stream_sub (empty)

125 

126 read: ReadCallable 

127 

128 if stream.group and stream.consumer: (empty)

129 group_create_id = "$" if self.last_id == ">" else self.last_id (empty)

130 try: (empty)

131 await client.xgroup_create( (empty)

132 name=stream.name, 

133 id=group_create_id, 

134 groupname=stream.group, 

135 mkstream=True, 

136 ) 

137 except ResponseError as e: (empty)

138 if "already exists" not in str(e): (empty)

139 raise 

140 

141 if stream.min_idle_time is None: (empty)

142 

143 def read( (empty)

144 _: str, 

145 ) -> Awaitable[ReadResponse]: 

146 return client.xreadgroup( (empty)

147 groupname=stream.group, 

148 consumername=stream.consumer, 

149 streams={stream.name: stream.last_id}, 

150 count=stream.max_records, 

151 block=stream.polling_interval, 

152 noack=stream.no_ack, 

153 ) 

154 

155 else: 

156 

157 async def read(_: str) -> ReadResponse: (empty)

158 stream_message = await client.xautoclaim( (empty)

159 name=self.stream_sub.name, 

160 groupname=self.stream_sub.group, 

161 consumername=self.stream_sub.consumer, 

162 min_idle_time=self.min_idle_time, 

163 start_id=self.autoclaim_start_id, 

164 count=1, 

165 ) 

166 stream_name = self.stream_sub.name.encode() (empty)

167 (next_id, messages, *_) = stream_message (empty)

168 

169 # Update start_id for next call 

170 self.autoclaim_start_id = next_id (empty)

171 

172 if next_id == b"0-0" and not messages: (empty)

173 await asyncio.sleep(stream.polling_interval / 1000) # ms to s (empty)

174 return () (empty)

175 

176 return ((stream_name, messages),) (empty)

177 

178 else: 

179 

180 def read( (empty)

181 last_id: str, 

182 ) -> Awaitable[ReadResponse]: 

183 return client.xread( (empty)

184 {stream.name: last_id}, 

185 block=stream.polling_interval, 

186 count=stream.max_records, 

187 ) 

188 

189 await super().start(read) (empty)

190 

191 @override (empty)

192 async def get_one( (empty)

193 self, 

194 *, 

195 timeout: float = 5.0, 

196 ) -> "RedisStreamMessage | None": 

197 assert not self.calls, ( (empty)

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

199 ) 

200 if self.stream_sub.group and self.stream_sub.consumer: (empty)

201 if self.min_idle_time is None: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true(empty)

202 stream_message = await self._client.xreadgroup( 

203 groupname=self.stream_sub.group, 

204 consumername=self.stream_sub.consumer, 

205 streams={self.stream_sub.name: self.last_id}, 

206 block=math.ceil(timeout * 1000), 

207 count=1, 

208 ) 

209 if not stream_message: 

210 return None 

211 

212 ((stream_name, ((message_id, raw_message),)),) = stream_message 

213 else: 

214 stream_message = await self._client.xautoclaim( (empty)

215 name=self.stream_sub.name, 

216 groupname=self.stream_sub.group, 

217 consumername=self.stream_sub.consumer, 

218 min_idle_time=self.min_idle_time, 

219 start_id=self.autoclaim_start_id, 

220 count=1, 

221 ) 

222 (next_id, messages, *_) = stream_message (empty)

223 # Update start_id for next call 

224 self.autoclaim_start_id = next_id (empty)

225 if not messages: (empty)

226 return None (empty)

227 stream_name = self.stream_sub.name.encode() (empty)

228 ((message_id, raw_message),) = messages (empty)

229 else: 

230 stream_message = await self._client.xread( (empty)

231 {self.stream_sub.name: self.last_id}, 

232 block=math.ceil(timeout * 1000), 

233 count=1, 

234 ) 

235 if not stream_message: (empty)

236 return None (empty)

237 

238 ((stream_name, ((message_id, raw_message),)),) = stream_message (empty)

239 

240 self.last_id = message_id.decode() (empty)

241 

242 redis_incoming_msg = DefaultStreamMessage( (empty)

243 type="stream", 

244 channel=stream_name.decode(), 

245 message_ids=[message_id], 

246 data=raw_message, 

247 ) 

248 

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

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

251 

252 msg: RedisStreamMessage = await process_msg( # type: ignore[assignment] (empty)

253 msg=redis_incoming_msg, 

254 middlewares=( 

255 m(redis_incoming_msg, context=context) for m in self._broker_middlewares 

256 ), 

257 parser=async_parser, 

258 decoder=async_decoder, 

259 ) 

260 return msg (empty)

261 

262 @override (empty)

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

264 assert not self.calls, ( (empty)

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

266 ) 

267 

268 timeout = 5 (empty)

269 

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

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

272 

273 while True: (empty)

274 if self.stream_sub.group and self.stream_sub.consumer: (empty)

275 if self.min_idle_time is None: 275 ↛ 276line 275 didn't jump to line 276 because the condition on line 275 was never true(empty)

276 stream_message = await self._client.xreadgroup( 

277 groupname=self.stream_sub.group, 

278 consumername=self.stream_sub.consumer, 

279 streams={self.stream_sub.name: self.last_id}, 

280 block=math.ceil(timeout * 1000), 

281 count=1, 

282 ) 

283 if not stream_message: 

284 continue 

285 

286 ((stream_name, ((message_id, raw_message),)),) = stream_message 

287 else: 

288 stream_message = await self._client.xautoclaim( (empty)

289 name=self.stream_sub.name, 

290 groupname=self.stream_sub.group, 

291 consumername=self.stream_sub.consumer, 

292 min_idle_time=self.min_idle_time, 

293 start_id=self.autoclaim_start_id, 

294 count=1, 

295 ) 

296 (next_id, messages, *_) = stream_message (empty)

297 # Update start_id for next call 

298 self.autoclaim_start_id = next_id (empty)

299 if not messages: 299 ↛ 300line 299 didn't jump to line 300 because the condition on line 299 was never true(empty)

300 continue 

301 stream_name = self.stream_sub.name.encode() (empty)

302 ((message_id, raw_message),) = messages (empty)

303 else: 

304 stream_message = await self._client.xread( (empty)

305 {self.stream_sub.name: self.last_id}, 

306 block=math.ceil(timeout * 1000), 

307 count=1, 

308 ) 

309 if not stream_message: 309 ↛ 310line 309 didn't jump to line 310 because the condition on line 309 was never true(empty)

310 continue 

311 

312 ((stream_name, ((message_id, raw_message),)),) = stream_message (empty)

313 

314 self.last_id = message_id.decode() (empty)

315 

316 redis_incoming_msg = DefaultStreamMessage( (empty)

317 type="stream", 

318 channel=stream_name.decode(), 

319 message_ids=[message_id], 

320 data=raw_message, 

321 ) 

322 

323 msg: RedisStreamMessage = await process_msg( # type: ignore[assignment] (empty)

324 msg=redis_incoming_msg, 

325 middlewares=( 

326 m(redis_incoming_msg, context=context) 

327 for m in self._broker_middlewares 

328 ), 

329 parser=async_parser, 

330 decoder=async_decoder, 

331 ) 

332 yield msg (empty)

333 

334 

335class StreamSubscriber(_StreamHandlerMixin): (empty)

336 def __init__( (empty)

337 self, 

338 config: "RedisSubscriberConfig", 

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

340 calls: "CallsCollection[Any]", 

341 ) -> None: 

342 parser = RedisStreamParser(config) (empty)

343 config.decoder = parser.decode_message (empty)

344 config.parser = parser.parse_message (empty)

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

346 

347 async def _get_msgs( (empty)

348 self, 

349 read: Callable[ (empty)

350 [str], (empty)

351 Awaitable[ (empty)

352 tuple[ (empty)

353 tuple[ (empty)

354 TopicName, (empty)

355 tuple[ (empty)

356 tuple[ (empty)

357 Offset, (empty)

358 dict[bytes, bytes], (empty)

359 ], 

360 ..., (empty)

361 ], 

362 ], 

363 ..., (empty)

364 ], 

365 ], 

366 ], 

367 ) -> None: (empty)

368 for stream_name, msgs in await read(self.last_id): (empty)

369 if msgs: (empty)

370 self.last_id = msgs[-1][0].decode() (empty)

371 

372 for message_id, raw_msg in msgs: (empty)

373 msg = DefaultStreamMessage( (empty)

374 type="stream", (empty)

375 channel=stream_name.decode(), (empty)

376 message_ids=[message_id], (empty)

377 data=raw_msg, (empty)

378 ) 

379 

380 await self.consume_one(msg) (empty)

381 

382 

383class StreamBatchSubscriber(_StreamHandlerMixin): (empty)

384 def __init__( (empty)

385 self, 

386 config: "RedisSubscriberConfig", 

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

388 calls: "CallsCollection[Any]", 

389 ) -> None: 

390 parser = RedisBatchStreamParser(config) (empty)

391 config.decoder = parser.decode_message (empty)

392 config.parser = parser.parse_message (empty)

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

394 

395 async def _get_msgs( (empty)

396 self, 

397 read: Callable[ (empty)

398 [str], (empty)

399 Awaitable[ (empty)

400 tuple[tuple[bytes, tuple[tuple[bytes, dict[bytes, bytes]], ...]], ...], (empty)

401 ], 

402 ], 

403 ) -> None: (empty)

404 for stream_name, msgs in await read(self.last_id): (empty)

405 if msgs: (empty)

406 self.last_id = msgs[-1][0].decode() (empty)

407 

408 data: list[dict[bytes, bytes]] = [] (empty)

409 ids: list[bytes] = [] (empty)

410 for message_id, i in msgs: (empty)

411 data.append(i) (empty)

412 ids.append(message_id) (empty)

413 

414 msg = BatchStreamMessage( (empty)

415 type="bstream", (empty)

416 channel=stream_name.decode(), (empty)

417 data=data, (empty)

418 message_ids=ids, (empty)

419 ) 

420 

421 await self.consume_one(msg) (empty)

422 

423 

424class StreamConcurrentSubscriber( (empty)

425 ConcurrentMixin["BrokerStreamMessage[Any]"], 

426 StreamSubscriber, 

427): 

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

429 await super().start() (empty)

430 self.start_consume_task() (empty)

431 

432 async def consume_one(self, msg: "BrokerStreamMessage[Any]") -> None: (empty)

433 await self._put_msg(msg) (empty)