Repository: dpkp/kafka-python
Pull request #2151 (2020-10-27 19:20:00)

KIP-482: add support for flexible versions / tagged fields
- Add support for varints
- Add support for compact collections (byte array, string, array)
- Add support for new request and response headers, supporting flexible
versions
- Add list of supported flexible versions, plus extraction script
- Add List / Alter partition reassignments apis 

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/dpkp/kafka-python/2151)
<!-- Reviewable:end -->


gabriel-tincu: The way the request / response headers are handled needs an update, as apparently it's possible to send requests that have an api version with flexible versions enabled with both flexible versions and without. So if someone ends up defining a new requests object but using the old collections, introducing the empty tag buffer will mess up the entire payload. The java code base uses the json files to generate the classes so maybe we should look into that as well?

jeffwidman: Is there a KIP associated with this change?
gabriel-tincu: > Is there a KIP associated with this change? I switched teams and only work infrequently with Kafka now, so haven't kept up with the latest news.

@jeffwidman 
New types as described by the protocol docs: https://kafka.apache.org/protocol.html#protocol_types
Tagged fields KIP: https://cwiki.apache.org/confluence/display/KAFKA/KIP-482%3A+The+Kafka+Protocol+should+Support+Optional+Tagged+Fields
Jira issue requiring version bump to use both tagged fields and compact data types for apis: https://issues.apache.org/jira/browse/KAFKA-9420
gabriel-tincu: @dpkp this is ready for another look, whenever you have the time

