OffsetAndTimestamp; FetchRequest v10

Author: dpkp

kafka/consumer/fetcher.py

@@ -18,7 +18,7 @@
 )
 from kafka.record import MemoryRecords
 from kafka.serializer import Deserializer
-from kafka.structs import TopicPartition, OffsetAndTimestamp
+from kafka.structs import TopicPartition, OffsetAndMetadata, OffsetAndTimestamp
 
 log = logging.getLogger(__name__)
 
@@ -28,7 +28,7 @@
 READ_COMMITTED = 1
 
 ConsumerRecord = collections.namedtuple("ConsumerRecord",
-    ["topic", "partition", "offset", "timestamp", "timestamp_type",
+    ["topic", "partition", "leader_epoch", "offset", "timestamp", "timestamp_type",
      "key", "value", "headers", "checksum", "serialized_key_size", "serialized_value_size", "serialized_header_size"])
 
 
@@ -198,9 +198,6 @@ def get_offsets_by_times(self, timestamps, timeout_ms):
         for tp in timestamps:
             if tp not in offsets:
                 offsets[tp] = None
-            else:
-                offset, timestamp = offsets[tp]
-                offsets[tp] = OffsetAndTimestamp(offset, timestamp)
         return offsets
 
     def beginning_offsets(self, partitions, timeout_ms):
@@ -215,7 +212,7 @@ def beginning_or_end_offset(self, partitions, timestamp, timeout_ms):
         timestamps = dict([(tp, timestamp) for tp in partitions])
         offsets = self._retrieve_offsets(timestamps, timeout_ms)
         for tp in timestamps:
-            offsets[tp] = offsets[tp][0]
+            offsets[tp] = offsets[tp].offset
         return offsets
 
     def _reset_offset(self, partition):
@@ -240,7 +237,7 @@ def _reset_offset(self, partition):
         offsets = self._retrieve_offsets({partition: timestamp})
 
         if partition in offsets:
-            offset = offsets[partition][0]
+            offset = offsets[partition].offset
 
             # we might lose the assignment while fetching the offset,
             # so check it is still active
@@ -261,8 +258,8 @@ def _retrieve_offsets(self, timestamps, timeout_ms=float("inf")):
                 available. Otherwise timestamp is treated as epoch milliseconds.
 
         Returns:
-            {TopicPartition: (int, int)}: Mapping of partition to
-                retrieved offset and timestamp. If offset does not exist for
+            {TopicPartition: OffsetAndTimestamp}: Mapping of partition to
+                retrieved offset, timestamp, and leader_epoch. If offset does not exist for
                 the provided timestamp, that partition will be missing from
                 this mapping.
         """
@@ -373,28 +370,29 @@ def _append(self, drained, part, max_records, update_offsets):
                 log.debug("Not returning fetched records for assigned partition"
                           " %s since it is no longer fetchable", tp)
 
-            elif fetch_offset == position:
+            elif fetch_offset == position.offset:
                 # we are ensured to have at least one record since we already checked for emptiness
                 part_records = part.take(max_records)
                 next_offset = part_records[-1].offset + 1
+                leader_epoch = part_records[-1].leader_epoch
 
                 log.log(0, "Returning fetched records at offset %d for assigned"
-                           " partition %s and update position to %s", position,
-                           tp, next_offset)
+                           " partition %s and update position to %s (leader epoch %s)", position.offset,
+                           tp, next_offset, leader_epoch)
 
                 for record in part_records:
                     drained[tp].append(record)
 
                 if update_offsets:
-                    self._subscriptions.assignment[tp].position = next_offset
+                    self._subscriptions.assignment[tp].position = OffsetAndMetadata(next_offset, b'', leader_epoch)
                 return len(part_records)
 
             else:
                 # these records aren't next in line based on the last consumed
                 # position, ignore them they must be from an obsolete request
                 log.debug("Ignoring fetched records for %s at offset %s since"
                           " the current position is %d", tp, part.fetch_offset,
-                          position)
+                          position.offset)
 
         part.discard()
         return 0
@@ -444,13 +442,13 @@ def _message_generator(self):
                     break
 
                 # Compressed messagesets may include earlier messages
-                elif msg.offset < self._subscriptions.assignment[tp].position:
+                elif msg.offset < self._subscriptions.assignment[tp].position.offset:
                     log.debug("Skipping message offset: %s (expecting %s)",
                               msg.offset,
-                              self._subscriptions.assignment[tp].position)
+                              self._subscriptions.assignment[tp].position.offset)
                     continue
 
-                self._subscriptions.assignment[tp].position = msg.offset + 1
+                self._subscriptions.assignment[tp].position = OffsetAndMetadata(msg.offset + 1, b'', -1)
                 yield msg
 
             self._next_partition_records = None
@@ -463,8 +461,9 @@ def _unpack_records(self, tp, records):
                 # Try DefaultsRecordBatch / message log format v2
                 # base_offset, last_offset_delta, and control batches
                 try:
-                    self._subscriptions.assignment[tp].last_offset_from_record_batch = batch.base_offset + \
-                                                                                       batch.last_offset_delta
+                    batch_offset = batch.base_offset + batch.last_offset_delta
+                    leader_epoch = batch.leader_epoch
+                    self._subscriptions.assignment[tp].last_offset_from_record_batch = batch_offset
                     # Control batches have a single record indicating whether a transaction
                     # was aborted or committed.
                     # When isolation_level is READ_COMMITTED (currently unsupported)
@@ -475,6 +474,7 @@ def _unpack_records(self, tp, records):
                         batch = records.next_batch()
                         continue
                 except AttributeError:
+                    leader_epoch = -1
                     pass
 
                 for record in batch:
@@ -491,7 +491,7 @@ def _unpack_records(self, tp, records):
                         len(h_key.encode("utf-8")) + (len(h_val) if h_val is not None else 0) for h_key, h_val in
                         headers) if headers else -1
                     yield ConsumerRecord(
-                        tp.topic, tp.partition, record.offset, record.timestamp,
+                        tp.topic, tp.partition, leader_epoch, record.offset, record.timestamp,
                         record.timestamp_type, key, value, headers, record.checksum,
                         key_size, value_size, header_size)
 
@@ -577,7 +577,9 @@ def _send_list_offsets_request(self, node_id, timestamps):
         version = self._client.api_version(ListOffsetsRequest, max_version=3)
         by_topic = collections.defaultdict(list)
         for tp, timestamp in six.iteritems(timestamps):
-            if version >= 1:
+            if version >= 4:
+                data = (tp.partition, leader_epoch, timestamp)
+            elif version >= 1:
                 data = (tp.partition, timestamp)
             else:
                 data = (tp.partition, timestamp, 1)
@@ -628,17 +630,18 @@ def _handle_list_offsets_response(self, future, response):
                             offset = UNKNOWN_OFFSET
                         else:
                             offset = offsets[0]
-                        log.debug("Handling v0 ListOffsetsResponse response for %s. "
-                                  "Fetched offset %s", partition, offset)
-                        if offset != UNKNOWN_OFFSET:
-                            timestamp_offset_map[partition] = (offset, None)
-                    else:
+                        timestamp = None
+                        leader_epoch = -1
+                    elif response.API_VERSION <= 3:
                         timestamp, offset = partition_info[2:]
-                        log.debug("Handling ListOffsetsResponse response for %s. "
-                                  "Fetched offset %s, timestamp %s",
-                                  partition, offset, timestamp)
-                        if offset != UNKNOWN_OFFSET:
-                            timestamp_offset_map[partition] = (offset, timestamp)
+                        leader_epoch = -1
+                    else:
+                        timestamp, offset, leader_epoch = partition_info[2:]
+                    log.debug("Handling ListOffsetsResponse response for %s. "
+                              "Fetched offset %s, timestamp %s, leader_epoch %s",
+                              partition, offset, timestamp, leader_epoch)
+                    if offset != UNKNOWN_OFFSET:
+                        timestamp_offset_map[partition] = OffsetAndTimestamp(offset, timestamp, leader_epoch)
                 elif error_type is Errors.UnsupportedForMessageFormatError:
                     # The message format on the broker side is before 0.10.0,
                     # we simply put None in the response.
@@ -686,7 +689,7 @@ def _create_fetch_requests(self):
         """
         # create the fetch info as a dict of lists of partition info tuples
         # which can be passed to FetchRequest() via .items()
-        version = self._client.api_version(FetchRequest, max_version=8)
+        version = self._client.api_version(FetchRequest, max_version=10)
         fetchable = collections.defaultdict(dict)
 
         for partition in self._fetchable_partitions():
@@ -695,12 +698,12 @@ def _create_fetch_requests(self):
             # advance position for any deleted compacted messages if required
             if self._subscriptions.assignment[partition].last_offset_from_record_batch:
                 next_offset_from_batch_header = self._subscriptions.assignment[partition].last_offset_from_record_batch + 1
-                if next_offset_from_batch_header > self._subscriptions.assignment[partition].position:
+                if next_offset_from_batch_header > self._subscriptions.assignment[partition].position.offset:
                     log.debug(
                         "Advance position for partition %s from %s to %s (last record batch location plus one)"
                         " to correct for deleted compacted messages and/or transactional control records",
-                        partition, self._subscriptions.assignment[partition].position, next_offset_from_batch_header)
-                    self._subscriptions.assignment[partition].position = next_offset_from_batch_header
+                        partition, self._subscriptions.assignment[partition].position.offset, next_offset_from_batch_header)
+                    self._subscriptions.assignment[partition].position = OffsetAndMetadata(next_offset_from_batch_header, b'', -1)
 
             position = self._subscriptions.assignment[partition].position
 
