Coverage for faststream / kafka / response.py: 98%

1from typing import TYPE_CHECKING, Any, Union (empty)

2 

3from typing_extensions import override (empty)

4 

5from faststream.response.publish_type import PublishType (empty)

6from faststream.response.response import ( (empty)

7 BatchPublishCommand, 

8 PublishCommand, 

9 Response, 

10 extract_per_message_keys_and_bodies, 

11 key_for_index, 

12) 

13 

14if TYPE_CHECKING: (empty)

15 from faststream._internal.basic_types import SendableMessage 

16 

17 

18class KafkaResponse(Response): (empty)

19 """Kafka-specific response object for outgoing messages. (empty)

20 

21 Can be used in two ways: 

22 1. As a return value from handler to send a response message 

23 2. Directly in publish_batch() to set per-message attributes (key, headers, etc.) 

24 

25 For publish operations, consider using the more semantic alias `KafkaPublishMessage`. 

26 """ 

27 

28 def __init__( (empty)

29 self, 

30 body: "SendableMessage", 

31 *, 

32 headers: dict[str, Any] | None = None, 

33 correlation_id: str | None = None, 

34 timestamp_ms: int | None = None, 

35 key: bytes | None = None, 

36 ) -> None: 

37 super().__init__( (empty)

38 body=body, 

39 headers=headers, 

40 correlation_id=correlation_id, 

41 ) 

42 

43 self.timestamp_ms = timestamp_ms (empty)

44 self.key = key (empty)

45 

46 @override (empty)

47 def get_publish_key(self) -> bytes | None: (empty)

48 """Return the Kafka message key for publishing.""" 

49 return self.key (empty)

50 

51 @override (empty)

52 def as_publish_command(self) -> "KafkaPublishCommand": (empty)

53 return KafkaPublishCommand( (empty)

54 self.body, 

55 headers=self.headers, 

56 correlation_id=self.correlation_id, 

57 _publish_type=PublishType.PUBLISH, 

58 # Kafka specific 

59 topic="", 

60 key=self.key, 

61 timestamp_ms=self.timestamp_ms, 

62 ) 

63 

64 

65class KafkaPublishCommand(BatchPublishCommand): (empty)

66 def __init__( (empty)

67 self, 

68 message: "SendableMessage", 

69 /, 

70 *messages: "SendableMessage", 

71 topic: str, 

72 _publish_type: PublishType, 

73 key: bytes | Any | None = None, 

74 partition: int | None = None, 

75 timestamp_ms: int | None = None, 

76 headers: dict[str, str] | None = None, 

77 correlation_id: str | None = None, 

78 reply_to: str = "", 

79 no_confirm: bool = False, 

80 timeout: float = 0.5, 

81 ) -> None: 

82 super().__init__( (empty)

83 message, 

84 *messages, 

85 destination=topic, 

86 reply_to=reply_to, 

87 correlation_id=correlation_id, 

88 headers=headers, 

89 _publish_type=_publish_type, 

90 ) 

91 

92 self.key = key (empty)

93 self.partition = partition (empty)

94 self.timestamp_ms = timestamp_ms (empty)

95 self.no_confirm = no_confirm (empty)

96 

97 # request option 

98 self.timeout = timeout (empty)

99 

100 # per-message keys support 

101 keys, normalized = extract_per_message_keys_and_bodies(self.batch_bodies) (empty)

102 if normalized is not None: (empty)

103 self.batch_bodies = normalized (empty)

104 self._per_message_keys = keys (empty)

105 

106 @classmethod (empty)

107 def from_cmd( (empty)

108 cls, 

109 cmd: Union["PublishCommand", "KafkaPublishCommand"], 

110 *, 

111 batch: bool = False, 

112 ) -> "KafkaPublishCommand": 

113 if isinstance(cmd, KafkaPublishCommand): (empty)

114 # NOTE: Should return a copy probably. 

115 return cmd (empty)

116 

117 body, extra_bodies = cls._parse_bodies(cmd.body, batch=batch) (empty)

118 

119 return cls( (empty)

120 body, 

121 *extra_bodies, 

122 topic=cmd.destination, 

123 correlation_id=cmd.correlation_id, 

124 headers=cmd.headers, 

125 reply_to=cmd.reply_to, 

126 _publish_type=cmd.publish_type, 

127 ) 

128 

129 def key_for(self, index: int) -> Any | None: (empty)

130 return key_for_index(self._per_message_keys, self.key, index) (empty)

131 

132 def headers_to_publish(self) -> dict[str, str]: (empty)

133 headers = {} (empty)

134 

135 if self.correlation_id: 135 ↛ 138line 135 didn't jump to line 138 because the condition on line 135 was always true(empty)

136 headers["correlation_id"] = self.correlation_id (empty)

137 

138 if self.reply_to: (empty)

139 headers["reply_to"] = self.reply_to (empty)

140 

141 return headers | self.headers (empty)

142 

143 

144# Semantic alias for publish operations 

145# More intuitive name when using in publish_batch() rather than as handler return value 

146KafkaPublishMessage = KafkaResponse (empty)