Repository: Parsely/pykafka
Pull request #621 (2016-12-02 19:56:01)

support 0.10's administrative API
This pull request builds on #620 to provide support for Kafka 0.10's administrative API. It adds support for `ListGroupsRequest` and `DescribeGroupsRequest`, as well as a convenience function `Cluster.get_managed_group_descriptions`.

These changes require a minor version increment.


Patch:
diff --git a/pykafka/broker.py b/pykafka/broker.py
index 562908a..76b9f1a 100644
--- a/pykafka/broker.py
+++ b/pykafka/broker.py
@@ -28,7 +28,8 @@ from .protocol import (
     MetadataResponse, OffsetCommitRequest, OffsetCommitResponse, OffsetFetchRequest,
     OffsetFetchResponse, ProduceResponse, JoinGroupRequest, JoinGroupResponse,
     SyncGroupRequest, SyncGroupResponse, HeartbeatRequest, HeartbeatResponse,
-    LeaveGroupRequest, LeaveGroupResponse)
+    LeaveGroupRequest, LeaveGroupResponse, ListGroupsRequest, ListGroupsResponse,
+    DescribeGroupsRequest, DescribeGroupsResponse)
 from .utils.compat import range, iteritems, get_bytes
 
 log = logging.getLogger(__name__)
@@ -466,3 +467,20 @@ class Broker(object):
             HeartbeatRequest(consumer_group, generation_id, member_id))
         self._handler.sleep()
         return future.get(HeartbeatResponse)
+
+    ########################
+    #  Administrative API  #
+    ########################
+    def list_groups(self):
+        """Send a ListGroupsRequest"""
+        future = self._req_handler.request(ListGroupsRequest())
+        return future.get(ListGroupsResponse)
+
+    def describe_groups(self, group_ids):
+        """Send a DescribeGroupsRequest
+
+        :param group_ids: A sequence of group identifiers for which to return descriptions
+        :type group_ids: sequence of str
+        """
+        future = self._req_handler.request(DescribeGroupsRequest(group_ids))
+        return future.get(DescribeGroupsResponse)
diff --git a/pykafka/cluster.py b/pykafka/cluster.py
index 11b49cc..7e6a7d7 100644
--- a/pykafka/cluster.py
+++ b/pykafka/cluster.py
@@ -371,6 +371,19 @@ class Cluster(object):
                 #       needed.
                 raise Exception('Broker host/port change detected! %s', broker)
 
+    def get_managed_group_descriptions(self):
+        """Return detailed descriptions of all managed consumer groups on this cluster
+
+        This function only returns descriptions for consumer groups created via the
+        Group Management API, which pykafka refers to as :class:`ManagedBalancedConsumer`s
+        """
+        descriptions = {}
+        for broker in itervalues(self.brokers):
+            res = broker.describe_groups([group.group_id for group
+                                          in itervalues(broker.list_groups().groups)])
+            descriptions.update(res.groups)
+        return descriptions
+
     def get_offset_manager(self, consumer_group):
         log.warning("WARNING: Cluster.get_offset_manager is deprecated since pykafka "
                     "2.3.0. Instead, use Cluster.get_group_coordinator.")
diff --git a/pykafka/protocol.py b/pykafka/protocol.py
index 2541eba..2829453 100644
--- a/pykafka/protocol.py
+++ b/pykafka/protocol.py
@@ -1615,3 +1615,156 @@ class LeaveGroupResponse(Response):
         fmt = 'h'
         response = struct_helpers.unpack_from(fmt, buff, 0)
         self.error_code = response[0]
