kafka-python 2.0.2__py2.py3-none-any.whl → 2.0.3__py2.py3-none-any.whl

kafka/__init__.py

@@ -4,7 +4,7 @@ __title__ = 'kafka'
 from kafka.version import __version__
 __author__ = 'Dana Powers'
 __license__ = 'Apache License 2.0'
-__copyright__ = 'Copyright 2016 Dana Powers, David Arthur, and Contributors'
+__copyright__ = 'Copyright 2025 Dana Powers, David Arthur, and Contributors'
 
 # Set default logging handler to avoid "No handler found" warnings.
 import logging

kafka/admin/client.py

@@ -20,7 +20,7 @@ from kafka.metrics import MetricConfig, Metrics
 from kafka.protocol.admin import (
     CreateTopicsRequest, DeleteTopicsRequest, DescribeConfigsRequest, AlterConfigsRequest, CreatePartitionsRequest,
     ListGroupsRequest, DescribeGroupsRequest, DescribeAclsRequest, CreateAclsRequest, DeleteAclsRequest,
-    DeleteGroupsRequest
+    DeleteGroupsRequest, DescribeLogDirsRequest
 )
 from kafka.protocol.commit import GroupCoordinatorRequest, OffsetFetchRequest
 from kafka.protocol.metadata import MetadataRequest
@@ -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
@@ -351,13 +355,14 @@ class KafkaAdminClient(object):
         }
         return groups_coordinators
 
-    def _send_request_to_node(self, node_id, request):
+    def _send_request_to_node(self, node_id, request, wakeup=True):
         """Send a Kafka protocol message to a specific broker.
 
         Returns a future that may be polled for status and results.
 
         :param node_id: The broker id to which to send the message.
         :param request: The message to send.
+        :param wakeup: Optional flag to disable thread-wakeup.
         :return: A future object that may be polled for status and results.
         :exception: The exception if the message could not be sent.
         """
@@ -365,7 +370,7 @@ class KafkaAdminClient(object):
             # poll until the connection to broker is ready, otherwise send()
             # will fail with NodeNotReadyError
             self._client.poll()
-        return self._client.send(node_id, request)
+        return self._client.send(node_id, request, wakeup)
 
     def _send_request_to_controller(self, request):
         """Send a Kafka protocol message to the cluster controller.
@@ -1205,7 +1210,7 @@ class KafkaAdminClient(object):
 
         :param response: an OffsetFetchResponse.
         :return: A dictionary composed of TopicPartition keys and
-            OffsetAndMetada values.
+            OffsetAndMetadata values.
         """
         if response.API_VERSION <= 3:
 
@@ -1219,7 +1224,7 @@ class KafkaAdminClient(object):
                         .format(response))
 
             # transform response into a dictionary with TopicPartition keys and
-            # OffsetAndMetada values--this is what the Java AdminClient returns
+            # OffsetAndMetadata values--this is what the Java AdminClient returns
             offsets = {}
             for topic, partitions in response.topics:
                 for partition, offset, metadata, error_code in partitions:
@@ -1340,3 +1345,19 @@ class KafkaAdminClient(object):
 
                 if future.failed():
                     raise future.exception  # pylint: disable-msg=raising-bad-type
+
+    def describe_log_dirs(self):
+        """Send a DescribeLogDirsRequest request to a broker.
+
+        :return: A message future
+        """
+        version = self._matching_api_version(DescribeLogDirsRequest)
+        if version <= 0:
+            request = DescribeLogDirsRequest[version]()
+            future = self._send_request_to_node(self._client.least_loaded_node(), request)
+            self._wait_for_futures([future])
+        else:
+            raise NotImplementedError(
+                "Support for DescribeLogDirsRequest_v{} has not yet been added to KafkaAdminClient."
+                    .format(version))
+        return future.value

kafka/codec.py

@@ -187,7 +187,7 @@ def _detect_xerial_stream(payload):
         The version is the version of this format as written by xerial,
         in the wild this is currently 1 as such we only support v1.
 
