Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/redis/event.py: 48%

1import asyncio 

2import threading 

3from abc import ABC, abstractmethod 

4from enum import Enum 

5from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union 

6 

7from redis.auth.token import TokenInterface 

8from redis.credentials import CredentialProvider, StreamingCredentialProvider 

9from redis.observability.recorder import ( 

10 init_connection_count, 

11 register_pools_connection_count, 

12) 

13from redis.utils import check_protocol_version, deprecated_function 

14 

15if TYPE_CHECKING: 

16 from redis.maint_notifications import ( 

17 MaintenanceNotification, 

18 MaintenanceState, 

19 MaintNotificationsConfig, 

20 ) 

21 

22 

23class EventListenerInterface(ABC): 

24 """ 

25 Represents a listener for given event object. 

26 """ 

27 

28 @abstractmethod 

29 def listen(self, event: object): 

30 pass 

31 

32 

33class AsyncEventListenerInterface(ABC): 

34 """ 

35 Represents an async listener for given event object. 

36 """ 

37 

38 @abstractmethod 

39 async def listen(self, event: object): 

40 pass 

41 

42 

43class EventDispatcherInterface(ABC): 

44 """ 

45 Represents a dispatcher that dispatches events to listeners 

46 associated with given event. 

47 """ 

48 

49 @abstractmethod 

50 def dispatch(self, event: object): 

51 pass 

52 

53 @abstractmethod 

54 async def dispatch_async(self, event: object): 

55 pass 

56 

57 @abstractmethod 

58 def register_listeners( 

59 self, 

60 mappings: Dict[ 

61 Type[object], 

62 List[Union[EventListenerInterface, AsyncEventListenerInterface]], 

63 ], 

64 ): 

65 """Register additional listeners.""" 

66 pass 

67 

68 @abstractmethod 

69 def unregister_listeners( 

70 self, 

71 mappings: Dict[ 

72 Type[object], 

73 List[Union[EventListenerInterface, AsyncEventListenerInterface]], 

74 ], 

75 ): 

76 """Remove previously registered listeners by identity.""" 

77 pass 

78 

79 

80class EventException(Exception): 

81 """ 

82 Exception wrapper that adds an event object into exception context. 

83 """ 

84 

85 def __init__(self, exception: Exception, event: object): 

86 self.exception = exception 

87 self.event = event 

88 super().__init__(exception) 

89 

90 

91class EventDispatcher(EventDispatcherInterface): 

92 # TODO: Make dispatcher to accept external mappings. 

93 def __init__( 

94 self, 

95 event_listeners: Optional[ 

96 Dict[Type[object], List[EventListenerInterface]] 

97 ] = None, 

98 ): 

99 """ 

100 Dispatcher that dispatches events to listeners associated with given event. 

101 """ 

102 self._event_listeners_mapping: Dict[ 

103 Type[object], List[EventListenerInterface] 

104 ] = { 

105 AfterConnectionReleasedEvent: [ 

106 ReAuthConnectionListener(), 

107 ], 

108 AfterPooledConnectionsInstantiationEvent: [ 

109 RegisterReAuthForPooledConnections(), 

110 ], 

111 AfterSingleConnectionInstantiationEvent: [ 

112 RegisterReAuthForSingleConnection() 

113 ], 

114 AfterPubSubConnectionInstantiationEvent: [RegisterReAuthForPubSub()], 

115 AfterAsyncClusterInstantiationEvent: [RegisterReAuthForAsyncClusterNodes()], 

116 AsyncAfterConnectionReleasedEvent: [ 

117 AsyncReAuthConnectionListener(), 

118 ], 

119 } 

120 

121 # Reentrant so a finalizer/listener that runs on the same thread 

122 # while the lock is held (e.g. a weakref.finalize callback fired 

123 # from cyclic GC during an allocation inside register_listeners / 

124 # unregister_listeners) can re-enter without deadlocking. 

125 self._lock = threading.RLock() 

126 self._async_lock = None 

127 

128 if event_listeners: 

129 self.register_listeners(event_listeners) 

130 

131 def dispatch(self, event: object): 

132 # Snapshot listeners under the lock, then release it before invoking 

133 # them. Holding the lock across listener execution would turn any 