@@ -718,19 +721,28 @@ def _create_fetch_requests(self):
             if version < 5:
                 partition_info = (
                     partition.partition,
-                    position,
+                    position.offset,
                     self.config['max_partition_fetch_bytes']
                 )
+            elif version <= 8:
+                partition_info = (
+                    partition.partition,
+                    position.offset,
+                    -1, # log_start_offset is used internally by brokers / replicas only
+                    self.config['max_partition_fetch_bytes'],
+                )
             else:
                 partition_info = (
                     partition.partition,
-                    position,
+                    position.leader_epoch,
+                    position.offset,
                     -1, # log_start_offset is used internally by brokers / replicas only
                     self.config['max_partition_fetch_bytes'],
                 )
+
             fetchable[node_id][partition] = partition_info
             log.debug("Adding fetch request for partition %s at offset %d",
-                      partition, position)
+                      partition, position.offset)
 
         requests = {}
         for node_id, next_partitions in six.iteritems(fetchable):
@@ -778,7 +790,10 @@ def _create_fetch_requests(self):
 
             fetch_offsets = {}
             for tp, partition_data in six.iteritems(next_partitions):
-                offset = partition_data[1]
+                if version <= 8:
+                    offset = partition_data[1]
+                else:
+                    offset = partition_data[2]
                 fetch_offsets[tp] = offset
 
             requests[node_id] = (request, fetch_offsets)
