kafka-python 3.0.0__py3-none-any.whl

kafka/record/memory_records.py

@@ -0,0 +1,255 @@
+# This class takes advantage of the fact that all formats v0, v1 and v2 of
+# messages storage has the same byte offsets for Length and Magic fields.
+# Lets look closely at what leading bytes all versions have:
+#
+# V0 and V1 (Offset is MessageSet part, other bytes are Message ones):
+#  Offset => Int64
+#  BytesLength => Int32
+#  CRC => Int32
+#  Magic => Int8
+#  ...
+#
+# V2:
+#  BaseOffset => Int64
+#  Length => Int32
+#  PartitionLeaderEpoch => Int32
+#  Magic => Int8
+#  ...
+#
+# So we can iterate over batches just by knowing offsets of Length. Magic is
+# used to construct the correct class for Batch itself.
+
+import struct
+
+from kafka.errors import CorruptRecordError, IllegalStateError, UnsupportedVersionError
+from kafka.record.abc import ABCRecords
+from kafka.record.legacy_records import LegacyRecordBatch, LegacyRecordBatchBuilder
+from kafka.record.default_records import DefaultRecordBatch, DefaultRecordBatchBuilder
+
+
+class MemoryRecords(ABCRecords):
+
+    LENGTH_OFFSET = struct.calcsize(">q")
+    LOG_OVERHEAD = struct.calcsize(">qi")
+    MAGIC_OFFSET = struct.calcsize(">qii")
+
+    # Minimum space requirements for Record V0
+    MIN_SLICE = LOG_OVERHEAD + LegacyRecordBatch.RECORD_OVERHEAD_V0
+
+    __slots__ = ("_buffer", "_pos", "_next_slice", "_remaining_bytes")
+
+    def __init__(self, bytes_data):
+        self._buffer = bytes_data
+        self._pos = 0
+        # We keep one slice ahead so `has_next` will return very fast
+        self._next_slice = None
+        self._remaining_bytes = None
+        self._cache_next()
+
+    def size_in_bytes(self):
+        return len(self._buffer)
+
+    def valid_bytes(self):
+        # We need to read the whole buffer to get the valid_bytes.
+        # NOTE: in Fetcher we do the call after iteration, so should be fast
+        if self._remaining_bytes is None:
+            next_slice = self._next_slice
+            pos = self._pos
+            while self._remaining_bytes is None:
+                self._cache_next()
+            # Reset previous iterator position
+            self._next_slice = next_slice
+            self._pos = pos
+        return len(self._buffer) - self._remaining_bytes
+
+    # NOTE: we cache offsets here as kwargs for a bit more speed, as cPython
+    # will use LOAD_FAST opcode in this case
+    def _cache_next(self, len_offset=LENGTH_OFFSET, log_overhead=LOG_OVERHEAD):
+        buffer = self._buffer
+        buffer_len = len(buffer)
+        pos = self._pos
+        remaining = buffer_len - pos
+        if remaining < log_overhead:
+            # Will be re-checked in Fetcher for remaining bytes.
+            self._remaining_bytes = remaining
+            self._next_slice = None
+            return
+
+        length, = struct.unpack_from(
+            ">i", buffer, pos + len_offset)
+
+        slice_end = pos + log_overhead + length
+        if slice_end > buffer_len:
+            # Will be re-checked in Fetcher for remaining bytes
+            self._remaining_bytes = remaining
+            self._next_slice = None
+            return
+
+        self._next_slice = memoryview(buffer)[pos: slice_end]
+        self._pos = slice_end
+
+    def has_next(self):
+        return self._next_slice is not None
+
+    # NOTE: same cache for LOAD_FAST as above
+    def next_batch(self, _min_slice=MIN_SLICE,
+                   _magic_offset=MAGIC_OFFSET):
+        next_slice = self._next_slice
+        if next_slice is None:
+            return None
+        if len(next_slice) < _min_slice:
+            raise CorruptRecordError(
+                "Record size is less than the minimum record overhead "
+                "({})".format(_min_slice - self.LOG_OVERHEAD))
+        self._cache_next()
+        magic, = struct.unpack_from(">b", next_slice, _magic_offset)
+        if magic <= 1:
+            return LegacyRecordBatch(next_slice, magic)
+        else:
+            return DefaultRecordBatch(next_slice)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if not self.has_next():
+            raise StopIteration
+        return self.next_batch()
+
+    next = __next__
+
+
+class MemoryRecordsBuilder:
+
+    __slots__ = ("_builder", "_batch_size", "_buffer", "_next_offset", "_closed",
+                 "_magic", "_bytes_written", "_producer_id", "_producer_epoch",
+                 "_base_sequence")
+
+    def __init__(self, magic, compression_type, batch_size, offset=0,
+                 transactional=False, producer_id=-1, producer_epoch=-1, base_sequence=-1):
+        assert magic in [0, 1, 2], "Not supported magic"
+        assert compression_type in [0, 1, 2, 3, 4], "Not valid compression type"
+        if magic >= 2:
+            assert not transactional or producer_id != -1, "Cannot write transactional messages without a valid producer ID"
+            assert producer_id == -1 or producer_epoch != -1, "Invalid negative producer epoch"
+            assert producer_id == -1 or base_sequence != -1, "Invalid negative sequence number used"
+
+            self._builder = DefaultRecordBatchBuilder(
+                magic=magic, compression_type=compression_type,
+                is_transactional=transactional, producer_id=producer_id,
+                producer_epoch=producer_epoch, base_sequence=base_sequence,
+                batch_size=batch_size)
+            self._producer_id = producer_id
+            self._producer_epoch = producer_epoch
+            self._base_sequence = base_sequence
+        else:
+            assert not transactional and producer_id == -1, "Idempotent messages are not supported for magic %s" % (magic,)
+            self._builder = LegacyRecordBatchBuilder(
+                magic=magic, compression_type=compression_type,
+                batch_size=batch_size)
+            self._producer_id = None
+            self._producer_epoch = None
+            self._base_sequence = None
+        self._batch_size = batch_size
+        self._buffer = None
+
+        self._next_offset = offset
+        self._closed = False
+        self._magic = magic
+        self._bytes_written = 0
+
+    def skip(self, offsets_to_skip):
+        # Exposed for testing compacted records
+        self._next_offset += offsets_to_skip
+
+    def append(self, timestamp, key, value, headers=None):
+        """ Append a message to the buffer.
+
+        Returns: RecordMetadata or None if unable to append
+        """
+        if headers is None:
+            headers = []
+        if self._closed:
+            return None
+
+        offset = self._next_offset
+        metadata = self._builder.append(offset, timestamp, key, value, headers)
+        # Return of None means there's no space to add a new message
+        if metadata is None:
+            return None
+
+        self._next_offset += 1
+        return metadata
+
+    def set_producer_state(self, producer_id, producer_epoch, base_sequence, is_transactional):
+        if self._magic < 2:
+            raise UnsupportedVersionError('Producer State requires Message format v2+')
+        elif self._closed:
+            # Sequence numbers are assigned when the batch is closed while the accumulator is being drained.
+            # If the resulting ProduceRequest to the partition leader failed for a retriable error, the batch will
+            # be re queued. In this case, we should not attempt to set the state again, since changing the pid and sequence
+            # once a batch has been sent to the broker risks introducing duplicates.
+            raise IllegalStateError("Trying to set producer state of an already closed batch. This indicates a bug on the client.")
+        self._builder.set_producer_state(producer_id, producer_epoch, base_sequence, is_transactional)
+        self._producer_id = producer_id
+        self._producer_epoch = producer_epoch
+        self._base_sequence = base_sequence
+
+    @property
+    def producer_id(self):
+        return self._producer_id
+
+    @property
+    def producer_epoch(self):
+        return self._producer_epoch
+
+    @property
+    def base_sequence(self):
+        return self._base_sequence
+
+    def records(self):
+        assert self._closed
+        return MemoryRecords(self._buffer)
+
+    def close(self):
+        # This method may be called multiple times on the same batch
+        # i.e., on retries
+        # we need to make sure we only close it out once
+        # otherwise compressed messages may be double-compressed
+        # see Issue 718
+        if not self._closed:
+            self._bytes_written = self._builder.size()
+            # Keep the buffer as bytearray to avoid a full-batch copy on
+            # close. Downstream consumers (MemoryRecords via memoryview and
+            # the protocol encoder via slice-assignment) handle bytearray
+            # without further copies.
+            self._buffer = self._builder.build()
+            if self._magic == 2:
+                self._producer_id = self._builder.producer_id
+                self._producer_epoch = self._builder.producer_epoch
+                self._base_sequence = self._builder.base_sequence
+            self._builder = None
+        self._closed = True
+
+    def size_in_bytes(self):
+        if not self._closed:
+            return self._builder.size()
+        else:
+            return len(self._buffer)
+
+    def compression_rate(self):
+        assert self._closed
+        return self.size_in_bytes() / self._bytes_written
+
+    def is_full(self):
+        if self._closed:
+            return True
+        else:
+            return self._builder.size() >= self._batch_size
+
+    def next_offset(self):
+        return self._next_offset
+
+    def buffer(self):
+        assert self._closed
+        return self._buffer