134 # listener that calls register_listeners / unregister_listeners / 

135 # dispatch back into the dispatcher into a deadlock. 

136 with self._lock: 

137 listeners = list(self._event_listeners_mapping.get(type(event), [])) 

138 for listener in listeners: 

139 listener.listen(event) 

140 

141 async def dispatch_async(self, event: object): 

142 if self._async_lock is None: 

143 self._async_lock = asyncio.Lock() 

144 

145 # Snapshot listeners under the lock, then release it before awaiting 

146 # them. See the note in dispatch(); the same rationale applies here 

147 # for dispatch_async re-entry from within a listener. 

148 async with self._async_lock: 

149 listeners = list(self._event_listeners_mapping.get(type(event), [])) 

150 for listener in listeners: 

151 await listener.listen(event) 

152 

153 def register_listeners( 

154 self, 

155 mappings: Dict[ 

156 Type[object], 

157 List[Union[EventListenerInterface, AsyncEventListenerInterface]], 

158 ], 

159 ): 

160 with self._lock: 

161 for event_type in mappings: 

162 if event_type in self._event_listeners_mapping: 

163 self._event_listeners_mapping[event_type] = list( 

164 set( 

165 self._event_listeners_mapping[event_type] 

166 + mappings[event_type] 

167 ) 

168 ) 

169 else: 

170 self._event_listeners_mapping[event_type] = mappings[event_type] 

171 

172 def unregister_listeners( 

173 self, 

174 mappings: Dict[ 

175 Type[object], 

176 List[Union[EventListenerInterface, AsyncEventListenerInterface]], 

177 ], 

178 ): 

179 with self._lock: 

180 for event_type, to_remove in mappings.items(): 

181 current = self._event_listeners_mapping.get(event_type) 

182 if not current: 

183 continue 

184 # Remove by identity to match register semantics and to avoid 

185 # reliance on listener __eq__ implementations. 

186 self._event_listeners_mapping[event_type] = [ 

187 listener 

188 for listener in current 

189 if all(listener is not target for target in to_remove) 

190 ] 

191 

192 

193class AfterConnectionReleasedEvent: 

194 """ 

195 Event that will be fired before each command execution. 

196 """ 

197 

198 def __init__(self, connection): 

199 self._connection = connection 

200 

201 @property 

202 def connection(self): 

203 return self._connection 

204 

205 

206class AsyncAfterConnectionReleasedEvent(AfterConnectionReleasedEvent): 

207 pass 

208 

209 

210class AfterSlotsCacheRefreshEvent: 

211 """ 

212 Event fired after NodesManager's slots cache is refreshed, either via a 

213 full re-initialization or a MOVED-driven slot re-mapping. Signal-only; 

214 carries no payload. Listeners typically reconcile per-node bookkeeping 

215 (e.g. ClusterPubSub shard subscriptions). 

216 """ 

217 

218 pass 

219 

220 

221class AsyncAfterSlotsCacheRefreshEvent(AfterSlotsCacheRefreshEvent): 

222 pass 

223 

224 

225class MaintenanceEvent: 

226 """ 

227 Base class of the events a maintenance notification pushed by the server 

228 produces while it relaxes timeouts: on a connection pool for a node handoff 

229 (``MaintenanceState.MOVING``), or on a single connection for a shard 

230 migration or failover (``MaintenanceState.MAINTENANCE``). 

231 

232 Dispatched synchronously, by both the sync and the async client, to the 

233 ``event_dispatcher`` of the ``MaintNotificationsConfig`` the handler acted 

234 with; listeners implement ``EventListenerInterface``. 

235 

236 :param connection_pool: The pool whose connections are affected. ``None`` 

237 when the notification was handled by a connection without a pool 

238 handler, such as an OSS cluster node connection. 

239 :param connection: The connection that received the notification. ``None`` 

240 when a pool-level handler had no connection to attribute it to. 

241 :param state: The ``MaintenanceState`` the notification puts its source in. 

242 :param notification: The ``MaintenanceNotification`` that was handled. 

243 :param config: The ``MaintNotificationsConfig`` the handler acted with. 

244 """ 

245 