-        Compat is there to claim the miniumum supported version that
+        Compat is there to claim the minimum supported version that
         can read a xerial block stream, presently in the wild this is
         1.
     """

kafka/conn.py

@@ -24,7 +24,7 @@ import kafka.errors as Errors
 from kafka.future import Future
 from kafka.metrics.stats import Avg, Count, Max, Rate
 from kafka.oauth.abstract import AbstractTokenProvider
-from kafka.protocol.admin import SaslHandShakeRequest, DescribeAclsRequest_v2
+from kafka.protocol.admin import SaslHandShakeRequest, DescribeAclsRequest_v2, DescribeClientQuotasRequest
 from kafka.protocol.commit import OffsetFetchRequest
 from kafka.protocol.offset import OffsetRequest
 from kafka.protocol.produce import ProduceRequest
@@ -78,7 +78,7 @@ except ImportError:
 try:
     import gssapi
     from gssapi.raw.misc import GSSError
-except ImportError:
+except (ImportError, OSError):
     #no gssapi available, will disable gssapi mechanism
     gssapi = None
     GSSError = None
@@ -496,7 +496,7 @@ class BrokerConnection(object):
         try:
             self._sock = self._ssl_context.wrap_socket(
                 self._sock,
-                server_hostname=self.host,
+                server_hostname=self.host.rstrip("."),
                 do_handshake_on_connect=False)
         except ssl.SSLError as e:
             log.exception('%s: Failed to wrap socket in SSLContext!', self)
@@ -916,7 +916,7 @@ class BrokerConnection(object):
         with self._lock:
             if self.state is ConnectionStates.DISCONNECTED:
                 return
-            log.info('%s: Closing connection. %s', self, error or '')
+            log.log(logging.ERROR if error else logging.INFO, '%s: Closing connection. %s', self, error or '')
             self._update_reconnect_backoff()
             self._sasl_auth_future = None
             self._protocol = KafkaProtocol(
@@ -1169,6 +1169,7 @@ class BrokerConnection(object):
         # in reverse order. As soon as we find one that works, return it
         test_cases = [
             # format (<broker version>, <needed struct>)
+            ((2, 6, 0), DescribeClientQuotasRequest[0]),
             ((2, 5, 0), DescribeAclsRequest_v2),
             ((2, 4, 0), ProduceRequest[8]),
             ((2, 3, 0), FetchRequest[11]),

kafka/consumer/fetcher.py

@@ -125,7 +125,7 @@ class Fetcher(six.Iterator):
                 log.debug("Sending FetchRequest to node %s", node_id)
                 future = self._client.send(node_id, request, wakeup=False)
                 future.add_callback(self._handle_fetch_response, request, time.time())
-                future.add_errback(log.error, 'Fetch to node %s failed: %s', node_id)
+                future.add_errback(self._handle_fetch_error, node_id)
                 futures.append(future)
         self._fetch_futures.extend(futures)
         self._clean_done_fetch_futures()
@@ -778,6 +778,14 @@ class Fetcher(six.Iterator):
             self._sensors.fetch_throttle_time_sensor.record(response.throttle_time_ms)
         self._sensors.fetch_latency.record((time.time() - send_time) * 1000)
 
+    def _handle_fetch_error(self, node_id, exception):
+        log.log(
+            logging.INFO if isinstance(exception, Errors.Cancelled) else logging.ERROR,
+            'Fetch to node %s failed: %s',
+            node_id,
+            exception
+        )
+
     def _parse_fetched_data(self, completed_fetch):
         tp = completed_fetch.topic_partition
         fetch_offset = completed_fetch.fetched_offset
@@ -817,8 +825,9 @@ class Fetcher(six.Iterator):
                               position)
                     unpacked = list(self._unpack_message_set(tp, records))
                     parsed_records = self.PartitionRecords(fetch_offset, tp, unpacked)
-                    last_offset = unpacked[-1].offset
-                    self._sensors.records_fetch_lag.record(highwater - last_offset)
+                    if unpacked:
+                        last_offset = unpacked[-1].offset
+                        self._sensors.records_fetch_lag.record(highwater - last_offset)
                     num_bytes = records.valid_bytes()
                     records_count = len(unpacked)
                 elif records.size_in_bytes() > 0:

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:
@@ -651,7 +653,7 @@ class KafkaConsumer(six.Iterator):
         # Poll for new data until the timeout expires
         start = time.time()
         remaining = timeout_ms
-        while True:
+        while not self._closed:
             records = self._poll_once(remaining, max_records, update_offsets=update_offsets)
             if records:
                 return records
@@ -660,7 +662,9 @@ class KafkaConsumer(six.Iterator):
             remaining = timeout_ms - elapsed_ms
 
             if remaining <= 0:
-                return {}
+                break
+
+        return {}
 
     def _poll_once(self, timeout_ms, max_records, update_offsets=True):
         """Do one round of polling. In addition to checking for new data, this does