kafka/record/util.py

@@ -0,0 +1,135 @@
+import binascii
+
+from kafka.record._crc32c import crc as crc32c_py
+try:
+    from crc32c import crc32c as crc32c_c
+except ImportError:
+    crc32c_c = None
+
+
+def encode_varint(value, write):
+    """ Encode an integer to a varint presentation. See
+    https://developers.google.com/protocol-buffers/docs/encoding?csw=1#varints
+    on how those can be produced.
+
+        Arguments:
+            value (int): Value to encode
+            write (function): Called per byte that needs to be writen
+
+        Returns:
+            int: Number of bytes written
+    """
+    value = (value << 1) ^ (value >> 63)
+
+    if value <= 0x7f:  # 1 byte
+        write(value)
+        return 1
+    if value <= 0x3fff:  # 2 bytes
+        write(0x80 | (value & 0x7f))
+        write(value >> 7)
+        return 2
+    if value <= 0x1fffff:  # 3 bytes
+        write(0x80 | (value & 0x7f))
+        write(0x80 | ((value >> 7) & 0x7f))
+        write(value >> 14)
+        return 3
+    if value <= 0xfffffff:  # 4 bytes
+        write(0x80 | (value & 0x7f))
+        write(0x80 | ((value >> 7) & 0x7f))
+        write(0x80 | ((value >> 14) & 0x7f))
+        write(value >> 21)
+        return 4
+    if value <= 0x7ffffffff:  # 5 bytes
+        write(0x80 | (value & 0x7f))
+        write(0x80 | ((value >> 7) & 0x7f))
+        write(0x80 | ((value >> 14) & 0x7f))
+        write(0x80 | ((value >> 21) & 0x7f))
+        write(value >> 28)
+        return 5
+    else:
+        # Return to general algorithm
+        bits = value & 0x7f
+        value >>= 7
+        i = 0
+        while value:
+            write(0x80 | bits)
+            bits = value & 0x7f
+            value >>= 7
+            i += 1
+    write(bits)
+    return i
+
+
+def size_of_varint(value):
+    """ Number of bytes needed to encode an integer in variable-length format.
+    """
+    value = (value << 1) ^ (value >> 63)
+    if value <= 0x7f:
+        return 1
+    if value <= 0x3fff:
+        return 2
+    if value <= 0x1fffff:
+        return 3
+    if value <= 0xfffffff:
+        return 4
+    if value <= 0x7ffffffff:
+        return 5
+    if value <= 0x3ffffffffff:
+        return 6
+    if value <= 0x1ffffffffffff:
+        return 7
+    if value <= 0xffffffffffffff:
+        return 8
+    if value <= 0x7fffffffffffffff:
+        return 9
+    return 10
+
+
+def decode_varint(buffer, pos=0):
+    """ Decode an integer from a varint presentation. See
+    https://developers.google.com/protocol-buffers/docs/encoding?csw=1#varints
+    on how those can be produced.
+
+        Arguments:
+            buffer (bytearray): buffer to read from.
+            pos (int): optional position to read from
+
+        Returns:
+            (int, int): Decoded int value and next read position
+    """
+    result = buffer[pos]
+    if not (result & 0x81):
+        return (result >> 1), pos + 1
+    if not (result & 0x80):
+        return (result >> 1) ^ (~0), pos + 1
+
+    result &= 0x7f
+    pos += 1
+    shift = 7
+    while 1:
+        b = buffer[pos]
+        result |= ((b & 0x7f) << shift)
+        pos += 1
+        if not (b & 0x80):
+            return ((result >> 1) ^ -(result & 1), pos)
+        shift += 7
+        if shift >= 64:
+            raise ValueError("Out of int64 range")
+
+
+_crc32c = crc32c_py
+if crc32c_c is not None:
+    _crc32c = crc32c_c
+
+
+def calc_crc32c(memview, _crc32c=_crc32c):
+    """ Calculate CRC-32C (Castagnoli) checksum over a memoryview of data
+    """
+    return _crc32c(memview)
+
+
+def calc_crc32(memview):
+    """ Calculate simple CRC-32 checksum over a memoryview of data
+    """
+    crc = binascii.crc32(memview) & 0xffffffff
+    return crc

