Coverage for faststream / kafka / publisher / producer.py: 86%

1from abc import abstractmethod (empty)

2from typing import TYPE_CHECKING, Any, Optional, Union (empty)

3 

4from typing_extensions import override (empty)

5 

6from faststream._internal.endpoint.utils import ParserComposition (empty)

7from faststream._internal.producer import ProducerProto (empty)

8from faststream.exceptions import FeatureNotSupportedException (empty)

9from faststream.kafka.exceptions import BatchBufferOverflowException (empty)

10from faststream.kafka.message import KafkaMessage (empty)

11from faststream.kafka.parser import AioKafkaParser (empty)

12from faststream.kafka.response import KafkaPublishCommand (empty)

13from faststream.message import encode_message (empty)

14 

15from .state import EmptyProducerState, ProducerState, RealProducer (empty)

16 

17if TYPE_CHECKING: (empty)

18 import asyncio 

19 

20 from aiokafka import AIOKafkaProducer 

21 from aiokafka.structs import RecordMetadata 

22 from fast_depends.library.serializer import SerializerProto 

23 

24 from faststream._internal.types import CustomCallable 

25 

26 

27class AioKafkaFastProducer(ProducerProto[KafkaPublishCommand]): (empty)

28 async def connect( (empty)

29 self, 

30 producer: "AIOKafkaProducer", (empty)

31 serializer: Optional["SerializerProto"], (empty)

32 ) -> None: ... (empty)

33 

34 async def disconnect(self) -> None: ... (empty)

35 

36 def __bool__(self) -> bool: (empty)

37 return False 

38 

39 @property (empty)

40 def closed(self) -> bool: (empty)

41 return True 

42 

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

44 return None 

45 

46 @abstractmethod (empty)

47 async def publish( (empty)

48 self, 

49 cmd: "KafkaPublishCommand", (empty)

50 ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]: ... (empty)

51 

52 @abstractmethod (empty)

53 async def publish_batch( (empty)

54 self, 

55 cmd: "KafkaPublishCommand", (empty)

56 ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]: ... (empty)

57 

58 async def request(self, cmd: "KafkaPublishCommand") -> Any: (empty)

59 msg = "Kafka doesn't support `request` method without test client." 

60 raise FeatureNotSupportedException(msg) 

61 

62 

63class AioKafkaFastProducerImpl(AioKafkaFastProducer): (empty)

64 """A class to represent Kafka producer.""" 

65 

66 def __init__( (empty)

67 self, 

68 parser: Optional["CustomCallable"], 

69 decoder: Optional["CustomCallable"], 

70 ) -> None: 

71 self._producer: ProducerState = EmptyProducerState() (empty)

72 self.serializer: SerializerProto | None = None (empty)

73 

74 # NOTE: register default parser to be compatible with request 

75 default = AioKafkaParser(msg_class=KafkaMessage, regex=None) (empty)

76 self._parser = ParserComposition(parser, default.parse_message) (empty)

77 self._decoder = ParserComposition(decoder, default.decode_message) (empty)

78 

79 async def connect( (empty)

80 self, 

81 producer: "AIOKafkaProducer", 

82 serializer: Optional["SerializerProto"], 

83 ) -> None: 

84 self.serializer = serializer (empty)

85 await producer.start() (empty)

86 self._producer = RealProducer(producer) (empty)

87 

88 async def disconnect(self) -> None: (empty)

89 await self._producer.stop() (empty)

90 self._producer = EmptyProducerState() (empty)

91 

92 def __bool__(self) -> bool: (empty)

93 return bool(self._producer) 

94 

95 @property (empty)

96 def closed(self) -> bool: (empty)

97 return self._producer.closed 

98 

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

100 await self._producer.flush() 

101 

102 @override (empty)

103 async def publish( (empty)

104 self, 

105 cmd: "KafkaPublishCommand", 

106 ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]: 

107 """Publish a message to a topic.""" 

108 message, content_type = encode_message(cmd.body, serializer=self.serializer) (empty)

109 

110 headers_to_send = { (empty)

111 "content-type": content_type or "", 

112 **cmd.headers_to_publish(), 

113 } 

114 

115 send_future = await self._producer.producer.send( (empty)

116 topic=cmd.destination, 

117 value=message, 

118 key=cmd.key, 

119 partition=cmd.partition, 

120 timestamp_ms=cmd.timestamp_ms, 

121 headers=[(i, (j or "").encode()) for i, j in headers_to_send.items()], 

122 ) 

123 

124 if not cmd.no_confirm: (empty)

125 return await send_future (empty)

126 return send_future (empty)

127 

128 @override (empty)

129 async def publish_batch( (empty)

130 self, 

131 cmd: "KafkaPublishCommand", 

132 ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]: 

133 """Publish a batch of messages to a topic.""" 

134 batch = self._producer.producer.create_batch() (empty)

135 

136 headers_to_send = cmd.headers_to_publish() (empty)

137 

138 for message_position, body in enumerate(cmd.batch_bodies): (empty)

139 message, content_type = encode_message(body, serializer=self.serializer) (empty)

140 

141 if content_type: 141 ↛ 147line 141 didn't jump to line 147 because the condition on line 141 was always true(empty)

142 final_headers = { (empty)

143 "content-type": content_type, 

144 **headers_to_send, 

145 } 

146 else: 

147 final_headers = headers_to_send.copy() 

148 

149 metadata = batch.append( (empty)

150 key=cmd.key_for(message_position), 

151 value=message, 

152 timestamp=cmd.timestamp_ms, 

153 headers=[(i, j.encode()) for i, j in final_headers.items()], 

154 ) 

155 if metadata is None: (empty)

156 raise BatchBufferOverflowException(message_position=message_position) (empty)

157 

158 send_future = await self._producer.producer.send_batch( (empty)

159 batch, 

160 cmd.destination, 

161 partition=cmd.partition, 

162 ) 

163 if not cmd.no_confirm: (empty)

164 return await send_future (empty)

165 return send_future (empty)

166 

167 

168class FakeAioKafkaFastProducer(AioKafkaFastProducer): (empty)

169 async def connect( (empty)

170 self, 

171 producer: "AIOKafkaProducer", 

172 serializer: Optional["SerializerProto"], 

173 ) -> None: 

174 raise NotImplementedError 

175 

176 async def disconnect(self) -> None: (empty)

177 raise NotImplementedError 

178 

179 def __bool__(self) -> bool: (empty)

180 return False 

181 

182 @property (empty)

183 def closed(self) -> bool: (empty)

184 raise NotImplementedError 

185 

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

187 raise NotImplementedError 

188 

189 async def publish( (empty)

190 self, 

191 cmd: "KafkaPublishCommand", 

192 ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]: 

193 raise NotImplementedError 

194 

195 async def publish_batch( (empty)

196 self, 

197 cmd: "KafkaPublishCommand", 

198 ) -> Union["asyncio.Future[RecordMetadata]", "RecordMetadata"]: 

199 raise NotImplementedError