diff options
-rw-r--r-- | kafka/consumer/fetcher.py | 13 | ||||
-rw-r--r-- | test/test_fetcher.py | 40 |
2 files changed, 51 insertions, 2 deletions
diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py index f9fcb37..c9bbb97 100644 --- a/kafka/consumer/fetcher.py +++ b/kafka/consumer/fetcher.py @@ -835,12 +835,21 @@ class Fetcher(six.Iterator): return parsed_records - class PartitionRecords(six.Iterator): + class PartitionRecords(object): def __init__(self, fetch_offset, tp, messages): self.fetch_offset = fetch_offset self.topic_partition = tp self.messages = messages - self.message_idx = 0 + # When fetching an offset that is in the middle of a + # compressed batch, we will get all messages in the batch. + # But we want to start 'take' at the fetch_offset + for i, msg in enumerate(messages): + if msg.offset == fetch_offset: + self.message_idx = i + break + else: + self.message_idx = 0 + self.messages = None # For truthiness evaluation we need to define __len__ or __nonzero__ def __len__(self): diff --git a/test/test_fetcher.py b/test/test_fetcher.py index 429071a..4547222 100644 --- a/test/test_fetcher.py +++ b/test/test_fetcher.py @@ -498,3 +498,43 @@ def test__parse_fetched_data__out_of_range(fetcher, topic, mocker): partition_record = fetcher._parse_fetched_data(completed_fetch) assert partition_record is None assert fetcher._subscriptions.assignment[tp].awaiting_reset is True + + +def test_partition_records_offset(): + """Test that compressed messagesets are handle correctly + when fetch offset is in the middle of the message list + """ + batch_start = 120 + batch_end = 130 + fetch_offset = 123 + tp = TopicPartition('foo', 0) + messages = [ConsumerRecord(tp.topic, tp.partition, i, + None, None, 'key', 'value', 'checksum', 0, 0) + for i in range(batch_start, batch_end)] + records = Fetcher.PartitionRecords(fetch_offset, None, messages) + assert len(records) > 0 + msgs = records.take(1) + assert msgs[0].offset == 123 + assert records.fetch_offset == 124 + msgs = records.take(2) + assert len(msgs) == 2 + assert len(records) > 0 + records.discard() + assert len(records) == 0 + + +def test_partition_records_empty(): + records = Fetcher.PartitionRecords(0, None, []) + assert len(records) == 0 + + +def test_partition_records_no_fetch_offset(): + batch_start = 0 + batch_end = 100 + fetch_offset = 123 + tp = TopicPartition('foo', 0) + messages = [ConsumerRecord(tp.topic, tp.partition, i, + None, None, 'key', 'value', 'checksum', 0, 0) + for i in range(batch_start, batch_end)] + records = Fetcher.PartitionRecords(fetch_offset, None, messages) + assert len(records) == 0 |