246 def __init__( 

247 self, 

248 connection_pool: Optional[object], 

249 connection: Optional[object], 

250 state: "MaintenanceState", 

251 notification: Optional["MaintenanceNotification"], 

252 config: "MaintNotificationsConfig", 

253 ): 

254 self._connection_pool = connection_pool 

255 self._connection = connection 

256 self._state = state 

257 self._notification = notification 

258 self._config = config 

259 

260 @property 

261 def connection_pool(self) -> Optional[object]: 

262 return self._connection_pool 

263 

264 @property 

265 def connection(self) -> Optional[object]: 

266 return self._connection 

267 

268 @property 

269 def state(self) -> "MaintenanceState": 

270 return self._state 

271 

272 @property 

273 def notification(self) -> Optional["MaintenanceNotification"]: 

274 return self._notification 

275 

276 @property 

277 def config(self) -> "MaintNotificationsConfig": 

278 return self._config 

279 

280 

281class MaintenanceStartedEvent(MaintenanceEvent): 

282 """ 

283 Event fired when a maintenance notification starts relaxing the timeouts of 

284 its source. Paired with a ``MaintenanceCompletedEvent`` carrying the same 

285 ``state`` and source. 

286 """ 

287 

288 pass 

289 

290 

291class MaintenanceCompletedEvent(MaintenanceEvent): 

292 """ 

293 Event fired when the timeout relaxation announced by a 

294 ``MaintenanceStartedEvent`` is reverted: the completion notification 

295 arrived, the handoff TTL expired, or the connection was closed while 

296 still under maintenance. Also fired, with ``MaintenanceState.MAINTENANCE``, 

297 when a handoff takes over a connection under maintenance: the connection 

298 stays relaxed, but under the handoff's ``MOVING`` pair from then on. In 

299 the last two cases ``notification`` is ``None``. 

300 

301 Not fired when a newer handoff superseded the one that completed; the 

302 source stays relaxed until the newer one completes. 

303 """ 

304 

305 pass 

306 

307 

308class ClientType(Enum): 

309 SYNC = ("sync",) 

310 ASYNC = ("async",) 

311 

312 

313class AfterPooledConnectionsInstantiationEvent: 

314 """ 

315 Event that will be fired after pooled connection instances was created. 

316 """ 

317 

318 def __init__( 

319 self, 

320 connection_pools: List, 

321 client_type: ClientType, 

322 credential_provider: Optional[CredentialProvider] = None, 

323 ): 

324 self._connection_pools = connection_pools 

325 self._client_type = client_type 

326 self._credential_provider = credential_provider 

327 

328 @property 

329 def connection_pools(self): 

330 return self._connection_pools 

331 

332 @property 

333 def client_type(self) -> ClientType: 

334 return self._client_type 

335 

336 @property 

337 def credential_provider(self) -> Union[CredentialProvider, None]: 

338 return self._credential_provider 

339 

340 

341class AfterSingleConnectionInstantiationEvent: 

342 """ 

343 Event that will be fired after single connection instances was created. 

344 

345 :param connection_lock: For sync client thread-lock should be provided, 

346 for async asyncio.Lock 

347 """ 

348 

349 def __init__( 

350 self, 

351 connection, 

352 client_type: ClientType, 

353 connection_lock: Union[threading.RLock, asyncio.Lock], 

354 ): 

355 self._connection = connection 

356 self._client_type = client_type 

357 self._connection_lock = connection_lock 

358 

359 @property 

360 def connection(self): 

361 return self._connection 

362 

363 @property 

364 def client_type(self) -> ClientType: 

365 return self._client_type 

366 

367 @property 

368 def connection_lock(self) -> Union[threading.RLock, asyncio.Lock]: 

369 return self._connection_lock 

370 

371 

372class AfterPubSubConnectionInstantiationEvent: 

373 def __init__( 

374 self, 

375 pubsub_connection, 

376 connection_pool, 

377 client_type: ClientType, 

378 connection_lock: Union[threading.RLock, asyncio.Lock], 

379 ): 

380 self._pubsub_connection = pubsub_connection 

381 self._connection_pool = connection_pool 

382 self._client_type = client_type 

383 self._connection_lock = connection_lock 

384 

385 @property 

386 def pubsub_connection(self): 

387 return self._pubsub_connection 

388 

389 @property 

390 def connection_pool(self): 