+
+
+###
+# Administrative API
+###
+class ListGroupsRequest(Request):
+    """A list groups request
+
+    Specification::
+
+    ListGroupsRequest =>
+    """
+    @property
+    def API_KEY(self):
+        """API_KEY for this request, from the Kafka docs"""
+        return 16
+
+    def get_bytes(self):
+        """Create a new list group request"""
+        output = bytearray(len(self))
+        self._write_header(output)
+        return output
+
+    def __len__(self):
+        """Length of the serialized message, in bytes"""
+        return self.HEADER_LEN
+
+
+GroupListing = namedtuple(
+    'GroupListing',
+    ['group_id', 'protocol_type']
+)
+
+
+class ListGroupsResponse(Response):
+    """A list groups response
+
+    Specification::
+
+    ListGroupsResponse => ErrorCode Groups
+      ErrorCode => int16
+      Groups => [GroupId ProtocolType]
+        GroupId => string
+        ProtocolType => string
+    """
+    def __init__(self, buff):
+        """Deserialize into a new Response
+
+        :param buff: Serialized message
+        :type buff: :class:`bytearray`
+        """
+        fmt = 'h [SS]'
+        response = struct_helpers.unpack_from(fmt, buff, 0)
+
+        self.error = response[0]
+        self.groups = {}
+        for group_info in response[1]:
+            listing = GroupListing(*group_info)
+            self.groups[listing.group_id] = listing
+
+
+class DescribeGroupsRequest(Request):
+    """A describe groups request
+
+    Specification::
+
+    DescribeGroupsRequest => [GroupId]
+      GroupId => string
+    """
+    def __init__(self, group_ids):
+        self.group_ids = group_ids
+
+    @property
+    def API_KEY(self):
+        """API_KEY for this request, from the Kafka docs"""
+        return 15
+
+    def get_bytes(self):
+        """Create a new list group request"""
+        output = bytearray(len(self))
+        self._write_header(output)
+        offset = self.HEADER_LEN
+        fmt = '!i'
+        struct.pack_into(fmt, output, offset, len(self.group_ids))
+        offset += struct.calcsize(fmt)
+        for group_id in self.group_ids:
+            fmt = '!h%ds' % len(group_id)
+            struct.pack_into(fmt, output, offset, len(group_id), group_id)
+            offset += struct.calcsize(fmt)
+        return output
+
+    def __len__(self):
+        """Length of the serialized message, in bytes"""
+        # header + len(group_ids)
+        size = self.HEADER_LEN + 4
+        for group_id in self.group_ids:
+            # len(group_id) + group_id
+            size += 2 + len(group_id)
+        return size
+
+
+GroupMember = namedtuple(
+    'GroupMember',
+    ['member_id', 'client_id', 'client_host', 'member_metadata', 'member_assignment']
+)
+
+
+DescribeGroupResponse = namedtuple(
+    'DescribeGroupResponse',
+    ['error_code', 'group_id', 'state', 'protocol_type', 'protocol', 'members']
+)
+
+
+class DescribeGroupsResponse(Response):
+    """A describe groups response
+
+    Specification::
+
+
+    DescribeGroupsResponse => [ErrorCode GroupId State ProtocolType Protocol Members]
+      ErrorCode => int16
+      GroupId => string
+      State => string
+      ProtocolType => string
+      Protocol => string
+      Members => [MemberId ClientId ClientHost MemberMetadata MemberAssignment]
+        MemberId => string
+        ClientId => string
+        ClientHost => string
+        MemberMetadata => bytes
+        MemberAssignment => bytes
+    """
+    def __init__(self, buff):
+        """Deserialize into a new Response
+
+        :param buff: Serialized message
+        :type buff: :class:`bytearray`
+        """
+        fmt = '[hSSSS [SSSYY ] ]'
+        response = struct_helpers.unpack_from(fmt, buff, 0)
+
+        self.groups = {}
+        for group_info in response:
+            members = {}
+            for member_info in group_info[5]:
+                member_metadata = ConsumerGroupProtocolMetadata.from_bytestring(
+                    member_info[3])
+                member_assignment = MemberAssignment.from_bytestring(member_info[4])
+                member = GroupMember(*(member_info[:3] + (member_metadata,
+                                                          member_assignment)))
+                members[member.member_id] = member
+            group = DescribeGroupResponse(*(group_info[:5] + (members,)))
+            self.groups[group.group_id] = group


Test patch:
diff --git a/tests/pykafka/test_protocol.py b/tests/pykafka/test_protocol.py
index 8117daa..2442cbf 100644
--- a/tests/pykafka/test_protocol.py
+++ b/tests/pykafka/test_protocol.py
@@ -722,5 +722,100 @@ class TestGroupMembershipAPI(unittest2.TestCase):
         self.assertEqual(response.error_code, 0)
 
 
