Add leader_epoch to OffsetAndMetadata struct

Author: dpkp

kafka/admin/client.py

@@ -1353,7 +1353,7 @@ def _list_consumer_group_offsets_send_request(self, group_id,
         Returns:
             A message future
         """
-        version = self._client.api_version(OffsetFetchRequest, max_version=3)
+        version = self._client.api_version(OffsetFetchRequest, max_version=5)
         if version <= 3:
             if partitions is None:
                 if version <= 1:
@@ -1386,7 +1386,7 @@ def _list_consumer_group_offsets_process_response(self, response):
             A dictionary composed of TopicPartition keys and
             OffsetAndMetadata values.
         """
-        if response.API_VERSION <= 3:
+        if response.API_VERSION <= 5:
 
             # OffsetFetchResponse_v1 lacks a top-level error_code
             if response.API_VERSION > 1:
@@ -1401,13 +1401,18 @@ def _list_consumer_group_offsets_process_response(self, response):
             # OffsetAndMetadata values--this is what the Java AdminClient returns
             offsets = {}
             for topic, partitions in response.topics:
-                for partition, offset, metadata, error_code in partitions:
+                for partition_data in partitions:
+                    if response.API_VERSION <= 4:
+                        partition, offset, metadata, error_code = partition_data
+                        leader_epoch = -1
+                    else:
+                        partition, offset, leader_epoch, metadata, error_code = partition_data
                     error_type = Errors.for_code(error_code)
                     if error_type is not Errors.NoError:
                         raise error_type(
                             "Unable to fetch consumer group offsets for topic {}, partition {}"
                             .format(topic, partition))
-                    offsets[TopicPartition(topic, partition)] = OffsetAndMetadata(offset, metadata)
+                    offsets[TopicPartition(topic, partition)] = OffsetAndMetadata(offset, metadata, leader_epoch)
         else:
             raise NotImplementedError(
                 "Support for OffsetFetchResponse_v{} has not yet been added to KafkaAdminClient."
@@ -1439,7 +1444,7 @@ def list_consumer_group_offsets(self, group_id, group_coordinator_id=None,
 
         Returns:
             dictionary: A dictionary with TopicPartition keys and
-            OffsetAndMetada values. Partitions that are not specified and for
+            OffsetAndMetadata values. Partitions that are not specified and for
             which the group_id does not have a recorded offset are omitted. An
             offset value of `-1` indicates the group_id has no offset for that
             TopicPartition. A `-1` can only happen for partitions that are

kafka/coordinator/consumer.py

@@ -649,7 +649,7 @@ def _send_offset_commit_request(self, offsets):
                     topic, [(
                         partition,
                         offset.offset,
-                        -1, # leader_epoch
+                        offset.leader_epoch,
                         offset.metadata
                     ) for partition, offset in six.iteritems(partitions)]
                 ) for topic, partitions in six.iteritems(offset_data)]
@@ -815,7 +815,6 @@ def _handle_offset_fetch_response(self, future, response):
                 else:
                     metadata, error_code = partition_data[2:]
                     leader_epoch = -1
-                # TODO: save leader_epoch!
                 tp = TopicPartition(topic, partition)
                 error_type = Errors.for_code(error_code)
                 if error_type is not Errors.NoError:
@@ -842,7 +841,7 @@ def _handle_offset_fetch_response(self, future, response):
                 elif offset >= 0:
                     # record the position with the offset
                     # (-1 indicates no committed offset to fetch)
-                    offsets[tp] = OffsetAndMetadata(offset, metadata)
+                    offsets[tp] = OffsetAndMetadata(offset, metadata, leader_epoch)
                 else:
                     log.debug("Group %s has no committed offset for partition"
                               " %s", self.group_id, tp)

kafka/structs.py

@@ -48,17 +48,17 @@
 """The Kafka offset commit API
 
 The Kafka offset commit API allows users to provide additional metadata
-(in the form of a string) when an offset is committed. This can be useful
+(in the form of raw bytes) when an offset is committed. This can be useful
 (for example) to store information about which node made the commit,
 what time the commit was made, etc.
 
 Keyword Arguments:
     offset (int): The offset to be committed
-    metadata (str): Non-null metadata
+    metadata (bytes): Non-null metadata
+    leader_epoch (int): The last known epoch from the leader / broker
 """
 OffsetAndMetadata = namedtuple("OffsetAndMetadata",
-    # TODO add leaderEpoch: OffsetAndMetadata(offset, leaderEpoch, metadata)
-    ["offset", "metadata"])
+    ["offset", "metadata", "leader_epoch"])
 
 
 """An offset and timestamp tuple

test/test_coordinator.py

@@ -230,13 +230,13 @@ def test_need_rejoin(coordinator):
 def test_refresh_committed_offsets_if_needed(mocker, coordinator):
     mocker.patch.object(ConsumerCoordinator, 'fetch_committed_offsets',
                         return_value = {
-                            TopicPartition('foobar', 0): OffsetAndMetadata(123, b''),
-                            TopicPartition('foobar', 1): OffsetAndMetadata(234, b'')})
+                            TopicPartition('foobar', 0): OffsetAndMetadata(123, b'', -1),
+                            TopicPartition('foobar', 1): OffsetAndMetadata(234, b'', -1)})
     coordinator._subscription.assign_from_user([TopicPartition('foobar', 0)])
     assert coordinator._subscription.needs_fetch_committed_offsets is True
     coordinator.refresh_committed_offsets_if_needed()
     assignment = coordinator._subscription.assignment
-    assert assignment[TopicPartition('foobar', 0)].committed == OffsetAndMetadata(123, b'')
+    assert assignment[TopicPartition('foobar', 0)].committed == OffsetAndMetadata(123, b'', -1)
     assert TopicPartition('foobar', 1) not in assignment
     assert coordinator._subscription.needs_fetch_committed_offsets is False
 
@@ -303,8 +303,8 @@ def test_close(mocker, coordinator):
 @pytest.fixture
 def offsets():
     return {
-        TopicPartition('foobar', 0): OffsetAndMetadata(123, b''),
-        TopicPartition('foobar', 1): OffsetAndMetadata(234, b''),
+        TopicPartition('foobar', 0): OffsetAndMetadata(123, b'', -1),
+        TopicPartition('foobar', 1): OffsetAndMetadata(234, b'', -1),
     }
 
 

test/test_fetcher.py

@@ -138,7 +138,7 @@ def test_update_fetch_positions(fetcher, topic, mocker):
     fetcher._reset_offset.reset_mock()
     fetcher._subscriptions.need_offset_reset(partition)
     fetcher._subscriptions.assignment[partition].awaiting_reset = False
-    fetcher._subscriptions.assignment[partition].committed = OffsetAndMetadata(123, b'')
+    fetcher._subscriptions.assignment[partition].committed = OffsetAndMetadata(123, b'', -1)
     mocker.patch.object(fetcher._subscriptions, 'seek')
     fetcher.update_fetch_positions([partition])
     assert fetcher._reset_offset.call_count == 0