@@ -807,7 +822,7 @@ def _handle_fetch_response(self, node_id, fetch_offsets, send_time, response):
                 tp = TopicPartition(topic, partition_data[0])
                 fetch_offset = fetch_offsets[tp]
                 completed_fetch = CompletedFetch(
-                    tp, fetch_offsets[tp],
+                    tp, fetch_offset,
                     response.API_VERSION,
                     partition_data[1:],
                     metric_aggregator
@@ -847,18 +862,18 @@ def _parse_fetched_data(self, completed_fetch):
                 # Note that the *response* may return a messageset that starts
                 # earlier (e.g., compressed messages) or later (e.g., compacted topic)
                 position = self._subscriptions.assignment[tp].position
-                if position is None or position != fetch_offset:
+                if position is None or position.offset != fetch_offset:
                     log.debug("Discarding fetch response for partition %s"
                               " since its offset %d does not match the"
                               " expected offset %d", tp, fetch_offset,
-                              position)
+                              position.offset)
                     return None
 
                 records = MemoryRecords(completed_fetch.partition_data[-1])
                 if records.has_next():
                     log.debug("Adding fetched record for partition %s with"
                               " offset %d to buffered record list", tp,
-                              position)
+                              position.offset)
                     unpacked = list(self._unpack_records(tp, records))
                     parsed_records = self.PartitionRecords(fetch_offset, tp, unpacked)
                     if unpacked:
@@ -889,10 +904,10 @@ def _parse_fetched_data(self, completed_fetch):
                 self._client.cluster.request_update()
             elif error_type is Errors.OffsetOutOfRangeError:
                 position = self._subscriptions.assignment[tp].position
-                if position is None or position != fetch_offset:
+                if position is None or position.offset != fetch_offset:
                     log.debug("Discarding stale fetch response for partition %s"
                               " since the fetched offset %d does not match the"
-                              " current offset %d", tp, fetch_offset, position)
+                              " current offset %d", tp, fetch_offset, position.offset)
                 elif self._subscriptions.has_default_offset_reset_policy():
                     log.info("Fetch offset %s is out of range for topic-partition %s", fetch_offset, tp)
                     self._subscriptions.need_offset_reset(tp)