+class TestAdministrativeAPI(unittest2.TestCase):
+    maxDiff = None
+
+    def test_list_groups_request(self):
+        req = protocol.ListGroupsRequest()
+        msg = req.get_bytes()
+        self.assertEqual(
+            msg,
+            bytearray(
+                b'\x00\x00\x00\x11\x00\x10\x00\x00\x00\x00\x00\x00\x00\x07pykafka'  # header
+            )
+        )
+
+    def test_list_groups_response(self):
+        response = protocol.ListGroupsResponse(
+            bytearray(
+                b'\x00\x00'  # error code
+                b'\x00\x00\x00\x01'  # len(groups)
+                    b'\x00\t'  # len(group_id)
+                        b'testgroup'  # group_id
+                    b'\x00\x08'  # len(protocol_type)
+                        b'consumer'  # protocol_type
+            )
+        )
+        self.assertEqual(len(response.groups), 1)
+        group = response.groups[b"testgroup"]
+        self.assertEqual(group.group_id, b"testgroup")
+        self.assertEqual(group.protocol_type, b"consumer")
+
+    def test_describe_groups_request(self):
+        req = protocol.DescribeGroupsRequest([b'testgroup'])
+        msg = req.get_bytes()
+        self.assertEqual(
+            msg,
+            bytearray(
+                b'\x00\x00\x00 \x00\x0f\x00\x00\x00\x00\x00\x00\x00\x07pykafka'  # header
+                b'\x00\x00\x00\x01'  # len(group_ids)
+                    b'\x00\t'  # len(group_id)
+                        b'testgroup'  # group_id
+            )
+        )
+
+    def test_describe_groups_response(self):
+        response = protocol.DescribeGroupsResponse(
+            bytearray(
+                b'\x00\x00\x00\x01'  # len(groups)
+                    b'\x00\x00'  # error code
+                    b'\x00\t'  # len(group_id)
+                        b'testgroup'  # group_id
+                    b'\x00\x06'  # len(state)
+                        b'Stable'  # state
+                    b'\x00\x08'  # len(protocol_type)
+                        b'consumer'  # protocol_type
+                    b'\x00\x19'  # len(protocol)
+                        b'pykafkaassignmentstrategy'  # protocol
+                    b'\x00\x00\x00\x01'  # len(members)
+                        b'\x00,'  # len(member_id)
+                            b'pykafka-d42426fb-c295-4cd9-b585-6dd79daf3afe'  # member_id
+                        b'\x00\x07'  # len(client_id)
+                            b'pykafka'  # client_id
+                        b'\x00\n'  # len(client_host)
+                            b'/127.0.0.1'  # client_host
+                        b'\x00\x00\x00"'  # len(member_metadata)
+                            b'\x00\x00\x00\x00\x00\x01\x00\ndummytopic\x00\x00\x00\x0ctestuserdata'
+                        b'\x00\x00\x00H'  # len(member_assignment)
+                            b'\x00\x01\x00\x00\x00\x01\x00\x14testtopic_replicated'  # member_assignment
+                            b'\x00\x00\x00\n\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00'
+                            b'\x02\x00\x00\x00\x08\x00\x00\x00\x03\x00\x00\x00\t'
+                            b'\x00\x00\x00\x04\x00\x00\x00\x05\x00\x00\x00\x06'
+                            b'\x00\x00\x00\x07\x00\x00\x00\x00'
+            )
+        )
+        self.assertTrue(b'testgroup' in response.groups)
+        group_response = response.groups[b'testgroup']
+        self.assertEqual(group_response.error_code, 0)
+        self.assertEqual(group_response.group_id, b'testgroup')
+        self.assertEqual(group_response.state, b'Stable')
+        self.assertEqual(group_response.protocol_type, b'consumer')
+        self.assertEqual(group_response.protocol, b'pykafkaassignmentstrategy')
+        member_id = b'pykafka-d42426fb-c295-4cd9-b585-6dd79daf3afe'
+        self.assertTrue(member_id in group_response.members)
+        member = group_response.members[member_id]
+        self.assertEqual(member.member_id, member_id)
+        self.assertEqual(member.client_id, b'pykafka')
+        self.assertEqual(member.client_host, b'/127.0.0.1')
+        metadata = member.member_metadata
+        self.assertEqual(metadata.version, 0)
+        self.assertEqual(metadata.topic_names, [b"dummytopic"])
+        self.assertEqual(metadata.user_data, b"testuserdata")
+        assignment = member.member_assignment
+        self.assertEqual(assignment.version, 1)
+        self.assertEqual(assignment.partition_assignment,
+                         [(b'testtopic_replicated', [0, 1, 2, 8, 3, 9, 4, 5, 6, 7])])
+
+
 if __name__ == '__main__':
     unittest2.main()