Repository: dpkp/kafka-python
Pull request #2505 (2025-02-27 18:09:40)

Support OffsetFetch v5 / OffsetCommit v6 (2.1 baseline)


Patch:
diff --git a/kafka/client_async.py b/kafka/client_async.py
index 96959d9..9e57efd 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -399,13 +399,23 @@ class KafkaClient(object):
         return False
 
     def _maybe_connect(self, node_id):
-        """Idempotent non-blocking connection attempt to the given node id."""
+        """Idempotent non-blocking connection attempt to the given node id.
+
+        Returns True if connection object exists and is connected / connecting
+        """
         with self._lock:
             conn = self._conns.get(node_id)
 
+            # Check if existing connection should be recreated because host/port changed
+            if conn is not None and self._should_recycle_connection(conn):
+                self._conns.pop(node_id).close()
+                conn = None
+
             if conn is None:
                 broker = self.cluster.broker_metadata(node_id)
-                assert broker, 'Broker id %s not in current metadata' % (node_id,)
+                if broker is None:
+                    log.debug('Broker id %s not in current metadata', node_id)
+                    return False
 
                 log.debug("Initiating connection to node %s at %s:%s",
                           node_id, broker.host, broker.port)
@@ -417,16 +427,11 @@ class KafkaClient(object):
                                         **self.config)
                 self._conns[node_id] = conn
 
-            # Check if existing connection should be recreated because host/port changed
-            elif self._should_recycle_connection(conn):
-                self._conns.pop(node_id)
-                return False
-
             elif conn.connected():
                 return True
 
             conn.connect()