Patch:
diff --git a/kafka/admin/client.py b/kafka/admin/client.py
index c58da0c..63a0f3b 100644
--- a/kafka/admin/client.py
+++ b/kafka/admin/client.py
@@ -146,6 +146,7 @@ class KafkaAdminClient(object):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        kafka_client (callable): Custom class / callable for creating KafkaClient instances
 
     """
     DEFAULT_CONFIG = {
@@ -186,6 +187,7 @@ class KafkaAdminClient(object):
         'metric_reporters': [],
         'metrics_num_samples': 2,
         'metrics_sample_window_ms': 30000,
+        'kafka_client': KafkaClient,
     }
 
     def __init__(self, **configs):
@@ -205,9 +207,11 @@ class KafkaAdminClient(object):
         reporters = [reporter() for reporter in self.config['metric_reporters']]
         self._metrics = Metrics(metric_config, reporters)
 
-        self._client = KafkaClient(metrics=self._metrics,
-                                   metric_group_prefix='admin',
-                                   **self.config)
+        self._client = self.config['kafka_client'](
+            metrics=self._metrics,
+            metric_group_prefix='admin',
+            **self.config
+        )
         self._client.check_version(timeout=(self.config['api_version_auto_timeout_ms'] / 1000))
 
         # Get auto-discovered version from client if necessary
diff --git a/kafka/consumer/group.py b/kafka/consumer/group.py
index 4fd57ae..a1d1dfa 100644
--- a/kafka/consumer/group.py
+++ b/kafka/consumer/group.py
@@ -244,6 +244,7 @@ class KafkaConsumer(six.Iterator):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        kafka_client (callable): Custom class / callable for creating KafkaClient instances
 
     Note:
         Configuration parameters are described in more detail at
@@ -306,6 +307,7 @@ class KafkaConsumer(six.Iterator):
         'sasl_kerberos_domain_name': None,
         'sasl_oauth_token_provider': None,
         'legacy_iterator': False, # enable to revert to < 1.4.7 iterator
+        'kafka_client': KafkaClient,
     }
     DEFAULT_SESSION_TIMEOUT_MS_0_9 = 30000
 
@@ -353,7 +355,7 @@ class KafkaConsumer(six.Iterator):
             log.warning('use api_version=%s [tuple] -- "%s" as str is deprecated',
                         str(self.config['api_version']), str_version)
 
-        self._client = KafkaClient(metrics=self._metrics, **self.config)
+        self._client = self.config['kafka_client'](metrics=self._metrics, **self.config)
 
         # Get auto-discovered version from client if necessary
         if self.config['api_version'] is None:
diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py
index cde26b0..ea010c5 100644
--- a/kafka/producer/kafka.py
+++ b/kafka/producer/kafka.py
@@ -280,6 +280,7 @@ class KafkaProducer(object):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        kafka_client (callable): Custom class / callable for creating KafkaClient instances
 
     Note:
         Configuration parameters are described in more detail at
@@ -332,7 +333,8 @@ class KafkaProducer(object):
         'sasl_plain_password': None,
         'sasl_kerberos_service_name': 'kafka',
         'sasl_kerberos_domain_name': None,
-        'sasl_oauth_token_provider': None
+        'sasl_oauth_token_provider': None,
+        'kafka_client': KafkaClient,
     }
 
     _COMPRESSORS = {
@@ -378,9 +380,10 @@ class KafkaProducer(object):
         reporters = [reporter() for reporter in self.config['metric_reporters']]
         self._metrics = Metrics(metric_config, reporters)
 
-        client = KafkaClient(metrics=self._metrics, metric_group_prefix='producer',
-                             wakeup_timeout_ms=self.config['max_block_ms'],
-                             **self.config)
+        client = self.config['kafka_client'](
+            metrics=self._metrics, metric_group_prefix='producer',
+            wakeup_timeout_ms=self.config['max_block_ms'],
+            **self.config)
 
         # Get auto-discovered version from client if necessary
         if self.config['api_version'] is None:
diff --git a/kafka/protocol/__init__.py b/kafka/protocol/__init__.py
index e739b5c..025447f 100644
--- a/kafka/protocol/__init__.py
+++ b/kafka/protocol/__init__.py
@@ -43,5 +43,7 @@ API_KEYS = {
     40: 'ExpireDelegationToken',
     41: 'DescribeDelegationToken',
     42: 'DeleteGroups',
+    45: 'AlterPartitionReassignments',
+    46: 'ListPartitionReassignments',
     48: 'DescribeClientQuotas',
 }
diff --git a/kafka/protocol/admin.py b/kafka/protocol/admin.py
index 63a3327..f9d61e5 100644
--- a/kafka/protocol/admin.py
+++ b/kafka/protocol/admin.py
@@ -1,7 +1,7 @@
 from __future__ import absolute_import
 
 from kafka.protocol.api import Request, Response
-from kafka.protocol.types import Array, Boolean, Bytes, Int8, Int16, Int32, Int64, Schema, String, Float64
+from kafka.protocol.types import Array, Boolean, Bytes, Int8, Int16, Int32, Int64, Schema, String, Float64, CompactString, CompactArray, TaggedFields
 
 
 class ApiVersionResponse_v0(Response):
@@ -963,3 +963,92 @@ DescribeClientQuotasRequest = [
 DescribeClientQuotasResponse = [
     DescribeClientQuotasResponse_v0,
 ]
+
+
+class AlterPartitionReassignmentsResponse_v0(Response):
+    API_KEY = 45
+    API_VERSION = 0
+    SCHEMA = Schema(
+        ("throttle_time_ms", Int32),
+        ("error_code", Int16),
+        ("error_message", CompactString("utf-8")),
+        ("responses", CompactArray(
+            ("name", CompactString("utf-8")),
+            ("partitions", CompactArray(
+                ("partition_index", Int32),
+                ("error_code", Int16),
+                ("error_message", CompactString("utf-8")),
+                ("tags", TaggedFields)
+            )),
+            ("tags", TaggedFields)
+        )),
+        ("tags", TaggedFields)
+    )
+
+
+class AlterPartitionReassignmentsRequest_v0(Request):
+    FLEXIBLE_VERSION = True
+    API_KEY = 45
+    API_VERSION = 0
+    RESPONSE_TYPE = AlterPartitionReassignmentsResponse_v0
+    SCHEMA = Schema(
+        ("timeout_ms", Int32),
+        ("topics", CompactArray(
+            ("name", CompactString("utf-8")),
+            ("partitions", CompactArray(
+                ("partition_index", Int32),
+                ("replicas", CompactArray(Int32)),
+                ("tags", TaggedFields)
+            )),
+            ("tags", TaggedFields)
+        )),
+        ("tags", TaggedFields)
+    )
+
+
+AlterPartitionReassignmentsRequest = [AlterPartitionReassignmentsRequest_v0]
+
+AlterPartitionReassignmentsResponse = [AlterPartitionReassignmentsResponse_v0]
+
+
+class ListPartitionReassignmentsResponse_v0(Response):
+    API_KEY = 46
+    API_VERSION = 0
+    SCHEMA = Schema(
+        ("throttle_time_ms", Int32),
+        ("error_code", Int16),
+        ("error_message", CompactString("utf-8")),
+        ("topics", CompactArray(
+            ("name", CompactString("utf-8")),
+            ("partitions", CompactArray(
+                ("partition_index", Int32),
+                ("replicas", CompactArray(Int32)),
+                ("adding_replicas", CompactArray(Int32)),
+                ("removing_replicas", CompactArray(Int32)),
+                ("tags", TaggedFields)
+            )),
+            ("tags", TaggedFields)
+        )),
+        ("tags", TaggedFields)
+    )
+
+
+class ListPartitionReassignmentsRequest_v0(Request):
+    FLEXIBLE_VERSION = True
+    API_KEY = 46
+    API_VERSION = 0
+    RESPONSE_TYPE = ListPartitionReassignmentsResponse_v0
+    SCHEMA = Schema(
+        ("timeout_ms", Int32),
+        ("topics", CompactArray(
+            ("name", CompactString("utf-8")),
+            ("partition_index", CompactArray(Int32)),
+            ("tags", TaggedFields)
+        )),
+        ("tags", TaggedFields)
+    )
+
+
+ListPartitionReassignmentsRequest = [ListPartitionReassignmentsRequest_v0]
+
+ListPartitionReassignmentsResponse = [ListPartitionReassignmentsResponse_v0]
diff --git a/kafka/protocol/api.py b/kafka/protocol/api.py
index 64276fc..f12cb97 100644
--- a/kafka/protocol/api.py
+++ b/kafka/protocol/api.py
@@ -3,7 +3,7 @@ from __future__ import absolute_import
 import abc
 
 from kafka.protocol.struct import Struct
-from kafka.protocol.types import Int16, Int32, String, Schema, Array
+from kafka.protocol.types import Int16, Int32, String, Schema, Array, TaggedFields
 
 
 class RequestHeader(Struct):
@@ -20,9 +20,40 @@ class RequestHeader(Struct):
         )
 
 
+class RequestHeaderV2(Struct):
+    # Flexible response / request headers end in field buffer
+    SCHEMA = Schema(
+        ('api_key', Int16),
+        ('api_version', Int16),
+        ('correlation_id', Int32),
+        ('client_id', String('utf-8')),
+        ('tags', TaggedFields),
+    )
+
+    def __init__(self, request, correlation_id=0, client_id='kafka-python', tags=None):
+        super(RequestHeaderV2, self).__init__(
+            request.API_KEY, request.API_VERSION, correlation_id, client_id, tags or {}
+        )
+
+
+class ResponseHeader(Struct):
+    SCHEMA = Schema(
+        ('correlation_id', Int32),
+    )
+
+
+class ResponseHeaderV2(Struct):
+    SCHEMA = Schema(
+        ('correlation_id', Int32),
+        ('tags', TaggedFields),
+    )
+
+
 class Request(Struct):
     __metaclass__ = abc.ABCMeta
 
+    FLEXIBLE_VERSION = False
+
     @abc.abstractproperty
     def API_KEY(self):
         """Integer identifier for api request"""
@@ -50,6 +81,16 @@ class Request(Struct):
     def to_object(self):
         return _to_object(self.SCHEMA, self)
 
+    def build_request_header(self, correlation_id, client_id):
+        if self.FLEXIBLE_VERSION:
+            return RequestHeaderV2(self, correlation_id=correlation_id, client_id=client_id)
+        return RequestHeader(self, correlation_id=correlation_id, client_id=client_id)
+
+    def parse_response_header(self, read_buffer):
+        if self.FLEXIBLE_VERSION:
+            return ResponseHeaderV2.decode(read_buffer)
+        return ResponseHeader.decode(read_buffer)
+
 
 class Response(Struct):
     __metaclass__ = abc.ABCMeta
diff --git a/kafka/protocol/parser.py b/kafka/protocol/parser.py
index cfee046..a9e7672 100644
--- a/kafka/protocol/parser.py
+++ b/kafka/protocol/parser.py
@@ -4,10 +4,9 @@ import collections
 import logging
 
 import kafka.errors as Errors
-from kafka.protocol.api import RequestHeader
 from kafka.protocol.commit import GroupCoordinatorResponse
 from kafka.protocol.frame import KafkaBytes
-from kafka.protocol.types import Int32
+from kafka.protocol.types import Int32, TaggedFields
 from kafka.version import __version__
 
 log = logging.getLogger(__name__)
@@ -59,9 +58,8 @@ class KafkaProtocol(object):
         log.debug('Sending request %s', request)
         if correlation_id is None:
             correlation_id = self._next_correlation_id()
-        header = RequestHeader(request,
-                               correlation_id=correlation_id,
-                               client_id=self._client_id)
+
+        header = request.build_request_header(correlation_id=correlation_id, client_id=self._client_id)
         message = b''.join([header.encode(), request.encode()])
         size = Int32.encode(len(message))
         data = size + message
@@ -135,17 +133,12 @@ class KafkaProtocol(object):
         return responses
 
     def _process_response(self, read_buffer):
-        recv_correlation_id = Int32.decode(read_buffer)
-        log.debug('Received correlation id: %d', recv_correlation_id)
-
         if not self.in_flight_requests:
-            raise Errors.CorrelationIdError(
-                'No in-flight-request found for server response'
-                ' with correlation ID %d'
-                % (recv_correlation_id,))
-
+            raise Errors.CorrelationIdError('No in-flight-request found for server response')
         (correlation_id, request) = self.in_flight_requests.popleft()
-
+        response_header = request.parse_response_header(read_buffer)
+        recv_correlation_id = response_header.correlation_id
+        log.debug('Received correlation id: %d', recv_correlation_id)
         # 0.8.2 quirk
         if (recv_correlation_id == 0 and
             correlation_id != 0 and
diff --git a/kafka/protocol/types.py b/kafka/protocol/types.py
index 2fde24f..0e3685d 100644
--- a/kafka/protocol/types.py
+++ b/kafka/protocol/types.py
@@ -210,3 +210,156 @@ class Array(AbstractType):
         if list_of_items is None:
             return 'NULL'
         return '[' + ', '.join([self.array_of.repr(item) for item in list_of_items]) + ']'
+
+
+class UnsignedVarInt32(AbstractType):
+    @classmethod
+    def decode(cls, data):
+        value, i = 0, 0
+        while True:
+            b, = struct.unpack('B', data.read(1))
+            if not (b & 0x80):
+                break
+            value |= (b & 0x7f) << i
+            i += 7
+            if i > 28:
+                raise ValueError('Invalid value {}'.format(value))
+        value |= b << i
+        return value
+
+    @classmethod
+    def encode(cls, value):
+        value &= 0xffffffff
+        ret = b''
+        while (value & 0xffffff80) != 0:
+            b = (value & 0x7f) | 0x80
+            ret += struct.pack('B', b)
+            value >>= 7
+        ret += struct.pack('B', value)
+        return ret
+
+
+class VarInt32(AbstractType):
+    @classmethod
+    def decode(cls, data):
+        value = UnsignedVarInt32.decode(data)
+        return (value >> 1) ^ -(value & 1)
+
+    @classmethod
+    def encode(cls, value):
+        # bring it in line with the java binary repr
+        value &= 0xffffffff
+        return UnsignedVarInt32.encode((value << 1) ^ (value >> 31))
+
+
+class VarInt64(AbstractType):
+    @classmethod
+    def decode(cls, data):
+        value, i = 0, 0
+        while True:
+            b = data.read(1)
+            if not (b & 0x80):
+                break
+            value |= (b & 0x7f) << i
+            i += 7
+            if i > 63:
+                raise ValueError('Invalid value {}'.format(value))
+        value |= b << i
+        return (value >> 1) ^ -(value & 1)
+
+    @classmethod
+    def encode(cls, value):
+        # bring it in line with the java binary repr
+        value &= 0xffffffffffffffff
+        v = (value << 1) ^ (value >> 63)
+        ret = b''
+        while (v & 0xffffffffffffff80) != 0:
+            b = (value & 0x7f) | 0x80
+            ret += struct.pack('B', b)
+            v >>= 7
+        ret += struct.pack('B', v)
+        return ret
+
+
+class CompactString(String):
+    def decode(self, data):
+        length = UnsignedVarInt32.decode(data) - 1
+        if length < 0:
+            return None
+        value = data.read(length)
+        if len(value) != length:
+            raise ValueError('Buffer underrun decoding string')
+        return value.decode(self.encoding)
+
+    def encode(self, value):
+        if value is None:
+            return UnsignedVarInt32.encode(0)
+        value = str(value).encode(self.encoding)
+        return UnsignedVarInt32.encode(len(value) + 1) + value
+
+
+class TaggedFields(AbstractType):
+    @classmethod
+    def decode(cls, data):
+        num_fields = UnsignedVarInt32.decode(data)
+        ret = {}
+        if not num_fields:
+            return ret
+        prev_tag = -1
+        for i in range(num_fields):
+            tag = UnsignedVarInt32.decode(data)
+            if tag <= prev_tag:
+                raise ValueError('Invalid or out-of-order tag {}'.format(tag))
+            prev_tag = tag
+            size = UnsignedVarInt32.decode(data)
+            val = data.read(size)
+            ret[tag] = val
+        return ret
+
+    @classmethod
+    def encode(cls, value):
+        ret = UnsignedVarInt32.encode(len(value))
+        for k, v in value.items():
+            # do we allow for other data types ?? It could get complicated really fast
+            assert isinstance(v, bytes), 'Value {} is not a byte array'.format(v)
+            assert isinstance(k, int) and k > 0, 'Key {} is not a positive integer'.format(k)
+            ret += UnsignedVarInt32.encode(k)
+            ret += v
+        return ret
+
+
+class CompactBytes(AbstractType):
+    @classmethod
+    def decode(cls, data):
+        length = UnsignedVarInt32.decode(data) - 1
+        if length < 0:
+            return None
+        value = data.read(length)
+        if len(value) != length:
+            raise ValueError('Buffer underrun decoding Bytes')
+        return value
+
+    @classmethod
+    def encode(cls, value):
+        if value is None:
+            return UnsignedVarInt32.encode(0)
+        else:
+            return UnsignedVarInt32.encode(len(value) + 1) + value
+
+
+class CompactArray(Array):
+
+    def encode(self, items):
+        if items is None:
+            return UnsignedVarInt32.encode(0)
+        return b''.join(
+            [UnsignedVarInt32.encode(len(items) + 1)] +
+            [self.array_of.encode(item) for item in items]
+        )
+
+    def decode(self, data):
+        length = UnsignedVarInt32.decode(data) - 1
+        if length == -1:
+            return None
+        return [self.array_of.decode(data) for _ in range(length)]
+


Test patch:
diff --git a/test/test_protocol.py b/test/test_protocol.py
index e295174..6a77e19 100644
--- a/test/test_protocol.py
+++ b/test/test_protocol.py
@@ -9,7 +9,7 @@ from kafka.protocol.commit import GroupCoordinatorRequest
 from kafka.protocol.fetch import FetchRequest, FetchResponse
 from kafka.protocol.message import Message, MessageSet, PartialMessage
 from kafka.protocol.metadata import MetadataRequest
-from kafka.protocol.types import Int16, Int32, Int64, String
+from kafka.protocol.types import Int16, Int32, Int64, String, UnsignedVarInt32, CompactString, CompactArray, CompactBytes
 
 
 def test_create_message():
@@ -282,3 +282,55 @@ def test_struct_unrecognized_kwargs():
 def test_struct_missing_kwargs():
     fr = FetchRequest[0](max_wait_time=100)
     assert fr.min_bytes is None
+
+
+def test_unsigned_varint_serde():
+    pairs = {
+        0: [0],
+        -1: [0xff, 0xff, 0xff, 0xff, 0x0f],
+        1: [1],
+        63: [0x3f],
+        -64: [0xc0, 0xff, 0xff, 0xff, 0x0f],
+        64: [0x40],
+        8191: [0xff, 0x3f],
+        -8192: [0x80, 0xc0, 0xff, 0xff, 0x0f],
+        8192: [0x80, 0x40],
+        -8193: [0xff, 0xbf, 0xff, 0xff, 0x0f],
+        1048575: [0xff, 0xff, 0x3f],
+
+    }
+    for value, expected_encoded in pairs.items():
+        value &= 0xffffffff
+        encoded = UnsignedVarInt32.encode(value)
+        assert encoded == b''.join(struct.pack('>B', x) for x in expected_encoded)
+        assert value == UnsignedVarInt32.decode(io.BytesIO(encoded))
+
+
+def test_compact_data_structs():
+    cs = CompactString()
+    encoded = cs.encode(None)
+    assert encoded == struct.pack('B', 0)
+    decoded = cs.decode(io.BytesIO(encoded))
+    assert decoded is None
+    assert b'\x01' == cs.encode('')
+    assert '' == cs.decode(io.BytesIO(b'\x01'))
+    encoded = cs.encode("foobarbaz")
+    assert cs.decode(io.BytesIO(encoded)) == "foobarbaz"
+
+    arr = CompactArray(CompactString())
+    assert arr.encode(None) == b'\x00'
+    assert arr.decode(io.BytesIO(b'\x00')) is None
+    enc = arr.encode([])
+    assert enc == b'\x01'
+    assert [] == arr.decode(io.BytesIO(enc))
+    encoded = arr.encode(["foo", "bar", "baz", "quux"])
+    assert arr.decode(io.BytesIO(encoded)) == ["foo", "bar", "baz", "quux"]
+
+    enc = CompactBytes.encode(None)
+    assert enc == b'\x00'
+    assert CompactBytes.decode(io.BytesIO(b'\x00')) is None
+    enc = CompactBytes.encode(b'')
+    assert enc == b'\x01'
+    assert CompactBytes.decode(io.BytesIO(b'\x01')) is b''
+    enc = CompactBytes.encode(b'foo')
+    assert CompactBytes.decode(io.BytesIO(enc)) == b'foo'