def test_compressed_buffer_close(compression): records = MessageSetBuffer(io.BytesIO(), 100000, compression_type=compression) orig_msg = Message(b'foobar') records.append(1234, orig_msg) records.close() msgset = MessageSet.decode(records.buffer()) assert len(msgset) == 1 (offset, size, msg) = msgset[0] assert offset == 0 assert msg.is_compressed() msgset = msg.decompress() (offset, size, msg) = msgset[0] assert not msg.is_compressed() assert offset == 1234 assert msg == orig_msg # Closing again should work fine records.close() msgset = MessageSet.decode(records.buffer()) assert len(msgset) == 1 (offset, size, msg) = msgset[0] assert offset == 0 assert msg.is_compressed() msgset = msg.decompress() (offset, size, msg) = msgset[0] assert not msg.is_compressed() assert offset == 1234 assert msg == orig_msg
def test_decode_message_set_partial(): encoded = b''.join([ struct.pack('>q', 0), # Msg Offset struct.pack('>i', 18), # Msg Size struct.pack('>i', 1474775406), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k1', # Key struct.pack('>i', 2), # Length of value b'v1', # Value struct.pack('>q', 1), # Msg Offset struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size) struct.pack('>i', -16383415), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k2', # Key struct.pack('>i', 8), # Length of value b'ar', # Value (truncated) ]) msgs = MessageSet.decode(encoded, bytes_to_read=len(encoded)) assert len(msgs) == 2 msg1, msg2 = msgs returned_offset1, message1_size, decoded_message1 = msg1 returned_offset2, message2_size, decoded_message2 = msg2 assert returned_offset1 == 0 message1 = Message(b'v1', key=b'k1') message1.encode() assert decoded_message1 == message1 assert returned_offset2 is None assert message2_size is None assert decoded_message2 == PartialMessage()
def test_encode_message_set(): messages = [ Message(b'v1', key=b'k1'), Message(b'v2', key=b'k2') ] encoded = MessageSet.encode([(0, msg.encode()) for msg in messages]) expect = b''.join([ struct.pack('>q', 0), # MsgSet Offset struct.pack('>i', 18), # Msg Size struct.pack('>i', 1474775406), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k1', # Key struct.pack('>i', 2), # Length of value b'v1', # Value struct.pack('>q', 0), # MsgSet Offset struct.pack('>i', 18), # Msg Size struct.pack('>i', -16383415), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k2', # Key struct.pack('>i', 2), # Length of value b'v2', # Value ]) expect = struct.pack('>i', len(expect)) + expect assert encoded == expect
def test_encode_message_set(): messages = [ Message(b'v1', key=b'k1'), Message(b'v2', key=b'k2') ] encoded = MessageSet.encode([(0, msg.encode()) for msg in messages], size=False) expect = b''.join([ struct.pack('>q', 0), # MsgSet Offset struct.pack('>i', 18), # Msg Size struct.pack('>i', 1474775406), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k1', # Key struct.pack('>i', 2), # Length of value b'v1', # Value struct.pack('>q', 0), # MsgSet Offset struct.pack('>i', 18), # Msg Size struct.pack('>i', -16383415), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k2', # Key struct.pack('>i', 2), # Length of value b'v2', # Value ]) assert encoded == expect
def test_decode_fetch_response_partial(): encoded = b''.join([ Int32.encode(1), # Num Topics (Array) String('utf-8').encode('foobar'), Int32.encode(2), # Num Partitions (Array) Int32.encode(0), # Partition id Int16.encode(0), # Error Code Int64.encode(1234), # Highwater offset Int32.encode(52), # MessageSet size Int64.encode(0), # Msg Offset Int32.encode(18), # Msg Size struct.pack('>i', 1474775406), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k1', # Key struct.pack('>i', 2), # Length of value b'v1', # Value Int64.encode(1), # Msg Offset struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size) struct.pack('>i', -16383415), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k2', # Key struct.pack('>i', 8), # Length of value b'ar', # Value (truncated) Int32.encode(1), Int16.encode(0), Int64.encode(2345), Int32.encode(52), # MessageSet size Int64.encode(0), # Msg Offset Int32.encode(18), # Msg Size struct.pack('>i', 1474775406), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k1', # Key struct.pack('>i', 2), # Length of value b'v1', # Value Int64.encode(1), # Msg Offset struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size) struct.pack('>i', -16383415), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k2', # Key struct.pack('>i', 8), # Length of value b'ar', # Value (truncated) ]) resp = FetchResponse[0].decode(io.BytesIO(encoded)) assert len(resp.topics) == 1 topic, partitions = resp.topics[0] assert topic == 'foobar' assert len(partitions) == 2 m1 = MessageSet.decode( partitions[0][3], bytes_to_read=len(partitions[0][3])) assert len(m1) == 2 assert m1[1] == (None, None, PartialMessage())
def decode_message_set(cls, raw_data): messages = MessageSet.decode(raw_data, bytes_to_read=len(raw_data)) for offset, _, message in messages: if isinstance(message, kafka.protocol.message.Message) and message.is_compressed(): inner_messages = message.decompress() for (inner_offset, _msg_size, inner_msg) in inner_messages: yield kafka.structs.OffsetAndMessage(inner_offset, inner_msg) else: yield kafka.structs.OffsetAndMessage(offset, message)
def test_decode_fetch_response_partial(): encoded = b''.join([ Int32.encode(1), # Num Topics (Array) String('utf-8').encode('foobar'), Int32.encode(2), # Num Partitions (Array) Int32.encode(0), # Partition id Int16.encode(0), # Error Code Int64.encode(1234), # Highwater offset Int32.encode(52), # MessageSet size Int64.encode(0), # Msg Offset Int32.encode(18), # Msg Size struct.pack('>i', 1474775406), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k1', # Key struct.pack('>i', 2), # Length of value b'v1', # Value Int64.encode(1), # Msg Offset struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size) struct.pack('>i', -16383415), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k2', # Key struct.pack('>i', 8), # Length of value b'ar', # Value (truncated) Int32.encode(1), Int16.encode(0), Int64.encode(2345), Int32.encode(52), # MessageSet size Int64.encode(0), # Msg Offset Int32.encode(18), # Msg Size struct.pack('>i', 1474775406), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k1', # Key struct.pack('>i', 2), # Length of value b'v1', # Value Int64.encode(1), # Msg Offset struct.pack('>i', 24), # Msg Size (larger than remaining MsgSet size) struct.pack('>i', -16383415), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k2', # Key struct.pack('>i', 8), # Length of value b'ar', # Value (truncated) ]) resp = FetchResponse[0].decode(io.BytesIO(encoded)) assert len(resp.topics) == 1 topic, partitions = resp.topics[0] assert topic == 'foobar' assert len(partitions) == 2 m1 = MessageSet.decode(partitions[0][3], bytes_to_read=len(partitions[0][3])) assert len(m1) == 2 assert m1[1] == (None, None, PartialMessage())
def test_buffer_close(): records = MessageSetBuffer(io.BytesIO(), 100000) orig_msg = Message(b'foobar') records.append(1234, orig_msg) records.close() msgset = MessageSet.decode(records.buffer()) assert len(msgset) == 1 (offset, size, msg) = msgset[0] assert offset == 1234 assert msg == orig_msg # Closing again should work fine records.close() msgset = MessageSet.decode(records.buffer()) assert len(msgset) == 1 (offset, size, msg) = msgset[0] assert offset == 1234 assert msg == orig_msg
def encode_produce_request(cls, payloads=(), acks=1, timeout=1000): """ Encode a ProduceRequest struct Arguments: payloads: list of ProduceRequestPayload acks: How "acky" you want the request to be 1: written to disk by the leader 0: immediate response -1: waits for all replicas to be in sync timeout: Maximum time (in ms) the server will wait for replica acks. This is _not_ a socket timeout Returns: ProduceRequest """ if acks not in (1, 0, -1): raise ValueError('ProduceRequest acks (%s) must be 1, 0, -1' % acks) topics = [] for topic, topic_payloads in group_by_topic_and_partition( payloads).items(): topic_msgs = [] for partition, payload in topic_payloads.items(): partition_msgs = [] for msg in payload.messages: m = kafka.protocol.message.Message( msg.value, key=msg.key, magic=msg.magic, attributes=msg.attributes) partition_msgs.append((0, m.encode())) topic_msgs.append((partition, MessageSet.encode(partition_msgs, prepend_size=False))) topics.append((topic, topic_msgs)) return kafka.protocol.produce.ProduceRequest[0](required_acks=acks, timeout=timeout, topics=topics)
def encode_produce_request(cls, payloads=(), acks=1, timeout=1000): """ Encode a ProduceRequest struct Arguments: payloads: list of ProduceRequestPayload acks: How "acky" you want the request to be 1: written to disk by the leader 0: immediate response -1: waits for all replicas to be in sync timeout: Maximum time (in ms) the server will wait for replica acks. This is _not_ a socket timeout Returns: ProduceRequest """ if acks not in (1, 0, -1): raise ValueError('ProduceRequest acks (%s) must be 1, 0, -1' % acks) topics = [] for topic, topic_payloads in group_by_topic_and_partition(payloads).items(): topic_msgs = [] for partition, payload in topic_payloads.items(): partition_msgs = [] for msg in payload.messages: m = kafka.protocol.message.Message( msg.value, key=msg.key, magic=msg.magic, attributes=msg.attributes ) partition_msgs.append((0, m.encode())) topic_msgs.append((partition, MessageSet.encode(partition_msgs, prepend_size=False))) topics.append((topic, topic_msgs)) return kafka.protocol.produce.ProduceRequest[0]( required_acks=acks, timeout=timeout, topics=topics )
def test_decode_message_set(): encoded = b''.join([ struct.pack('>q', 0), # MsgSet Offset struct.pack('>i', 18), # Msg Size struct.pack('>i', 1474775406), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k1', # Key struct.pack('>i', 2), # Length of value b'v1', # Value struct.pack('>q', 1), # MsgSet Offset struct.pack('>i', 18), # Msg Size struct.pack('>i', -16383415), # CRC struct.pack('>bb', 0, 0), # Magic, flags struct.pack('>i', 2), # Length of key b'k2', # Key struct.pack('>i', 2), # Length of value b'v2', # Value ]) msgs = MessageSet.decode(encoded, bytes_to_read=len(encoded)) assert len(msgs) == 2 msg1, msg2 = msgs returned_offset1, message1_size, decoded_message1 = msg1 returned_offset2, message2_size, decoded_message2 = msg2 assert returned_offset1 == 0 message1 = Message(b'v1', key=b'k1') message1.encode() assert decoded_message1 == message1 assert returned_offset2 == 1 message2 = Message(b'v2', key=b'k2') message2.encode() assert decoded_message2 == message2
Array(('partition', Int32), ('messages', MessageSet)))))) class ProduceRequest_v1(Struct): API_KEY = API_KEY API_VERSION = 1 SCHEMA = ProduceRequest_v0.SCHEMA class ProduceRequest_v2(Struct): API_KEY = API_KEY API_VERSION = 2 SCHEMA = ProduceRequest_v1.SCHEMA send( ProduceRequest_v0(required_acks=0, timeout=1000, topics=[('topic1', [(0, MessageSet())])])) send( ProduceRequest_v1(required_acks=0, timeout=1000, topics=[('topic1', [(0, MessageSet())])])) send( ProduceRequest_v2(required_acks=0, timeout=1000, topics=[('topic1', [(0, MessageSet())])]))