Coverage for faststream / confluent / parser.py: 90%

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

2 

3from faststream.message import StreamMessage, decode_message (empty)

4 

5from .message import FAKE_CONSUMER, KafkaMessage (empty)

6 

7if TYPE_CHECKING: (empty)

8 from collections.abc import Sequence 

9 

10 from confluent_kafka import Message 

11 

12 from faststream._internal.basic_types import DecodedMessage 

13 

14 from .message import ConsumerProtocol 

15 

16 # Type of headers returned by confluent_kafka Message.headers() 

17 _HeadersInput = ( 

18 dict[str, str | bytes | None] 

19 | list[tuple[str, str | bytes | None]] 

20 | tuple[tuple[str, str | bytes | None], ...] 

21 ) 

22 

23 

24class AsyncConfluentParser: (empty)

25 """A class to parse Kafka messages.""" 

26 

27 def __init__(self, is_manual: bool = False) -> None: (empty)

28 self.is_manual = is_manual (empty)

29 self._consumer: ConsumerProtocol = FAKE_CONSUMER (empty)

30 

31 def _setup(self, consumer: "ConsumerProtocol") -> None: (empty)

32 self._consumer = consumer (empty)

33 

34 async def parse_message( (empty)

35 self, 

36 message: "Message", 

37 ) -> KafkaMessage: 

38 """Parses a Kafka message.""" 

39 headers = _parse_msg_headers(cast("_HeadersInput", message.headers() or ())) (empty)

40 

41 body = message.value() or b"" (empty)

42 offset = message.offset() (empty)

43 _, timestamp = message.timestamp() (empty)

44 

45 return KafkaMessage( (empty)

46 body=body, 

47 headers=headers, 

48 reply_to=headers.get("reply_to", ""), 

49 content_type=headers.get("content-type"), 

50 message_id=f"{offset}-{timestamp}", 

51 correlation_id=headers.get("correlation_id"), 

52 raw_message=message, 

53 consumer=self._consumer, 

54 is_manual=self.is_manual, 

55 ) 

56 

57 async def parse_batch( (empty)

58 self, 

59 message: tuple["Message", ...], (empty)

60 ) -> KafkaMessage: (empty)

61 """Parses a batch of messages from a Kafka consumer.""" 

62 body: list[Any] = [] (empty)

63 batch_headers: list[dict[str, str]] = [] (empty)

64 

65 first = message[0] (empty)

66 last = message[-1] (empty)

67 

68 for m in message: (empty)

69 body.append(m.value() or b"") (empty)

70 batch_headers.append( (empty)

71 _parse_msg_headers(cast("_HeadersInput", m.headers() or ())) (empty)

72 ) 

73 

74 headers = next(iter(batch_headers), {}) (empty)

75 

76 _, first_timestamp = first.timestamp() (empty)

77 

78 return KafkaMessage( (empty)

79 body=body, (empty)

80 headers=headers, (empty)

81 batch_headers=batch_headers, (empty)

82 reply_to=headers.get("reply_to", ""), (empty)

83 content_type=headers.get("content-type"), (empty)

84 message_id=f"{first.offset()}-{last.offset()}-{first_timestamp}", (empty)

85 correlation_id=headers.get("correlation_id"), (empty)

86 raw_message=message, (empty)

87 consumer=self._consumer, (empty)

88 is_manual=self.is_manual, (empty)

89 ) 

90 

91 async def decode_message( (empty)

92 self, 

93 msg: "StreamMessage[Message]", 

94 ) -> "DecodedMessage": 

95 """Decodes a message.""" 

96 return decode_message(msg) (empty)

97 

98 async def decode_batch( (empty)

99 self, 

100 msg: "StreamMessage[tuple[Message, ...]]", (empty)

101 ) -> "DecodedMessage": (empty)

102 """Decode a batch of messages.""" 

103 return [decode_message(await self.parse_message(m)) for m in msg.raw_message] (empty)

104 

105 

106def _parse_msg_headers(headers: "_HeadersInput") -> dict[str, str]: (empty)

107 if isinstance(headers, dict): 107 ↛ 108line 107 didn't jump to line 108 because the condition on line 107 was never true(empty)

108 seq: Sequence[tuple[str, bytes | str | None]] = list(headers.items()) 

109 else: 

110 seq = headers (empty)

111 return { (empty)

112 i: (j if isinstance(j, str) else (j.decode() if j is not None else "")) 

113 for i, j in seq 

114 if j is not None 

115 }