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

1from abc import abstractmethod (empty)

2from typing import TYPE_CHECKING, Any, Optional (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.confluent.parser import AsyncConfluentParser (empty)

9from faststream.confluent.response import KafkaPublishCommand (empty)

10from faststream.exceptions import FeatureNotSupportedException (empty)

11from faststream.message import encode_message (empty)

12 

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

14 

15if TYPE_CHECKING: (empty)

16 import asyncio 

17 

18 from confluent_kafka import Message 

19 from fast_depends.library.serializer import SerializerProto 

20 

21 from faststream._internal.types import CustomCallable 

22 from faststream.confluent.helpers.client import AsyncConfluentProducer 

23 

24 

25class AsyncConfluentFastProducer(ProducerProto[KafkaPublishCommand]): (empty)

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

27 

28 def connect( (empty)

29 self, 

30 producer: "AsyncConfluentProducer", (empty)

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

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

33 

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

35 return False 

36 

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

38 return None 

39 

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

41 return None 

42 

43 @abstractmethod (empty)

44 async def ping(self, timeout: float) -> bool: (empty)

45 return False 

46 

47 @override (empty)

48 @abstractmethod (empty)

49 async def publish( (empty)

50 self, 

51 cmd: "KafkaPublishCommand", (empty)

52 ) -> "asyncio.Future[Message | None] | Message | None": ... (empty)

53 

54 @override (empty)

55 @abstractmethod (empty)

56 async def publish_batch(self, cmd: "KafkaPublishCommand") -> None: ... (empty)

57 

58 @override (empty)

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

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

61 raise FeatureNotSupportedException(msg) 

62 

63 

64class FakeConfluentFastProducer(AsyncConfluentFastProducer): (empty)

65 def connect( (empty)

66 self, 

67 producer: "AsyncConfluentProducer", 

68 serializer: Optional["SerializerProto"], 

69 ) -> None: 

70 raise NotImplementedError 

71 

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

73 raise NotImplementedError 

74 

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

76 raise NotImplementedError 

77 

78 async def ping(self, timeout: float) -> bool: (empty)

79 raise NotImplementedError 

80 

81 @override (empty)

82 async def publish( (empty)

83 self, 

84 cmd: "KafkaPublishCommand", 

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

86 raise NotImplementedError 

87 

88 @override (empty)

89 async def publish_batch(self, cmd: "KafkaPublishCommand") -> None: (empty)

90 raise NotImplementedError 

91 

92 

93class AsyncConfluentFastProducerImpl(AsyncConfluentFastProducer): (empty)

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

95 

96 def __init__( (empty)

97 self, 

98 parser: Optional["CustomCallable"], 

99 decoder: Optional["CustomCallable"], 

100 ) -> None: 

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

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

103 

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

105 default = AsyncConfluentParser() (empty)

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

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

108 

109 def connect( (empty)

110 self, 

111 producer: "AsyncConfluentProducer", 

112 serializer: Optional["SerializerProto"], 

113 ) -> None: 

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

115 self.serializer = serializer (empty)

116 

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

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

119 self._producer = EmptyProducerState() (empty)

120 

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

122 return bool(self._producer) (empty)

123 

124 async def ping(self, timeout: float) -> bool: (empty)

125 return await self._producer.ping(timeout=timeout) (empty)

126 

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

128 await self._producer.flush() 

129 

130 @override (empty)

131 async def publish( (empty)

132 self, 

133 cmd: "KafkaPublishCommand", 

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

135 """Publish a message to a topic.""" 

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

137 

138 headers_to_send = { (empty)

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

140 **cmd.headers_to_publish(), 

141 } 

142 

143 return await self._producer.producer.send( (empty)

144 topic=cmd.destination, 

145 value=message, 

146 key=cmd.key, 

147 partition=cmd.partition, 

148 timestamp_ms=cmd.timestamp_ms, 

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

150 no_confirm=cmd.no_confirm, 

151 ) 

152 

153 @override (empty)

154 async def publish_batch(self, cmd: "KafkaPublishCommand") -> None: (empty)

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

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

157 

158 headers_to_send = cmd.headers_to_publish() (empty)

159 

160 for message_position, msg in enumerate(cmd.batch_bodies): (empty)

161 message, content_type = encode_message(msg, serializer=self.serializer) (empty)

162 

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

164 final_headers = { (empty)

165 "content-type": content_type, 

166 **headers_to_send, 

167 } 

168 else: 

169 final_headers = headers_to_send.copy() 

170 

171 batch.append( (empty)

172 key=cmd.key_for(message_position), 

173 value=message, 

174 timestamp=cmd.timestamp_ms, 

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

176 ) 

177 

178 await self._producer.producer.send_batch( (empty)

179 batch, 

180 cmd.destination, 

181 partition=cmd.partition, 

182 no_confirm=cmd.no_confirm, 

183 )