Fix tagged fields handling

* Use RequestHeader_v1/ResponseHeader_v1 for flexible version
* Fix DeleteRecords for version 2+

Author: ods

aiokafka/admin/client.py

@@ -649,8 +649,9 @@ async def delete_records(
                 timeout_ms or self._request_timeout_ms,
             )
             response = await self._client.send(leader, request)
-            for topic, partitions in response.topics:
-                for partition_index, low_watermark, error_code in partitions:
+            # Starting with v2, DeleteRecordsResponse contains extra field with tags
+            for topic, partitions, *_ in response.topics:
+                for partition_index, low_watermark, error_code, *_ in partitions:
                     if error_code:
                         err = for_code(error_code)
                         raise err

aiokafka/conn.py

@@ -4,6 +4,7 @@
 import functools
 import hashlib
 import hmac
+import io
 import logging
 import random
 import socket
@@ -24,7 +25,6 @@
     SaslAuthenticateRequest,
     SaslHandShakeRequest,
 )
-from aiokafka.protocol.api import RequestHeader
 from aiokafka.protocol.commit import (
     GroupCoordinatorResponse_v0 as GroupCoordinatorResponse,
 )
@@ -459,10 +459,8 @@ def send(self, request, expect_response=True):
             )
 
         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 = header.encode() + request.encode()
         size = struct.pack(">i", len(message))
@@ -480,7 +478,7 @@ def send(self, request, expect_response=True):
             return self._writer.drain()
         fut = self._loop.create_future()
         self._requests.append(
-            (correlation_id, request.RESPONSE_TYPE, fut),
+            (correlation_id, request, fut),
         )
         return wait_for(fut, self._request_timeout)
 
@@ -569,39 +567,41 @@ async def _read(self_ref):
             del self
 
     def _handle_frame(self, resp):
-        correlation_id, resp_type, fut = self._requests[0]
+        correlation_id, request, fut = self._requests[0]
 
         if correlation_id is None:  # Is a SASL packet, just pass it though
             if not fut.done():
                 fut.set_result(resp)
         else:
-            (recv_correlation_id,) = struct.unpack_from(">i", resp, 0)
+            resp = io.BytesIO(resp)
+            response_header = request.parse_response_header(resp)
+            resp_type = request.RESPONSE_TYPE
 
             if (
                 self._api_version == (0, 8, 2)
                 and resp_type is GroupCoordinatorResponse
                 and correlation_id != 0
-                and recv_correlation_id == 0
+                and response_header.correlation_id == 0
             ):
                 log.warning(
                     "Kafka 0.8.2 quirk -- GroupCoordinatorResponse"
-                    " coorelation id does not match request. This"
+                    " correlation id does not match request. This"
                     " should go away once at least one topic has been"
                     " initialized on the broker"
                 )
 
-            elif correlation_id != recv_correlation_id:
+            elif response_header.correlation_id != correlation_id:
                 error = Errors.CorrelationIdError(
                     f"Correlation ids do not match: sent {correlation_id},"
-                    f" recv {recv_correlation_id}"
+                    f" recv {response_header.correlation_id}"
                 )
                 if not fut.done():
                     fut.set_exception(error)
                 self.close(reason=CloseReason.OUT_OF_SYNC)
                 return
 
             if not fut.done():
-                response = resp_type.decode(resp[4:])
+                response = resp_type.decode(resp)
                 log.debug("%s Response %d: %s", self, correlation_id, response)
                 fut.set_result(response)
 

aiokafka/protocol/admin.py

@@ -1385,17 +1385,32 @@ class DeleteRecordsRequest_v2(Request):
         ("tags", TaggedFields),
     )
 
+    def __init__(self, topics, timeout_ms, tags=None):
+        super().__init__(
+            [
+                (
+                    topic,
+                    [
+                        (partition, before_offset, {})
+                        for partition, before_offset in partitions
+                    ],
+                    {},
+                )
+                for (topic, partitions) in topics
+            ],
+            timeout_ms,
+            tags or {},
+        )
+
 
 DeleteRecordsRequest = [
     DeleteRecordsRequest_v0,
     DeleteRecordsRequest_v1,
-    # FIXME: We have some problems with `TaggedFields`
-    # DeleteRecordsRequest_v2,
+    DeleteRecordsRequest_v2,
 ]
 
 DeleteRecordsResponse = [
     DeleteRecordsResponse_v0,
     DeleteRecordsResponse_v1,
-    # FIXME: We have some problems with `TaggedFields`
-    # DeleteRecordsResponse_v2,
+    DeleteRecordsResponse_v2,
 ]

aiokafka/protocol/types.py

@@ -313,7 +313,6 @@ def encode(self, value):
         return UnsignedVarInt32.encode(len(value) + 1) + value
 
 
-# FIXME: TaggedFields doesn't seem to work properly so they should be avoided
 class TaggedFields(AbstractType):
     @classmethod
     def decode(cls, data):