kafka/coordinator/assignors/sticky/sticky_assignor.py

@@ -648,15 +648,19 @@ class StickyPartitionAssignor(AbstractPartitionAssignor):
 
     @classmethod
     def metadata(cls, topics):
-        if cls.member_assignment is None:
+        return cls._metadata(topics, cls.member_assignment, cls.generation)
+
+    @classmethod
+    def _metadata(cls, topics, member_assignment_partitions, generation=-1):
+        if member_assignment_partitions is None:
             log.debug("No member assignment available")
             user_data = b''
         else:
             log.debug("Member assignment is available, generating the metadata: generation {}".format(cls.generation))
             partitions_by_topic = defaultdict(list)
-            for topic_partition in cls.member_assignment:   # pylint: disable=not-an-iterable
+            for topic_partition in member_assignment_partitions:
                 partitions_by_topic[topic_partition.topic].append(topic_partition.partition)
-            data = StickyAssignorUserDataV1(six.iteritems(partitions_by_topic), cls.generation)
+            data = StickyAssignorUserDataV1(six.viewitems(partitions_by_topic), generation)
             user_data = data.encode()
         return ConsumerProtocolMemberMetadata(cls.version, list(topics), user_data)
 

kafka/coordinator/base.py

@@ -952,7 +952,7 @@ class HeartbeatThread(threading.Thread):
                 # disable here to prevent propagating an exception to this
                 # heartbeat thread
                 # must get client._lock, or maybe deadlock at heartbeat 
-                # failure callbak in consumer poll
+                # failure callback in consumer poll
                 self.coordinator._client.poll(timeout_ms=0)
 
         with self.coordinator._lock:
@@ -990,6 +990,11 @@ class HeartbeatThread(threading.Thread):
                 # foreground thread has stalled in between calls to
                 # poll(), so we explicitly leave the group.
                 log.warning('Heartbeat poll expired, leaving group')
+                ### XXX
+                # maybe_leave_group acquires client + coordinator lock;
+                # if we hold coordinator lock before calling, we risk deadlock
+                # release() is safe here because this is the last code in the current context
+                self.coordinator._lock.release()
                 self.coordinator.maybe_leave_group()
 
             elif not self.coordinator.heartbeat.should_heartbeat():

kafka/producer/kafka.py

@@ -233,7 +233,7 @@ class KafkaProducer(object):
             should verify that the certificate matches the brokers hostname.
             default: true.
         ssl_cafile (str): optional filename of ca file to use in certificate
-            veriication. default: none.
+            verification. default: none.
         ssl_certfile (str): optional filename of file in pem format containing
             the client certificate, as well as any ca certificates needed to
             establish the certificate's authenticity. default: none.
