Repository: Parsely/pykafka
Pull request #765 (2018-01-11 12:39:18)

add api_version to OffsetFetchRequest/OffsetFetchResponse
According to [kafka protocol description](https://kafka.apache.org/protocol#The_Messages_OffsetFetch)
OffsetFetch Request (version 2) fetches offsets for all _topics_ if topics array parameter is `Null`.
We use this functionality to get fetch offsets for all topics/partitions for particular consumer group, including offsets for unassigned/not_curently_consumed partitions.
Tested with kafka 0.11.0.


Patch:
diff --git a/pykafka/broker.py b/pykafka/broker.py
index 8becad8..c6d03f4 100644
--- a/pykafka/broker.py
+++ b/pykafka/broker.py
@@ -183,7 +183,8 @@ class Broker(object):
                    source_host=source_host,
                    source_port=source_port,
                    ssl_config=ssl_config,
-                   broker_version=broker_version)
+                   broker_version=broker_version,
+                   api_versions=api_versions)
 
     @property
     def connected(self):
@@ -432,8 +433,10 @@ class Broker(object):
         """
         if not self.offsets_channel_connected:
             self.connect_offsets_channel()
-        req = OffsetFetchRequest(consumer_group, partition_requests=preqs)
-        return self._offsets_channel_req_handler.request(req).get(OffsetFetchResponse)
+        request_class = OffsetFetchRequest.get_version_impl(self._api_versions)
+        response_class = OffsetFetchResponse.get_version_impl(self._api_versions)
+        req = request_class(consumer_group, partition_requests=preqs)
+        return self._offsets_channel_req_handler.request(req).get(response_class)
 
     ##########################
     #  Group Membership API  #
diff --git a/pykafka/protocol.py b/pykafka/protocol.py
index 0bcb29f..603a067 100644
--- a/pykafka/protocol.py
+++ b/pykafka/protocol.py
@@ -1481,8 +1481,13 @@ class OffsetFetchRequest(Request):
             TopicName => string
             Partition => int32
     """
+    API_VERSION = 0
     API_KEY = 9
 
+    @classmethod
+    def get_versions(cls):
+        return {0: OffsetFetchRequest, 1: OffsetFetchRequestV1, 2: OffsetFetchRequestV2}
+
     def __init__(self, consumer_group, partition_requests=[]):
         """Create a new offset fetch request
 
@@ -1495,6 +1500,9 @@ class OffsetFetchRequest(Request):
         for t in partition_requests:
             self._reqs[t.topic_name].append(t.partition_id)
 
+    def _reqs_len(self):
+        return len(self._reqs)
+
     def __len__(self):
         """Length of the serialized message, in bytes"""
         # Header + consumer group + len(topics)
@@ -1513,12 +1521,12 @@ class OffsetFetchRequest(Request):
         :rtype: :class:`bytearray`
         """
         output = bytearray(len(self))
-        self._write_header(output, api_version=1)
+        self._write_header(output, api_version=self.API_VERSION)
         offset = self.HEADER_LEN
         fmt = '!h%dsi' % len(self.consumer_group)
         struct.pack_into(fmt, output, offset,
                          len(self.consumer_group), self.consumer_group,
-                         len(self._reqs))
+                         self._reqs_len())
         offset += struct.calcsize(fmt)
         for topic_name, partitions in iteritems(self._reqs):
             fmt = '!h%dsi' % len(topic_name)
@@ -1532,6 +1540,18 @@ class OffsetFetchRequest(Request):
         return output
 
 
