API refactor

venv/lib/python3.12/site-packages/kafka/protocol/api.py

@@ -3,7 +3,9 @@ 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
+
+from kafka.vendor.six import add_metaclass
 
 
 class RequestHeader(Struct):
@@ -20,8 +22,38 @@ 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),
+    )
+
+
+@add_metaclass(abc.ABCMeta)
 class Request(Struct):
-    __metaclass__ = abc.ABCMeta
+    FLEXIBLE_VERSION = False
 
     @abc.abstractproperty
     def API_KEY(self):
@@ -50,9 +82,15 @@ class Request(Struct):
     def to_object(self):
         return _to_object(self.SCHEMA, self)
 
+    def build_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)
 
+
+@add_metaclass(abc.ABCMeta)
 class Response(Struct):
-    __metaclass__ = abc.ABCMeta
+    FLEXIBLE_VERSION = False
 
     @abc.abstractproperty
     def API_KEY(self):
@@ -72,6 +110,12 @@ class Response(Struct):
     def to_object(self):
         return _to_object(self.SCHEMA, self)
 
+    @classmethod
+    def parse_header(cls, read_buffer):
+        if cls.FLEXIBLE_VERSION:
+            return ResponseHeaderV2.decode(read_buffer)
+        return ResponseHeader.decode(read_buffer)
+
 
 def _to_object(schema, data):
     obj = {}