Coverage for faststream / kafka / helpers / rebalance_listener.py: 87%

1import asyncio (empty)

2import logging (empty)

3from typing import TYPE_CHECKING, Any, Optional (empty)

4 

5from aiokafka import ConsumerRebalanceListener (empty)

6 

7from faststream._internal.utils.functions import call_or_await (empty)

8 

9if TYPE_CHECKING: (empty)

10 from aiokafka import AIOKafkaConsumer, TopicPartition 

11 

12 from faststream._internal.basic_types import LoggerProto 

13 

14 

15def make_logging_listener( (empty)

16 *, 

17 consumer: "AIOKafkaConsumer", 

18 logger: Optional["LoggerProto"], 

19 log_extra: dict[str, Any], 

20 listener: Optional["ConsumerRebalanceListener"], 

21) -> Optional["ConsumerRebalanceListener"]: 

22 if logger is None: 22 ↛ 23line 22 didn't jump to line 23 because the condition on line 22 was never true(empty)

23 return listener 

24 

25 logging_listener = _LoggingListener( (empty)

26 consumer=consumer, 

27 logger=logger, 

28 log_extra=log_extra, 

29 ) 

30 if listener is None: (empty)

31 return logging_listener (empty)

32 

33 return _LoggingListenerFacade( (empty)

34 logging_listener=logging_listener, 

35 listener=listener, 

36 ) 

37 

38 

39class _LoggingListener(ConsumerRebalanceListener): # type: ignore[misc] (empty)

40 _log_unassigned_consumer_delay_seconds = 60 * 2 (empty)

41 _log_unassigned_consumer_task: asyncio.Task[None] | None (empty)

42 

43 def __init__( (empty)

44 self, 

45 *, 

46 consumer: "AIOKafkaConsumer", 

47 logger: "LoggerProto", 

48 log_extra: dict[str, Any], 

49 ) -> None: 

50 self.consumer = consumer (empty)

51 self.logger = logger (empty)

52 self.log_extra = log_extra (empty)

53 self._log_unassigned_consumer_task = None (empty)

54 

55 async def on_partitions_revoked(self, revoked: set["TopicPartition"]) -> None: (empty)

56 pass (empty)

57 

58 async def log_unassigned_consumer(self) -> None: (empty)

59 await asyncio.sleep(self._log_unassigned_consumer_delay_seconds) (empty)

60 self.logger.log( (empty)

61 logging.WARNING, 

62 f"Consumer in group {self.consumer._group_id} has had no partition " 

63 f"assignments for {self._log_unassigned_consumer_delay_seconds} seconds: " 

64 f"topics {self.consumer._subscription.topics} may have fewer partitions " 

65 f"than consumers.", 

66 extra=self.log_extra, 

67 ) 

68 

69 async def on_partitions_assigned(self, assigned: set["TopicPartition"]) -> None: (empty)

70 self.logger.log( (empty)

71 logging.INFO, 

72 f"Consumer {self.consumer._coordinator.member_id} assigned to partitions: " 

73 f"({', '.join(str(tp) for tp in assigned)})", 

74 extra=self.log_extra, 

75 ) 

76 

77 if not assigned: (empty)

78 self.logger.log( (empty)

79 logging.WARNING, 

80 f"Consumer in group {self.consumer._group_id} has no partition assignments - this " 

81 f"could be temporary, e.g. during a rolling update. A separate warning will be logged if " 

82 f"this condition persists for {self._log_unassigned_consumer_delay_seconds} seconds.", 

83 extra=self.log_extra, 

84 ) 

85 

86 self._log_unassigned_consumer_task: asyncio.Task[None] | None = ( (empty)

87 asyncio.create_task(self.log_unassigned_consumer()) 

88 ) 

89 

90 elif self._log_unassigned_consumer_task: 90 ↛ 91line 90 didn't jump to line 91 because the condition on line 90 was never true(empty)

91 self._log_unassigned_consumer_task.cancel() 

92 self._log_unassigned_consumer_task = None 

93 

94 

95class _LoggingListenerFacade(ConsumerRebalanceListener): # type: ignore[misc] (empty)

96 def __init__( (empty)

97 self, 

98 *, 

99 logging_listener: _LoggingListener, 

100 listener: ConsumerRebalanceListener, 

101 ) -> None: 

102 self.logging_listener = logging_listener (empty)

103 self.listener = listener (empty)

104 

105 async def on_partitions_revoked(self, revoked: set["TopicPartition"]) -> None: (empty)

106 await call_or_await(self.listener.on_partitions_revoked, revoked) (empty)

107 

108 async def on_partitions_assigned(self, assigned: set["TopicPartition"]) -> None: (empty)

109 await self.logging_listener.on_partitions_revoked(assigned) (empty)

110 await call_or_await(self.listener.on_partitions_assigned, assigned) (empty)