+class OffsetFetchRequestV1(OffsetFetchRequest):
+    API_VERSION = 1
+
+
+class OffsetFetchRequestV2(OffsetFetchRequestV1):
+    API_VERSION = 2
+
+    def _reqs_len(self):
+        # v2 allows a null array to select all topics
+        return len(self._reqs) or -1
+
+
 OffsetFetchPartitionResponse = namedtuple(
     'OffsetFetchPartitionResponse',
     ['offset', 'metadata', 'err']
@@ -1539,17 +1559,26 @@ OffsetFetchPartitionResponse = namedtuple(
 
 
 class OffsetFetchResponse(Response):
-    """An offset fetch response
+    """An offset fetch response v0
 
     Specification::
 
-        OffsetFetchResponse => [TopicName [Partition Offset Metadata ErrorCode]]
-            TopicName => string
-            Partition => int32
-            Offset => int64
-            Metadata => string
-            ErrorCode => int16
+    OffsetFetch Response (Version: 0) => [responses]
+        responses => topic [partition_responses]
+            topic => STRING
+            partition_responses => partition offset metadata error_code
+                partition => INT32
+                offset => INT64
+                metadata => NULLABLE_STRING
+                error_code => INT16
     """
+    API_VERSION = 0
+    API_KEY = 9
+
+    @classmethod
+    def get_versions(cls):
+        return {0: OffsetFetchResponse, 1: OffsetFetchResponseV1, 2: OffsetFetchResponseV2}
+
     def __init__(self, buff):
         """Deserialize into a new Response
 
@@ -1558,9 +1587,11 @@ class OffsetFetchResponse(Response):
         """
         fmt = '[S [iqSh ] ]'
         response = struct_helpers.unpack_from(fmt, buff, 0)
+        self._populate_partition_responses(response)
 
+    def _populate_partition_responses(self, partition_responses):
         self.topics = {}
-        for topic_name, partitions in response:
+        for topic_name, partitions in partition_responses:
             self.topics[topic_name] = {}
             for partition in partitions:
                 pres = OffsetFetchPartitionResponse(partition[1],
@@ -1569,6 +1600,47 @@ class OffsetFetchResponse(Response):
                 self.topics[topic_name][partition[0]] = pres
 
 
+class OffsetFetchResponseV1(OffsetFetchResponse):
+    """An offset fetch response v1 (all the same as v0)
+
+    Specification::
+
+    OffsetFetch Response (Version: 1) => [responses]
+        responses => topic [partition_responses]
+            topic => STRING
+            partition_responses => partition offset metadata error_code
+                partition => INT32
+                offset => INT64
+                metadata => NULLABLE_STRING
+                error_code => INT16
+    """
+    API_VERSION = 1
+
+
+class OffsetFetchResponseV2(OffsetFetchResponseV1):
+    """An offset fetch response v2
+
+    Specification::
+
+    OffsetFetch Response (Version: 2) => [responses] error_code
+        responses => topic [partition_responses]
+            topic => STRING
+            partition_responses => partition offset metadata error_code
+                partition => INT32
+                offset => INT64
+                metadata => NULLABLE_STRING
+                error_code => INT16
+        error_code => INT16 (new since v1)
+    """
+    API_VERSION = 2
+
+    def __init__(self, buff):
+        fmt = '[S [iqSh ] ] h'
+        response = struct_helpers.unpack_from(fmt, buff, 0)
+        partition_responses, self.err = response
+        self._populate_partition_responses(partition_responses)
+
+
 ###
 # Group Membership API
 ###
@@ -2412,7 +2484,7 @@ API_VERSIONS_080 = {
     6: ApiVersionsSpec(6, 0, 0),
     7: ApiVersionsSpec(7, 0, 0),
     8: ApiVersionsSpec(8, 0, 1),
-    9: ApiVersionsSpec(9, 0, 0),
+    9: ApiVersionsSpec(9, 0, 1),
     10: ApiVersionsSpec(10, 0, 0),
     11: ApiVersionsSpec(11, 0, 0),
     12: ApiVersionsSpec(12, 0, 0),
@@ -2431,7 +2503,7 @@ API_VERSIONS_090 = {
     6: ApiVersionsSpec(6, 0, 0),
     7: ApiVersionsSpec(7, 0, 0),
     8: ApiVersionsSpec(8, 0, 1),
-    9: ApiVersionsSpec(9, 0, 0),
+    9: ApiVersionsSpec(9, 0, 1),
     10: ApiVersionsSpec(10, 0, 0),
     11: ApiVersionsSpec(11, 0, 0),
     12: ApiVersionsSpec(12, 0, 0),


Test patch:
diff --git a/tests/pykafka/test_protocol.py b/tests/pykafka/test_protocol.py
index 96d324b..a3c675e 100644
--- a/tests/pykafka/test_protocol.py
+++ b/tests/pykafka/test_protocol.py
@@ -906,11 +906,19 @@ class TestOffsetCommitFetchAPI(unittest2.TestCase):
         self.assertEqual(
             msg,
             bytearray(
-                b'\x00\x00\x00.\x00\t\x00\x01\x00\x00\x00\x00\x00\x07pykafka'  # header
+                # header
+                b'\x00\x00\x00\x2e'  # len(buffer)
+                b'\x00\x09'  # ApiKey
+                b'\x00\x00'  # api version
+                b'\x00\x00\x00\x00'  # correlation id
+                b'\x00\x07'  # len(client id)
+                    b'pykafka'  # client id
+                # end header
+
                 b'\x00\x04'  # len(consumer group)
                     b'test'  # consumer group # noqa
                 b'\x00\x00\x00\x01'  # len(topics)
-                    b'\x00\t'  # len(topic name)
+                    b'\x00\x09'  # len(topic name)
                         b'testtopic'  # topic name
                     b'\x00\x00\x00\x01'  # len(partitions)
                         b'\x00\x00\x00\x00'  # partition
@@ -934,6 +942,49 @@ class TestOffsetCommitFetchAPI(unittest2.TestCase):
         self.assertEqual(response.topics[b'emmett.dummy'][0].offset, 1)
 
 
+class TestOffsetCommitFetchAPIV2(unittest2.TestCase):
+    maxDiff = None
+
+    def test_offset_fetch_request(self):
+        req = protocol.OffsetFetchRequestV2(b'test', partition_requests=[ ])
+        msg = req.get_bytes()
+        self.assertEqual(
+            msg,
+            bytearray(
+                # header
+                b'\x00\x00\x00\x1b'  # len(buffer)
+                b'\x00\x09'  # ApiKey
+                b'\x00\x02'  # api version
+                b'\x00\x00\x00\x00'  # correlation id
+                b'\x00\x07'  # len(client id)
+                    b'pykafka'  # client id
+                # end header
+
+                b'\x00\x04'  # len(consumer group)
+                    b'test'  # consumer group # noqa
+                b'\xff\xff\xff\xff'  # len(topics) = -1 for empty topics list
+            )
+        )
+
+    def test_offset_fetch_response(self):
+        response = protocol.OffsetFetchResponseV2(
+            buffer(
+                b'\x00\x00\x00\x01'  # len(topics)
+                    b'\x00\x0c'  # len(topic name) # noqa
+                        b'emmett.dummy'  # topic name
+                    b'\x00\x00\x00\x01'  # len(partitions)
+                        b'\x00\x00\x00\x00'  # partition
+                            b'\x00\x00\x00\x00\x00\x00\x00\x01'  # offset
+                            b'\x00\x00'  # len(metadata)
+                            b'\x00\x00'  # error code
+                b'\x00\x00'  # response level error_code
+            )
+        )
+        self.assertEqual(response.topics[b'emmett.dummy'][0].metadata, b'')
+        self.assertEqual(response.topics[b'emmett.dummy'][0].offset, 1)
+        self.assertEqual(response.err, 0)
+
+
 class TestGroupMembershipAPI(unittest2.TestCase):
     maxDiff = None