kafka/serializer/__init__.py

@@ -0,0 +1,4 @@
+from .abstract import Serializer, Deserializer
+from .default import DefaultSerializer
+from .json import JsonSerializer
+from .wrapper import DeserializeWrapper, SerializeWrapper

kafka/serializer/abstract.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+from typing import List, Tuple, Any
+
+
+class Serializer(ABC):
+    @abstractmethod
+    def serialize(self, topic: str, headers: List[Tuple[str, bytes]], data: Any):
+        pass
+
+    def close(self):
+        pass
+
+
+class Deserializer(ABC):
+    @abstractmethod
+    def deserialize(self, topic: str, headers: List[Tuple[str, bytes]], data: bytes):
+        pass
+
+    def close(self):
+        pass

kafka/serializer/default.py

@@ -0,0 +1,16 @@
+from .abstract import Serializer, Deserializer
+
+
+class DefaultSerializer(Serializer, Deserializer):
+    def __init__(self, encoding='utf-8'):
+        self.encoding = encoding
+
+    def serialize(self, topic, headers, data):
+        if type(data) in (bytes, bytearray, memoryview, type(None)):
+            return data
+        return data.encode(self.encoding)
+
+    def deserialize(self, topic, headers, data):
+        if data is None:
+            return None
+        return data.decode(self.encoding)