391 return self._connection_pool 

392 

393 @property 

394 def client_type(self) -> ClientType: 

395 return self._client_type 

396 

397 @property 

398 def connection_lock(self) -> Union[threading.RLock, asyncio.Lock]: 

399 return self._connection_lock 

400 

401 

402class AfterAsyncClusterInstantiationEvent: 

403 """ 

404 Event that will be fired after async cluster instance was created. 

405 

406 Async cluster doesn't use connection pools, 

407 instead ClusterNode object manages connections. 

408 """ 

409 

410 def __init__( 

411 self, 

412 nodes: dict, 

413 credential_provider: Optional[CredentialProvider] = None, 

414 ): 

415 self._nodes = nodes 

416 self._credential_provider = credential_provider 

417 

418 @property 

419 def nodes(self) -> dict: 

420 return self._nodes 

421 

422 @property 

423 def credential_provider(self) -> Union[CredentialProvider, None]: 

424 return self._credential_provider 

425 

426 

427class OnCommandsFailEvent: 

428 """ 

429 Event fired whenever a command fails during the execution. 

430 """ 

431 

432 def __init__( 

433 self, 

434 commands: tuple, 

435 exception: Exception, 

436 ): 

437 self._commands = commands 

438 self._exception = exception 

439 

440 @property 

441 def commands(self) -> tuple: 

442 return self._commands 

443 

444 @property 

445 def exception(self) -> Exception: 

446 return self._exception 

447 

448 

449class AsyncOnCommandsFailEvent(OnCommandsFailEvent): 

450 pass 

451 

452 

453class ReAuthConnectionListener(EventListenerInterface): 

454 """ 

455 Listener that performs re-authentication of given connection. 

456 """ 

457 

458 def listen(self, event: AfterConnectionReleasedEvent): 

459 event.connection.re_auth() 

460 

461 

462class AsyncReAuthConnectionListener(AsyncEventListenerInterface): 

463 """ 

464 Async listener that performs re-authentication of given connection. 

465 """ 

466 

467 async def listen(self, event: AsyncAfterConnectionReleasedEvent): 

468 await event.connection.re_auth() 

469 

470 

471class RegisterReAuthForPooledConnections(EventListenerInterface): 

472 """ 

473 Listener that registers a re-authentication callback for pooled connections. 

474 Required by :class:`StreamingCredentialProvider`. 

475 """ 

476 

477 def __init__(self): 

478 self._event = None 

479 

480 def listen(self, event: AfterPooledConnectionsInstantiationEvent): 

481 if isinstance(event.credential_provider, StreamingCredentialProvider): 

482 self._event = event 

483 

484 if event.client_type == ClientType.SYNC: 

485 event.credential_provider.on_next(self._re_auth) 

486 event.credential_provider.on_error(self._raise_on_error) 

487 else: 

488 event.credential_provider.on_next(self._re_auth_async) 

489 event.credential_provider.on_error(self._raise_on_error_async) 

490 

491 def _re_auth(self, token): 

492 for pool in self._event.connection_pools: 

493 pool.re_auth_callback(token) 

494 

495 async def _re_auth_async(self, token): 

496 for pool in self._event.connection_pools: 

497 await pool.re_auth_callback(token) 

498 

499 def _raise_on_error(self, error: Exception): 

500 raise EventException(error, self._event) 

501 

502 async def _raise_on_error_async(self, error: Exception): 

503 raise EventException(error, self._event) 

504 

505 

506class RegisterReAuthForSingleConnection(EventListenerInterface): 

507 """ 

508 Listener that registers a re-authentication callback for single connection. 

509 Required by :class:`StreamingCredentialProvider`. 

510 """ 

511 

512 def __init__(self): 

513 self._event = None 

514 

515 def listen(self, event: AfterSingleConnectionInstantiationEvent): 

516 if isinstance( 

517 event.connection.credential_provider, StreamingCredentialProvider 

518 ): 

519 self._event = event 

520 

521 if event.client_type == ClientType.SYNC: 

522 event.connection.credential_provider.on_next(self._re_auth) 

523 event.connection.credential_provider.on_error(self._raise_on_error) 

524 else: 

525 event.connection.credential_provider.on_next(self._re_auth_async) 