kafka/consumer/group.py

@@ -17,7 +17,7 @@
 from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
 from kafka.metrics import MetricConfig, Metrics
 from kafka.protocol.list_offsets import OffsetResetStrategy
-from kafka.structs import TopicPartition
+from kafka.structs import OffsetAndMetadata, TopicPartition
 from kafka.version import __version__
 
 log = logging.getLogger(__name__)
@@ -737,11 +737,11 @@ def position(self, partition):
         if not isinstance(partition, TopicPartition):
             raise TypeError('partition must be a TopicPartition namedtuple')
         assert self._subscription.is_assigned(partition), 'Partition is not assigned'
-        offset = self._subscription.assignment[partition].position
+        position = self._subscription.assignment[partition].position
         if offset is None:
             self._update_fetch_positions([partition])
-            offset = self._subscription.assignment[partition].position
-        return offset
+            position = self._subscription.assignment[partition].position
+        return position.offset
 
     def highwater(self, partition):
         """Last known highwater offset for a partition.
@@ -1144,7 +1144,7 @@ def _message_generator_v2(self):
                     log.debug("Not returning fetched records for partition %s"
                               " since it is no longer fetchable", tp)
                     break
-                self._subscription.assignment[tp].position = record.offset + 1
+                self._subscription.assignment[tp].position = OffsetAndMetadata(record.offset + 1, b'', -1)
                 yield record
 
     def _message_generator(self):

kafka/consumer/subscription_state.py

@@ -319,7 +319,7 @@ def all_consumed_offsets(self):
         all_consumed = {}
         for partition, state in six.iteritems(self.assignment):
             if state.has_valid_position:
-                all_consumed[partition] = OffsetAndMetadata(state.position, b'', -1)
+                all_consumed[partition] = state.position
         return all_consumed
 
     def need_offset_reset(self, partition, offset_reset_strategy=None):
@@ -379,7 +379,7 @@ def __init__(self):
         self.paused = False # whether this partition has been paused by the user
         self.awaiting_reset = False # whether we are awaiting reset
         self.reset_strategy = None # the reset strategy if awaitingReset is set
-        self._position = None # offset exposed to the user
+        self._position = None # OffsetAndMetadata exposed to the user
         self.highwater = None
         self.drop_pending_record_batch = False
         # The last message offset hint available from a record batch with
@@ -388,6 +388,7 @@ def __init__(self):
 
     def _set_position(self, offset):
         assert self.has_valid_position, 'Valid position required'
+        assert isinstance(offset, OffsetAndMetadata)
         self._position = offset
 
     def _get_position(self):
@@ -403,7 +404,7 @@ def await_reset(self, strategy):
         self.has_valid_position = False
 
     def seek(self, offset):
-        self._position = offset
+        self._position = OffsetAndMetadata(offset, b'', -1)
         self.awaiting_reset = False
         self.reset_strategy = None
         self.has_valid_position = True

kafka/record/default_records.py

@@ -136,6 +136,10 @@ def __init__(self, buffer):
     def base_offset(self):
         return self._header_data[0]
 
+    @property
+    def leader_epoch(self):
+        return self._header_data[2]
+
     @property
     def magic(self):
         return self._header_data[3]

kafka/structs.py

@@ -66,9 +66,10 @@
 Keyword Arguments:
     offset (int): An offset
     timestamp (int): The timestamp associated to the offset
+    leader_epoch (int): The last known epoch from the leader / broker
 """
 OffsetAndTimestamp = namedtuple("OffsetAndTimestamp",
-    ["offset", "timestamp"])
+    ["offset", "timestamp", "leader_epoch"])
 
 MemberInformation = namedtuple("MemberInformation",
     ["member_id", "client_id", "client_host", "member_metadata", "member_assignment"])

test/test_consumer_integration.py

@@ -259,8 +259,8 @@ def test_kafka_consumer_offsets_search_many_partitions(kafka_consumer, kafka_pro
     })
 
     assert offsets == {
-        tp0: OffsetAndTimestamp(p0msg.offset, send_time),
-        tp1: OffsetAndTimestamp(p1msg.offset, send_time)
+        tp0: OffsetAndTimestamp(p0msg.offset, send_time, -1),
+        tp1: OffsetAndTimestamp(p1msg.offset, send_time, -1)
     }
 
     offsets = consumer.beginning_offsets([tp0, tp1])