kafka-python 3.0.0__py3-none-any.whl

kafka/protocol/parser.py

@@ -0,0 +1,191 @@
+import collections
+import logging
+import struct
+
+import kafka.errors as Errors
+from kafka.protocol.metadata import FindCoordinatorResponse
+from kafka.protocol.frame import KafkaBytes
+from kafka.protocol.schemas.fields.codecs import Int32
+from kafka.version import __version__
+
+log = logging.getLogger(__name__)
+
+
+class KafkaProtocol:
+    """Manage the kafka network protocol
+
+    Use an instance of KafkaProtocol to manage bytes send/recv'd
+    from a network socket to a broker.
+
+    Arguments:
+        client_id (str): identifier string to be included in each request
+        ident (str): Optional log-prefix identifier.
+        receive_message_max_bytes (int): Maximum allowed message frame size.
+            Default: 100000000 (100MB).
+    """
+    def __init__(self, **config):
+        self._ident = config.get('ident', '')
+        self._client_id = config.get('client_id', self._gen_client_id())
+        self._max_frame_size = config.get('receive_message_max_bytes', 100000000)
+        self._correlation_id = 0
+        self._header = KafkaBytes(4)
+        self._rbuffer = None
+        self._receiving = False
+        self.in_flight_requests = collections.deque()
+        self.bytes_to_send = []
+
+    def _next_correlation_id(self):
+        self._correlation_id = (self._correlation_id + 1) % 2**31
+        return self._correlation_id
+
+    def _gen_client_id(self):
+        return 'kafka-python' + __version__
+
+    def send_request(self, request, correlation_id=None):
+        """Encode and queue a kafka api request for sending.
+
+        Arguments:
+            request : An un-encoded kafka request.
+            correlation_id (int, optional): Optionally specify an ID to
+                correlate requests with responses. If not provided, an ID will
+                be generated automatically.
+
+        Returns:
+            correlation_id
+        """
+        if correlation_id is None:
+            correlation_id = self._next_correlation_id()
+
+        log.debug('%s Sending request %d %s', self._ident, correlation_id, request)
+        request.with_header(correlation_id=correlation_id, client_id=self._client_id)
+        data = request.encode(framed=True, header=True)
+        self.bytes_to_send.append(data)
+        if request.expect_response():
+            self.in_flight_requests.append(request.header)
+        return correlation_id
+
+    def send_bytes(self):
+        """Retrieve all pending bytes to send on the network"""
+        # Short-circuit the common single-request case to avoid an extra
+        # full-request copy through b''.join.
+        n = len(self.bytes_to_send)
+        if n == 0:
+            return b''
+        if n == 1:
+            data = self.bytes_to_send[0]
+            self.bytes_to_send = []
+        else:
+            data = b''.join(self.bytes_to_send)
+            self.bytes_to_send = []
+        log.debug('%s Send: %r', self._ident, data)
+        return data
+
+    def receive_bytes(self, data):
+        """Process bytes received from the network.
+
+        Arguments:
+            data (bytes): any length bytes received from a network connection
+                to a kafka broker.
+
+        Returns:
+            responses (list of (correlation_id, response)): any/all completed
+                responses, decoded from bytes to python objects.
+
+        Raises:
+             KafkaProtocolError: if the bytes received could not be decoded.
+             CorrelationIdError: if the response does not match the request
+                 correlation id.
+        """
+        i = 0
+        n = len(data)
+        responses = []
+        if data:
+            log.debug('%s Recv: %r', self._ident, data)
+        while i < n:
+
+            # Not receiving is the state of reading the payload header
+            if not self._receiving:
+                bytes_to_read = min(4 - self._header.tell(), n - i)
+                self._header.write(data[i:i+bytes_to_read])
+                i += bytes_to_read
+
+                if self._header.tell() == 4:
+                    self._header.seek(0)
+                    nbytes = Int32.decode(self._header)
+                    self._validate_frame_size(nbytes)
+                    # reset buffer and switch state to receiving payload bytes
+                    self._rbuffer = KafkaBytes(nbytes)
+                    self._receiving = True
+                elif self._header.tell() > 4:
+                    raise Errors.KafkaError('this should not happen - are you threading?')
+
+            if self._receiving:
+                total_bytes = len(self._rbuffer)
+                staged_bytes = self._rbuffer.tell()
+                bytes_to_read = min(total_bytes - staged_bytes, n - i)
+                self._rbuffer.write(data[i:i+bytes_to_read])
+                i += bytes_to_read
+
+                staged_bytes = self._rbuffer.tell()
+                if staged_bytes > total_bytes:
+                    raise Errors.KafkaError('Receive buffer has more bytes than expected?')
+
+                if staged_bytes != total_bytes:
+                    break
+
+                self._receiving = False
+                self._rbuffer.seek(0)
+                resp = self._process_response(self._rbuffer)
+                responses.append(resp)
+                self._reset_buffer()
+        return responses
+
+    def _validate_frame_size(self, nbytes):
+        if nbytes < 0 or nbytes > self._max_frame_size:
+            raise Errors.InvalidReceiveError('Invalid frame length: %d' % nbytes)
+
+    def _process_response(self, read_buffer):
+        if not self.in_flight_requests:
+            raise Errors.CorrelationIdError('No in-flight-request found for server response')
+        header = self.in_flight_requests.popleft()
+        correlation_id = header.correlation_id
+        response_type = header.get_response_class()
+        if response_type is None:
+            log.error('Unable to find ResponseType for api=%d version=%d',
+                      header.api_key, header.api_version)
+            raise Errors.KafkaProtocolError('Unable to find response type for api %d v%d' % (header.api_key, header.api_version))
+        response_header = response_type.parse_header(read_buffer)
+        recv_correlation_id = response_header.correlation_id
+        # ignore correlation id mismatch for 0.8.2 quirk
+        if (recv_correlation_id == 0 and correlation_id != 0 and
+            response_type is FindCoordinatorResponse and header.api_version == 0):
+            log.warning('Kafka 0.8.2 quirk -- FindCoordinatorResponse'
+                        ' 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:
+            # return or raise?
+            raise Errors.CorrelationIdError(
+                'Correlation IDs do not match: sent %d, recv %d'
+                % (correlation_id, recv_correlation_id))
+
+        # decode response
+        try:
+            response = response_type.decode(read_buffer)
+        except (ValueError, struct.error):
+            read_buffer.seek(0)
+            buf = read_buffer.read()
+            log.error('Response %d [ResponseType: %s RequestHeader: %s]:'
+                      ' Unable to decode %d-byte buffer: %r',
+                      correlation_id, response_type,
+                      header, len(buf), buf)
+            raise Errors.KafkaProtocolError('Unable to decode response')
+
+        log.debug('%s Received response %d %s', self._ident, correlation_id, response)
+        return (correlation_id, response)
+
+    def _reset_buffer(self):
+        self._receiving = False
+        self._header.seek(0)
+        self._rbuffer = None

kafka/protocol/producer/__init__.py

@@ -0,0 +1,7 @@
+from .produce import *
+from .produce import __all__ as produce_all
+
+from .transaction import *
+from .transaction import __all__ as transaction_all
+
+__all__ = produce_all + transaction_all

kafka/protocol/producer/produce.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from ..api_message import ApiMessage
+
+
+class ProduceRequest(ApiMessage):
+    def expect_response(self) -> bool:
+        if self.acks == 0: # pylint: disable=no-member  # ty: ignore[unresolved-attribute]
+            return False
+        return True
+
+class ProduceResponse(ApiMessage): pass
+
+
+__all__ = [
+    'ProduceRequest', 'ProduceResponse',
+]

kafka/protocol/producer/produce.pyi

@@ -0,0 +1,197 @@
+# Generated by generate_stubs.py (Python 3.14)
+import uuid
+from typing import Any, Self
+
+from kafka.protocol.api_message import ApiMessage
+from kafka.protocol.api_data import ApiData
+from kafka.protocol.data_container import DataContainer
+
+__all__ = ['ProduceRequest', 'ProduceResponse']
+
+class ProduceRequest(ApiMessage):
+    class TopicProduceData(DataContainer):
+        class PartitionProduceData(DataContainer):
+            index: int
+            records: bytes | ApiData | None
+            def __init__(
+                self,
+                *args: Any,
+                index: int = ...,
+                records: bytes | ApiData | None = ...,
+                version: int | None = None,
+                **kwargs: Any,
+            ) -> None: ...
+            @property
+            def version(self) -> int | None: ...
+            def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+        name: str
+        topic_id: uuid.UUID
+        partition_data: list[PartitionProduceData]
+        def __init__(
+            self,
+            *args: Any,
+            name: str = ...,
+            topic_id: uuid.UUID = ...,
+            partition_data: list[PartitionProduceData] = ...,
+            version: int | None = None,
+            **kwargs: Any,
+        ) -> None: ...
+        @property
+        def version(self) -> int | None: ...
+        def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+    transactional_id: str | None
+    acks: int
+    timeout_ms: int
+    topic_data: list[TopicProduceData]
+    def __init__(
+        self,
+        *args: Any,
+        transactional_id: str | None = ...,
+        acks: int = ...,
+        timeout_ms: int = ...,
+        topic_data: list[TopicProduceData] = ...,
+        version: int | None = None,
+        **kwargs: Any,
+    ) -> None: ...
+    @property
+    def version(self) -> int | None: ...
+    def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+    name: str
+    type: str
+    API_KEY: int
+    API_VERSION: int
+    valid_versions: tuple[int, int]
+    min_version: int
+    max_version: int
+    @property
+    def header(self) -> Any: ...
+    @classmethod
+    def is_request(cls) -> bool: ...
+    def expect_response(self) -> bool: ...
+    def with_header(self, correlation_id: int = 0, client_id: str = "kafka-python") -> None: ...
+    def expect_response(self) -> bool: ...
+
+class ProduceResponse(ApiMessage):
+    class TopicProduceResponse(DataContainer):
+        class PartitionProduceResponse(DataContainer):
+            class BatchIndexAndErrorMessage(DataContainer):
+                batch_index: int
+                batch_index_error_message: str | None
+                def __init__(
+                    self,
+                    *args: Any,
+                    batch_index: int = ...,
+                    batch_index_error_message: str | None = ...,
+                    version: int | None = None,
+                    **kwargs: Any,
+                ) -> None: ...
+                @property
+                def version(self) -> int | None: ...
+                def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+            class LeaderIdAndEpoch(DataContainer):
+                leader_id: int
+                leader_epoch: int
+                def __init__(
+                    self,
+                    *args: Any,
+                    leader_id: int = ...,
+                    leader_epoch: int = ...,
+                    version: int | None = None,
+                    **kwargs: Any,
+                ) -> None: ...
+                @property
+                def version(self) -> int | None: ...
+                def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+            index: int
+            error_code: int
+            base_offset: int
+            log_append_time_ms: int
+            log_start_offset: int
+            record_errors: list[BatchIndexAndErrorMessage]
+            error_message: str | None
+            current_leader: LeaderIdAndEpoch
+            def __init__(
+                self,
+                *args: Any,
+                index: int = ...,
+                error_code: int = ...,
+                base_offset: int = ...,
+                log_append_time_ms: int = ...,
+                log_start_offset: int = ...,
+                record_errors: list[BatchIndexAndErrorMessage] = ...,
+                error_message: str | None = ...,
+                current_leader: LeaderIdAndEpoch = ...,
+                version: int | None = None,
+                **kwargs: Any,
+            ) -> None: ...
+            @property
+            def version(self) -> int | None: ...
+            def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+        name: str
+        topic_id: uuid.UUID
+        partition_responses: list[PartitionProduceResponse]
+        def __init__(
+            self,
+            *args: Any,
+            name: str = ...,
+            topic_id: uuid.UUID = ...,
+            partition_responses: list[PartitionProduceResponse] = ...,
+            version: int | None = None,
+            **kwargs: Any,
+        ) -> None: ...
+        @property
+        def version(self) -> int | None: ...
+        def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+    class NodeEndpoint(DataContainer):
+        node_id: int
+        host: str
+        port: int
+        rack: str | None
+        def __init__(
+            self,
+            *args: Any,
+            node_id: int = ...,
+            host: str = ...,
+            port: int = ...,
+            rack: str | None = ...,
+            version: int | None = None,
+            **kwargs: Any,
+        ) -> None: ...
+        @property
+        def version(self) -> int | None: ...
+        def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+
+    responses: list[TopicProduceResponse]
+    throttle_time_ms: int
+    node_endpoints: list[NodeEndpoint]
+    def __init__(
+        self,
+        *args: Any,
+        responses: list[TopicProduceResponse] = ...,
+        throttle_time_ms: int = ...,
+        node_endpoints: list[NodeEndpoint] = ...,
+        version: int | None = None,
+        **kwargs: Any,
+    ) -> None: ...
+    @property
+    def version(self) -> int | None: ...
+    def to_dict(self, meta: bool = False, json: bool = True) -> dict: ...
+    name: str
+    type: str
+    API_KEY: int
+    API_VERSION: int
+    valid_versions: tuple[int, int]
+    min_version: int
+    max_version: int
+    @property
+    def header(self) -> Any: ...
+    @classmethod
+    def is_request(cls) -> bool: ...
+    def expect_response(self) -> bool: ...
+    def with_header(self, correlation_id: int = 0, client_id: str = "kafka-python") -> None: ...

kafka/protocol/producer/transaction.py

@@ -0,0 +1,30 @@
+from ..api_message import ApiMessage
+
+
+class InitProducerIdRequest(ApiMessage): pass
+class InitProducerIdResponse(ApiMessage): pass
+
+class AddPartitionsToTxnRequest(ApiMessage): pass
+class AddPartitionsToTxnResponse(ApiMessage): pass
+
+class AddOffsetsToTxnRequest(ApiMessage): pass
+class AddOffsetsToTxnResponse(ApiMessage): pass
+
+class EndTxnRequest(ApiMessage): pass
+class EndTxnResponse(ApiMessage): pass
+
+class TxnOffsetCommitRequest(ApiMessage): pass
+class TxnOffsetCommitResponse(ApiMessage): pass
+
+class WriteTxnMarkersRequest(ApiMessage): pass
+class WriteTxnMarkersResponse(ApiMessage): pass
+
+
+__all__ = [
+    'InitProducerIdRequest', 'InitProducerIdResponse',
+    'AddPartitionsToTxnRequest', 'AddPartitionsToTxnResponse',
+    'AddOffsetsToTxnRequest', 'AddOffsetsToTxnResponse',
+    'EndTxnRequest', 'EndTxnResponse',
+    'TxnOffsetCommitRequest', 'TxnOffsetCommitResponse',
+    'WriteTxnMarkersRequest', 'WriteTxnMarkersResponse',
+]