kafka/serializer/json.py

@@ -0,0 +1,17 @@
+import json
+
+from .abstract import Serializer, Deserializer
+
+
+class JsonSerializer(Serializer, Deserializer):
+    def serialize(self, topic, headers, data):
+        if data is None:
+            return None
+        s = json.dumps(data)
+        return super().serialize(topic, headers, s)
+
+    def deserialize(self, topic, headers, data):
+        if data is None:
+            return None
+        s = super().deserialize(topic, headers, data)
+        return json.loads(s)

kafka/serializer/wrapper.py

@@ -0,0 +1,21 @@
+from .abstract import Deserializer, Serializer
+
+
+class DeserializeWrapper(Deserializer):
+    def __init__(self, fn):
+        self.fn = fn
+
+    def deserialize(self, topic, headers, data):
+        if self.fn is None:
+            return data
+        return self.fn(data)
+
+
+class SerializeWrapper(Serializer):
+    def __init__(self, fn):
+        self.fn = fn
+
+    def serialize(self, topic, headers, data):
+        if self.fn is None:
+            return data
+        return self.fn(data)

kafka/structs.py

@@ -0,0 +1,69 @@
+""" Other useful structs """
+
+from collections import namedtuple
+
+
+TopicPartition = namedtuple("TopicPartition",
+    ["topic", "partition"])
+TopicPartition.__doc__ = """A topic and partition tuple
+
+Keyword Arguments:
+    topic (str): A topic name
+    partition (int): A partition id
+"""
+
+
+TopicPartitionReplica = namedtuple("TopicPartitionReplica",
+    ["topic", "partition", "broker_id"])
+TopicPartitionReplica.__doc__ = """A topic / partition / broker replica tuple
+
+Keyword Arguments:
+    topic (str): A topic name
+    partition (int): A partition id
+    broker_id (int): The node_id of the broker hosting the replica
+"""
+
+
+OffsetAndMetadata = namedtuple("OffsetAndMetadata",
+    ["offset", "metadata", "leader_epoch"], defaults=[None, '', -1])
+OffsetAndMetadata.__doc__ = """Container for committed group offset data.
+
+The Kafka offset commit API allows users to provide additional metadata
+(in the form of a string) when an offset is committed. This can be useful
+(for example) to store information about which node made the commit,
+what time the commit was made, etc.
+
+Keyword Arguments:
+    offset (int): The offset to be committed
+    metadata (str): Non-null metadata
+    leader_epoch (int): The last known epoch from the leader / broker
+"""
+
+
+OffsetAndTimestamp = namedtuple("OffsetAndTimestamp",
+    ["offset", "timestamp", "leader_epoch"])
+OffsetAndTimestamp.__doc__ = """An offset and timestamp tuple
+
+Keyword Arguments:
+    offset (int): An offset
+    timestamp (int): The timestamp associated to the offset
+    leader_epoch (int): The last known epoch from the leader / broker
+"""
+
+
+ConsumerGroupMetadata = namedtuple("ConsumerGroupMetadata",
+    ["group_id", "generation_id", "member_id", "group_instance_id"],
+    defaults=[-1, '', None])
+ConsumerGroupMetadata.__doc__ = """A snapshot of a consumer's group membership (KIP-447).
+
+Passed to KafkaProducer.send_offsets_to_transaction() so the broker can fence
+stale consumer instances when committing offsets inside a transaction. The
+broker uses member_id + generation_id + group_instance_id to verify the
+producer is acting on behalf of the current group generation.
+
+Keyword Arguments:
+    group_id (str): The consumer group id.
+    generation_id (int): The current generation id (-1 if unjoined).
+    member_id (str): The current member id ('' if unjoined).
+    group_instance_id (str): The static membership instance id, or None.
+"""