526 event.connection.credential_provider.on_error( 

527 self._raise_on_error_async 

528 ) 

529 

530 def _re_auth(self, token): 

531 with self._event.connection_lock: 

532 self._event.connection.send_command( 

533 "AUTH", token.try_get("oid"), token.get_value() 

534 ) 

535 self._event.connection.read_response() 

536 

537 async def _re_auth_async(self, token): 

538 async with self._event.connection_lock: 

539 await self._event.connection.send_command( 

540 "AUTH", token.try_get("oid"), token.get_value() 

541 ) 

542 await self._event.connection.read_response() 

543 

544 def _raise_on_error(self, error: Exception): 

545 raise EventException(error, self._event) 

546 

547 async def _raise_on_error_async(self, error: Exception): 

548 raise EventException(error, self._event) 

549 

550 

551class RegisterReAuthForAsyncClusterNodes(EventListenerInterface): 

552 def __init__(self): 

553 self._event = None 

554 

555 def listen(self, event: AfterAsyncClusterInstantiationEvent): 

556 if isinstance(event.credential_provider, StreamingCredentialProvider): 

557 self._event = event 

558 event.credential_provider.on_next(self._re_auth) 

559 event.credential_provider.on_error(self._raise_on_error) 

560 

561 async def _re_auth(self, token: TokenInterface): 

562 for key in self._event.nodes: 

563 await self._event.nodes[key].re_auth_callback(token) 

564 

565 async def _raise_on_error(self, error: Exception): 

566 raise EventException(error, self._event) 

567 

568 

569class RegisterReAuthForPubSub(EventListenerInterface): 

570 def __init__(self): 

571 self._connection = None 

572 self._connection_pool = None 

573 self._client_type = None 

574 self._connection_lock = None 

575 self._event = None 

576 

577 def listen(self, event: AfterPubSubConnectionInstantiationEvent): 

578 if isinstance( 

579 event.pubsub_connection.credential_provider, StreamingCredentialProvider 

580 ) and check_protocol_version(event.pubsub_connection.get_protocol(), 3): 

581 self._event = event 

582 self._connection = event.pubsub_connection 

583 self._connection_pool = event.connection_pool 

584 self._client_type = event.client_type 

585 self._connection_lock = event.connection_lock 

586 

587 if self._client_type == ClientType.SYNC: 

588 self._connection.credential_provider.on_next(self._re_auth) 

589 self._connection.credential_provider.on_error(self._raise_on_error) 

590 else: 

591 self._connection.credential_provider.on_next(self._re_auth_async) 

592 self._connection.credential_provider.on_error( 

593 self._raise_on_error_async 

594 ) 

595 

596 def _re_auth(self, token: TokenInterface): 

597 with self._connection_lock: 

598 self._connection.send_command( 

599 "AUTH", token.try_get("oid"), token.get_value() 

600 ) 

601 self._connection.read_response() 

602 

603 self._connection_pool.re_auth_callback(token) 

604 

605 async def _re_auth_async(self, token: TokenInterface): 

606 async with self._connection_lock: 

607 await self._connection.send_command( 

608 "AUTH", token.try_get("oid"), token.get_value() 

609 ) 

610 await self._connection.read_response() 

611 

612 await self._connection_pool.re_auth_callback(token) 

613 

614 def _raise_on_error(self, error: Exception): 

615 raise EventException(error, self._event) 

616 

617 async def _raise_on_error_async(self, error: Exception): 

618 raise EventException(error, self._event) 

619 

620 

621class InitializeConnectionCountObservability(EventListenerInterface): 

622 """ 

623 Listener that initializes connection count observability. 

624 """ 

625 

626 @deprecated_function( 

627 reason="Connection count is now tracked via record_connection_count(). " 

628 "This functionality will be removed in the next major version", 

629 version="7.4.0", 

630 ) 

631 def listen(self, event: AfterPooledConnectionsInstantiationEvent): 

632 # Initialize gauge only once, subsequent calls won't have an affect. 

633 # Note: init_connection_count() and register_pools_connection_count() 

634 # are deprecated and will emit their own warnings. 

635 init_connection_count() 

636 

637 # Register pools for connection count observability. 

638 register_pools_connection_count(event.connection_pools)