-            return conn.connected()
+            return not conn.disconnected()
 
     def ready(self, node_id, metadata_priority=True):
         """Check whether a node is connected and ok to send more requests.
@@ -621,7 +626,10 @@ class KafkaClient(object):
 
                 # Attempt to complete pending connections
                 for node_id in list(self._connecting):
-                    self._maybe_connect(node_id)
+                    # False return means no more connection progress is possible
+                    # Connected nodes will update _connecting via state_change callback
+                    if not self._maybe_connect(node_id):
+                        self._connecting.remove(node_id)
 
                 # If we got a future that is already done, don't block in _poll
                 if future is not None and future.is_done:
@@ -965,7 +973,12 @@ class KafkaClient(object):
             if try_node is None:
                 self._lock.release()
                 raise Errors.NoBrokersAvailable()
-            self._maybe_connect(try_node)
+            if not self._maybe_connect(try_node):
+                if try_node == node_id:
+                    raise Errors.NodeNotReadyError("Connection failed to %s" % node_id)
+                else:
+                    continue
+
             conn = self._conns[try_node]
 
             # We will intentionally cause socket failures
diff --git a/kafka/coordinator/base.py b/kafka/coordinator/base.py
index 8607e48..b6caabc 100644
--- a/kafka/coordinator/base.py
+++ b/kafka/coordinator/base.py
@@ -14,10 +14,8 @@ from kafka import errors as Errors
 from kafka.future import Future
 from kafka.metrics import AnonMeasurable
 from kafka.metrics.stats import Avg, Count, Max, Rate
-from kafka.protocol.commit import OffsetCommitRequest
 from kafka.protocol.find_coordinator import FindCoordinatorRequest
-from kafka.protocol.group import (HeartbeatRequest, JoinGroupRequest,
-                            LeaveGroupRequest, SyncGroupRequest)
+from kafka.protocol.group import HeartbeatRequest, JoinGroupRequest, LeaveGroupRequest, SyncGroupRequest, DEFAULT_GENERATION_ID, UNKNOWN_MEMBER_ID
 
 log = logging.getLogger('kafka.coordinator')
 
@@ -34,10 +32,7 @@ class Generation(object):
         self.member_id = member_id
         self.protocol = protocol
 
-Generation.NO_GENERATION = Generation(
-    OffsetCommitRequest[2].DEFAULT_GENERATION_ID,
-    JoinGroupRequest[0].UNKNOWN_MEMBER_ID,
-    None)
+Generation.NO_GENERATION = Generation(DEFAULT_GENERATION_ID, UNKNOWN_MEMBER_ID, None)
 
 
 class UnjoinedGroupException(Errors.KafkaError):
diff --git a/kafka/coordinator/consumer.py b/kafka/coordinator/consumer.py
index 5f62f73..5850d1a 100644
--- a/kafka/coordinator/consumer.py
+++ b/kafka/coordinator/consumer.py
@@ -575,7 +575,7 @@ class ConsumerCoordinator(BaseCoordinator):
             offset_data[tp.topic][tp.partition] = offset
 
         if self._subscription.partitions_auto_assigned():
-            generation = self.generation()
+            generation = self.generation() or Generation.NO_GENERATION
         else:
             generation = Generation.NO_GENERATION
 
@@ -585,8 +585,35 @@ class ConsumerCoordinator(BaseCoordinator):
         if self.config['api_version'] >= (0, 9) and generation is None:
             return Future().failure(Errors.CommitFailedError())
 
-        version = self._client.api_version(OffsetCommitRequest, max_version=2)
-        if version == 2:
+        version = self._client.api_version(OffsetCommitRequest, max_version=6)
+        if version == 0:
+            request = OffsetCommitRequest[version](
+                self.group_id,
+                [(
+                    topic, [(
+                        partition,
+                        offset.offset,
+                        offset.metadata
+                    ) for partition, offset in six.iteritems(partitions)]
+                ) for topic, partitions in six.iteritems(offset_data)]
+            )
+        elif version == 1:
+            request = OffsetCommitRequest[version](
+                self.group_id,
+                # This api version was only used in v0.8.2, prior to join group apis
+                # so this always ends up as NO_GENERATION
+                generation.generation_id,
+                generation.member_id,
+                [(
+                    topic, [(
+                        partition,
+                        offset.offset,
+                        -1, # timestamp, unused
+                        offset.metadata
+                    ) for partition, offset in six.iteritems(partitions)]
+                ) for topic, partitions in six.iteritems(offset_data)]
+            )
+        elif version <= 4:
             request = OffsetCommitRequest[version](
                 self.group_id,
                 generation.generation_id,
@@ -600,25 +627,29 @@ class ConsumerCoordinator(BaseCoordinator):
                     ) for partition, offset in six.iteritems(partitions)]
                 ) for topic, partitions in six.iteritems(offset_data)]
             )
-        elif version == 1:
+        elif version <= 5:
             request = OffsetCommitRequest[version](
-                self.group_id, -1, '',
+                self.group_id,
+                generation.generation_id,
+                generation.member_id,
                 [(
                     topic, [(
                         partition,
                         offset.offset,
-                        -1,
                         offset.metadata
                     ) for partition, offset in six.iteritems(partitions)]
                 ) for topic, partitions in six.iteritems(offset_data)]
             )
-        elif version == 0:
+        else:
             request = OffsetCommitRequest[version](
                 self.group_id,
+                generation.generation_id,
+                generation.member_id,
                 [(
                     topic, [(
                         partition,
                         offset.offset,
+                        -1, # leader_epoch
                         offset.metadata
                     ) for partition, offset in six.iteritems(partitions)]
                 ) for topic, partitions in six.iteritems(offset_data)]
@@ -634,6 +665,8 @@ class ConsumerCoordinator(BaseCoordinator):
         return future
 
     def _handle_offset_commit_response(self, offsets, future, send_time, response):
+        if response.API_VERSION >= 3 and response.throttle_time_ms > 0:
+            log.warning()
         # TODO look at adding request_latency_ms to response (like java kafka)
         self.consumer_sensors.commit_latency.record((time.time() - send_time) * 1000)
         unauthorized_topics = set()
@@ -735,7 +768,9 @@ class ConsumerCoordinator(BaseCoordinator):
         for tp in partitions:
             topic_partitions[tp.topic].add(tp.partition)
 
-        version = self._client.api_version(OffsetFetchRequest, max_version=1)
+        version = self._client.api_version(OffsetFetchRequest, max_version=5)
+        # Starting in version 2, the request can contain a null topics array to indicate that offsets should be fetched
+        # TODO: support
         request = OffsetFetchRequest[version](
             self.group_id,
             list(topic_partitions.items())
@@ -749,9 +784,23 @@ class ConsumerCoordinator(BaseCoordinator):
         return future
 
     def _handle_offset_fetch_response(self, future, response):
+        if response.API_VERSION >= 3 and response.throttle_time_ms > 0:
+            log.warning()
+
+        if response.API_VERSION >= 2 and response.error_code != Errors.NoError.errno:
+            error_type = Errors.for_code(response.error_code)
+            # TODO: handle...
+
         offsets = {}
         for topic, partitions in response.topics:
-            for partition, offset, metadata, error_code in partitions:
+            for partition_data in partitions:
+                partition, offset = partition_data[:2]
+                if response.API_VERSION >= 5:
+                    leader_epoch, metadata, error_code = partition_data[2:]
+                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:
diff --git a/kafka/protocol/commit.py b/kafka/protocol/commit.py
index f5828ba..53c2466 100644
--- a/kafka/protocol/commit.py
+++ b/kafka/protocol/commit.py
@@ -41,6 +41,24 @@ class OffsetCommitResponse_v3(Response):
     )
 
 
+class OffsetCommitResponse_v4(Response):
+    API_KEY = 8
+    API_VERSION = 4
+    SCHEMA = OffsetCommitResponse_v3.SCHEMA
+
+
+class OffsetCommitResponse_v5(Response):
+    API_KEY = 8
+    API_VERSION = 5
+    SCHEMA = OffsetCommitResponse_v4.SCHEMA
+
+
+class OffsetCommitResponse_v6(Response):
+    API_KEY = 8
+    API_VERSION = 6
+    SCHEMA = OffsetCommitResponse_v5.SCHEMA
+
+
 class OffsetCommitRequest_v0(Request):
     API_KEY = 8
     API_VERSION = 0  # Zookeeper-backed storage
@@ -76,13 +94,13 @@ class OffsetCommitRequest_v1(Request):
 
 class OffsetCommitRequest_v2(Request):
     API_KEY = 8
-    API_VERSION = 2  # added retention_time, dropped timestamp
+    API_VERSION = 2
     RESPONSE_TYPE = OffsetCommitResponse_v2
     SCHEMA = Schema(
         ('consumer_group', String('utf-8')),
         ('consumer_group_generation_id', Int32),
         ('consumer_id', String('utf-8')),
-        ('retention_time', Int64),
+        ('retention_time', Int64), # added retention_time, dropped timestamp
         ('topics', Array(
             ('topic', String('utf-8')),
             ('partitions', Array(
@@ -90,7 +108,6 @@ class OffsetCommitRequest_v2(Request):
                 ('offset', Int64),
                 ('metadata', String('utf-8'))))))
     )
-    DEFAULT_GENERATION_ID = -1
     DEFAULT_RETENTION_TIME = -1
 
 
@@ -99,15 +116,63 @@ class OffsetCommitRequest_v3(Request):
     API_VERSION = 3
     RESPONSE_TYPE = OffsetCommitResponse_v3
     SCHEMA = OffsetCommitRequest_v2.SCHEMA
+    DEFAULT_RETENTION_TIME = -1
+
+
+class OffsetCommitRequest_v4(Request):
+    API_KEY = 8
+    API_VERSION = 4
+    RESPONSE_TYPE = OffsetCommitResponse_v4
+    SCHEMA = OffsetCommitRequest_v3.SCHEMA
+    DEFAULT_RETENTION_TIME = -1
+
+
+class OffsetCommitRequest_v5(Request):
+    API_KEY = 8
+    API_VERSION = 5 # drops retention_time
+    RESPONSE_TYPE = OffsetCommitResponse_v5
+    SCHEMA = Schema(
+        ('consumer_group', String('utf-8')),
+        ('consumer_group_generation_id', Int32),
+        ('consumer_id', String('utf-8')),
+        ('topics', Array(
+            ('topic', String('utf-8')),
+            ('partitions', Array(
+                ('partition', Int32),
+                ('offset', Int64),
+                ('metadata', String('utf-8'))))))
+    )
+
+
+class OffsetCommitRequest_v6(Request):
+    API_KEY = 8
+    API_VERSION = 6
+    RESPONSE_TYPE = OffsetCommitResponse_v6
+    SCHEMA = Schema(
+        ('consumer_group', String('utf-8')),
+        ('consumer_group_generation_id', Int32),
+        ('consumer_id', String('utf-8')),
+        ('topics', Array(
+            ('topic', String('utf-8')),
+            ('partitions', Array(
+                ('partition', Int32),
+                ('offset', Int64),
+                ('leader_epoch', Int32), # added for fencing / kip-320. default -1
+                ('metadata', String('utf-8'))))))
+    )
 
 
 OffsetCommitRequest = [
     OffsetCommitRequest_v0, OffsetCommitRequest_v1,
-    OffsetCommitRequest_v2, OffsetCommitRequest_v3
+    OffsetCommitRequest_v2, OffsetCommitRequest_v3,
+    OffsetCommitRequest_v4, OffsetCommitRequest_v5,
+    OffsetCommitRequest_v6,
 ]
 OffsetCommitResponse = [
     OffsetCommitResponse_v0, OffsetCommitResponse_v1,
-    OffsetCommitResponse_v2, OffsetCommitResponse_v3
+    OffsetCommitResponse_v2, OffsetCommitResponse_v3,
+    OffsetCommitResponse_v4, OffsetCommitResponse_v5,
+    OffsetCommitResponse_v6,
 ]
 
 
@@ -163,6 +228,29 @@ class OffsetFetchResponse_v3(Response):
     )
 
 
+class OffsetFetchResponse_v4(Response):
+    API_KEY = 9
+    API_VERSION = 4
+    SCHEMA = OffsetFetchResponse_v3.SCHEMA
+
+
+class OffsetFetchResponse_v5(Response):
+    API_KEY = 9
+    API_VERSION = 5
+    SCHEMA = Schema(
+        ('throttle_time_ms', Int32),
+        ('topics', Array(
+            ('topic', String('utf-8')),
+            ('partitions', Array(
+                ('partition', Int32),
+                ('offset', Int64),
+                ('leader_epoch', Int32),
+                ('metadata', String('utf-8')),
+                ('error_code', Int16))))),
+        ('error_code', Int16)
+    )
+
+
 class OffsetFetchRequest_v0(Request):
     API_KEY = 9
     API_VERSION = 0  # zookeeper-backed storage
@@ -199,11 +287,27 @@ class OffsetFetchRequest_v3(Request):
     SCHEMA = OffsetFetchRequest_v2.SCHEMA
 
 
+class OffsetFetchRequest_v4(Request):
+    API_KEY = 9
+    API_VERSION = 4
+    RESPONSE_TYPE = OffsetFetchResponse_v4
+    SCHEMA = OffsetFetchRequest_v3.SCHEMA
+
+
+class OffsetFetchRequest_v5(Request):
+    API_KEY = 9
+    API_VERSION = 5
+    RESPONSE_TYPE = OffsetFetchResponse_v5
+    SCHEMA = OffsetFetchRequest_v4.SCHEMA
+
+
 OffsetFetchRequest = [
     OffsetFetchRequest_v0, OffsetFetchRequest_v1,
     OffsetFetchRequest_v2, OffsetFetchRequest_v3,
+    OffsetFetchRequest_v4, OffsetFetchRequest_v5,
 ]
 OffsetFetchResponse = [
     OffsetFetchResponse_v0, OffsetFetchResponse_v1,
     OffsetFetchResponse_v2, OffsetFetchResponse_v3,
+    OffsetFetchResponse_v4, OffsetFetchResponse_v5,
 ]
diff --git a/kafka/protocol/group.py b/kafka/protocol/group.py
index ee06141..3b32590 100644
--- a/kafka/protocol/group.py
+++ b/kafka/protocol/group.py
@@ -5,6 +5,10 @@ from kafka.protocol.struct import Struct
 from kafka.protocol.types import Array, Bytes, Int16, Int32, Schema, String
 
 
+DEFAULT_GENERATION_ID = -1
+UNKNOWN_MEMBER_ID = ''
+
+
 class JoinGroupResponse_v0(Response):
     API_KEY = 11
     API_VERSION = 0
@@ -61,7 +65,6 @@ class JoinGroupRequest_v0(Request):
             ('protocol_name', String('utf-8')),
             ('protocol_metadata', Bytes)))
     )
-    UNKNOWN_MEMBER_ID = ''
 
 
 class JoinGroupRequest_v1(Request):
@@ -78,7 +81,6 @@ class JoinGroupRequest_v1(Request):
             ('protocol_name', String('utf-8')),
             ('protocol_metadata', Bytes)))
     )
-    UNKNOWN_MEMBER_ID = ''
 
 
 class JoinGroupRequest_v2(Request):
@@ -86,7 +88,6 @@ class JoinGroupRequest_v2(Request):
     API_VERSION = 2
     RESPONSE_TYPE = JoinGroupResponse_v2
     SCHEMA = JoinGroupRequest_v1.SCHEMA
-    UNKNOWN_MEMBER_ID = ''
 
 
 class JoinGroupRequest_v3(Request):


Test patch:
diff --git a/test/test_client_async.py b/test/test_client_async.py
index ccdd570..16ee429 100644
--- a/test/test_client_async.py
+++ b/test/test_client_async.py
@@ -71,19 +71,14 @@ def test_can_connect(cli, conn):
 
 
 def test_maybe_connect(cli, conn):
-    try:
-        # Node not in metadata, raises AssertionError
-        cli._maybe_connect(2)
-    except AssertionError:
-        pass
-    else:
-        assert False, 'Exception not raised'
+    # Node not in metadata, return False
+    assert not cli._maybe_connect(2)
 
     # New node_id creates a conn object
     assert 0 not in cli._conns
     conn.state = ConnectionStates.DISCONNECTED
     conn.connect.side_effect = lambda: conn._set_conn_state(ConnectionStates.CONNECTING)
-    assert cli._maybe_connect(0) is False
+    assert cli._maybe_connect(0) is True
     assert cli._conns[0] is conn
 
 
diff --git a/test/test_coordinator.py b/test/test_coordinator.py
index 0c4ee6d..c0e7c6d 100644
--- a/test/test_coordinator.py
+++ b/test/test_coordinator.py
@@ -439,7 +439,11 @@ def test_send_offset_commit_request_fail(mocker, patched_coord, offsets):
 @pytest.mark.parametrize('api_version,req_type', [
     ((0, 8, 1), OffsetCommitRequest[0]),
     ((0, 8, 2), OffsetCommitRequest[1]),
-    ((0, 9), OffsetCommitRequest[2])])
+    ((0, 9), OffsetCommitRequest[2]),
+    ((0, 11), OffsetCommitRequest[3]),
+    ((2, 0), OffsetCommitRequest[4]),
+    ((2, 1), OffsetCommitRequest[6]),
+])
 def test_send_offset_commit_request_versions(patched_coord, offsets,
                                              api_version, req_type):
     expect_node = 0
@@ -499,13 +503,27 @@ def test_send_offset_commit_request_success(mocker, patched_coord, offsets):
      Errors.InvalidTopicError, False),
     (OffsetCommitResponse[0]([('foobar', [(0, 29), (1, 29)])]),
      Errors.TopicAuthorizationFailedError, False),
+    (OffsetCommitResponse[0]([('foobar', [(0, 0), (1, 0)])]),
+     None, False),
+    (OffsetCommitResponse[1]([('foobar', [(0, 0), (1, 0)])]),
+     None, False),
+    (OffsetCommitResponse[2]([('foobar', [(0, 0), (1, 0)])]),
+     None, False),
+    (OffsetCommitResponse[3](0, [('foobar', [(0, 0), (1, 0)])]),
+     None, False),
+    (OffsetCommitResponse[4](0, [('foobar', [(0, 0), (1, 0)])]),
+     None, False),
+    (OffsetCommitResponse[5](0, [('foobar', [(0, 0), (1, 0)])]),
+     None, False),
+    (OffsetCommitResponse[6](0, [('foobar', [(0, 0), (1, 0)])]),
+     None, False),
 ])
 def test_handle_offset_commit_response(mocker, patched_coord, offsets,
                                        response, error, dead):
     future = Future()
     patched_coord._handle_offset_commit_response(offsets, future, time.time(),
                                                  response)
-    assert isinstance(future.exception, error)
+    assert isinstance(future.exception, error) if error else True
     assert patched_coord.coordinator_id is (None if dead else 0)
 
 
@@ -534,7 +552,12 @@ def test_send_offset_fetch_request_fail(mocker, patched_coord, partitions):
 @pytest.mark.parametrize('api_version,req_type', [
     ((0, 8, 1), OffsetFetchRequest[0]),
     ((0, 8, 2), OffsetFetchRequest[1]),
-    ((0, 9), OffsetFetchRequest[1])])
+    ((0, 9), OffsetFetchRequest[1]),
+    ((0, 10, 2), OffsetFetchRequest[2]),
+    ((0, 11), OffsetFetchRequest[3]),
+    ((2, 0), OffsetFetchRequest[4]),
+    ((2, 1), OffsetFetchRequest[5]),
+])
 def test_send_offset_fetch_request_versions(patched_coord, partitions,
                                             api_version, req_type):
     # assuming fixture sets coordinator=0, least_loaded_node=1
@@ -583,6 +606,16 @@ def test_send_offset_fetch_request_success(patched_coord, partitions):
      Errors.TopicAuthorizationFailedError, False),
     (OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 0), (1, 234, b'', 0)])]),
      None, False),
+    (OffsetFetchResponse[1]([('foobar', [(0, 123, b'', 0), (1, 234, b'', 0)])]),
+     None, False),
+    (OffsetFetchResponse[2]([('foobar', [(0, 123, b'', 0), (1, 234, b'', 0)])], 0),
+     None, False),
+    (OffsetFetchResponse[3](0, [('foobar', [(0, 123, b'', 0), (1, 234, b'', 0)])], 0),
+     None, False),
+    (OffsetFetchResponse[4](0, [('foobar', [(0, 123, b'', 0), (1, 234, b'', 0)])], 0),
+     None, False),
+    (OffsetFetchResponse[5](0, [('foobar', [(0, 123, -1, b'', 0), (1, 234, -1, b'', 0)])], 0),
+     None, False),
 ])
 def test_handle_offset_fetch_response(patched_coord, offsets,
                                       response, error, dead):