@@ -280,10 +280,11 @@ 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
-        https://kafka.apache.org/0100/configuration.html#producerconfigs
+        https://kafka.apache.org/0100/documentation/#producerconfigs
     """
     DEFAULT_CONFIG = {
         'bootstrap_servers': 'localhost',
@@ -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:

kafka/protocol/__init__.py

@@ -43,4 +43,7 @@ API_KEYS = {
     40: 'ExpireDelegationToken',
     41: 'DescribeDelegationToken',
     42: 'DeleteGroups',
+    45: 'AlterPartitionReassignments',
+    46: 'ListPartitionReassignments',
+    48: 'DescribeClientQuotas',
 }

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
+from kafka.protocol.types import Array, Boolean, Bytes, Int8, Int16, Int32, Int64, Schema, String, Float64, CompactString, CompactArray, TaggedFields
 
 
 class ApiVersionResponse_v0(Response):
@@ -719,7 +719,7 @@ class DescribeConfigsResponse_v1(Response):
                 ('config_names', String('utf-8')),
                 ('config_value', String('utf-8')),
                 ('read_only', Boolean),
-                ('is_default', Boolean),
+                ('config_source', Int8),
                 ('is_sensitive', Boolean),
                 ('config_synonyms', Array(
                     ('config_name', String('utf-8')),
@@ -790,6 +790,48 @@ DescribeConfigsResponse = [
 ]
 
 
+class DescribeLogDirsResponse_v0(Response):
+    API_KEY = 35
+    API_VERSION = 0
+    FLEXIBLE_VERSION = True
+    SCHEMA = Schema(
+        ('throttle_time_ms', Int32),
+        ('log_dirs', Array(
+            ('error_code', Int16),
+            ('log_dir', String('utf-8')),
+            ('topics', Array(
+                ('name', String('utf-8')),
+                ('partitions', Array(
+                    ('partition_index', Int32),
+                    ('partition_size', Int64),
+                    ('offset_lag', Int64),
+                    ('is_future_key', Boolean)
+                ))
+            ))
+        ))
+    )
+
+
+class DescribeLogDirsRequest_v0(Request):
+    API_KEY = 35
+    API_VERSION = 0
+    RESPONSE_TYPE = DescribeLogDirsResponse_v0
+    SCHEMA = Schema(
+                     ('topics', Array(
+                         ('topic', String('utf-8')),
+                         ('partitions', Int32)
+                         ))
+                 )
+
+
+DescribeLogDirsResponse = [
+    DescribeLogDirsResponse_v0,
+]
+DescribeLogDirsRequest = [
+    DescribeLogDirsRequest_v0,
+]
+
+
 class SaslAuthenticateResponse_v0(Response):
     API_KEY = 36
     API_VERSION = 0
@@ -923,3 +965,132 @@ DeleteGroupsRequest = [
 DeleteGroupsResponse = [
     DeleteGroupsResponse_v0, DeleteGroupsResponse_v1
 ]
+
+
+class DescribeClientQuotasResponse_v0(Response):
+    API_KEY = 48
+    API_VERSION = 0
+    SCHEMA = Schema(
+        ('throttle_time_ms', Int32),
+        ('error_code', Int16),
+        ('error_message', String('utf-8')),
+        ('entries', Array(
+            ('entity', Array(
+                ('entity_type', String('utf-8')),
+                ('entity_name', String('utf-8')))),
+            ('values', Array(
+                ('name', String('utf-8')),
+                ('value', Float64))))),
+    )
+
+
+class DescribeClientQuotasRequest_v0(Request):
+    API_KEY = 48
+    API_VERSION = 0
+    RESPONSE_TYPE = DescribeClientQuotasResponse_v0
+    SCHEMA = Schema(
+        ('components', Array(
+            ('entity_type', String('utf-8')),
+            ('match_type', Int8),
+            ('match', String('utf-8')),
+        )),
+        ('strict', Boolean)
+    )
+
+
+DescribeClientQuotasRequest = [
+    DescribeClientQuotasRequest_v0,
+]
+
+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]

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

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