Use GroupMember namedtuple for join group response members

Author: dpkp

kafka/coordinator/base.py

@@ -16,7 +16,10 @@
 from kafka.metrics import AnonMeasurable
 from kafka.metrics.stats import Avg, Count, Max, Rate
 from kafka.protocol.find_coordinator import FindCoordinatorRequest
-from kafka.protocol.group import HeartbeatRequest, JoinGroupRequest, LeaveGroupRequest, SyncGroupRequest, DEFAULT_GENERATION_ID, UNKNOWN_MEMBER_ID
+from kafka.protocol.group import (
+    HeartbeatRequest, JoinGroupRequest, LeaveGroupRequest, SyncGroupRequest,
+    DEFAULT_GENERATION_ID, UNKNOWN_MEMBER_ID, GroupMember,
+)
 from kafka.util import Timer
 
 log = logging.getLogger('kafka.coordinator')
@@ -216,10 +219,10 @@ def _perform_assignment(self, leader_id, protocol, members):
         Arguments:
             leader_id (str): The id of the leader (which is this member)
             protocol (str): the chosen group protocol (assignment strategy)
-            members (list): [(member_id, group_instance_id, metadata_bytes)] from
-                JoinGroupResponse. metadata_bytes are associated with the chosen
-                group protocol, and the Coordinator subclass is responsible for
-                decoding metadata_bytes based on that protocol.
+            members (list): [GroupMember] from JoinGroupResponse.
+                metadata_bytes are associated with the chosen group protocol,
+                and the Coordinator subclass is responsible for decoding
+                metadata_bytes based on that protocol.
 
         Returns:
             dict: {member_id: assignment}; assignment must either be bytes
@@ -702,9 +705,11 @@ def _on_join_leader(self, response):
             Future: resolves to member assignment encoded-bytes
         """
         try:
+            members = [GroupMember(*member) if response.API_VERSION >= 5 else GroupMember(member[0], None, member[1])
+                       for member in response.members]
             group_assignment = self._perform_assignment(response.leader_id,
                                                         response.group_protocol,
-                                                        response.members)
+                                                        members)
             for member_id, assignment in six.iteritems(group_assignment):
                 if not isinstance(assignment, bytes):
                     group_assignment[member_id] = assignment.encode()

kafka/coordinator/consumer.py

@@ -337,16 +337,11 @@ def _perform_assignment(self, leader_id, assignment_strategy, members):
         member_subscriptions = {}
         all_subscribed_topics = set()
         for member in members:
-            if len(member) == 3:
-                member_id, group_instance_id, metadata_bytes = member
-            else:
-                member_id, metadata_bytes = member
-                group_instance_id = None
             subscription = Subscription(
-                ConsumerProtocol[0].METADATA.decode(metadata_bytes),
-                group_instance_id
+                ConsumerProtocol[0].METADATA.decode(member.metadata_bytes),
+                member.group_instance_id
             )
-            member_subscriptions[member_id] = subscription
+            member_subscriptions[member.member_id] = subscription
             all_subscribed_topics.update(subscription.subscription)
 
         # the leader will begin watching for changes to any of the topics

kafka/protocol/group.py

@@ -1,5 +1,7 @@
 from __future__ import absolute_import
 
+import collections
+
 from kafka.protocol.api import Request, Response
 from kafka.protocol.struct import Struct
 from kafka.protocol.types import Array, Bytes, Int16, Int32, Schema, String
@@ -8,6 +10,9 @@
 DEFAULT_GENERATION_ID = -1
 UNKNOWN_MEMBER_ID = ''
 
+GroupMember = collections.namedtuple("GroupMember", ["member_id", "group_instance_id", "metadata_bytes"])
+GroupMember.__new__.__defaults__ = (None,) * len(GroupMember._fields)
+
 
 class JoinGroupResponse_v0(Response):
     API_KEY = 11

test/test_coordinator.py

@@ -21,6 +21,7 @@
 from kafka.protocol.commit import (
     OffsetCommitRequest, OffsetCommitResponse,
     OffsetFetchRequest, OffsetFetchResponse)
+from kafka.protocol.group import GroupMember
 from kafka.protocol.metadata import MetadataResponse
 from kafka.structs import OffsetAndMetadata, TopicPartition
 from kafka.util import WeakMethod
@@ -209,7 +210,7 @@ def test_perform_assignment(mocker, coordinator):
 
     ret = coordinator._perform_assignment(
         'member-foo', 'roundrobin',
-        [(member, subscription.encode())
+        [GroupMember(member, None, subscription.encode())
          for member, subscription in group_subscriptions.items()])
 
     assert RoundRobinPartitionAssignor.assign.call_count == 1