message.py

from __future__ import absolute_import

import io
import time

from kafka.codec import (has_gzip, has_snappy, has_lz4, has_zstd,
                     gzip_decode, snappy_decode, zstd_decode,
                     lz4_decode, lz4_decode_old_kafka)
from kafka.protocol.frame import KafkaBytes
from kafka.protocol.struct import Struct
from kafka.protocol.types import (
    Int8, Int32, Int64, Bytes, Schema, AbstractType
)
from kafka.util import crc32, WeakMethod


class Message(Struct):
    SCHEMAS = [
        Schema(
            ('crc', Int32),
            ('magic', Int8),
            ('attributes', Int8),
            ('key', Bytes),
            ('value', Bytes)),
        Schema(
            ('crc', Int32),
            ('magic', Int8),
            ('attributes', Int8),
            ('timestamp', Int64),
            ('key', Bytes),
            ('value', Bytes)),
    ]
    SCHEMA = SCHEMAS[1]
    CODEC_MASK = 0x07
    CODEC_GZIP = 0x01
    CODEC_SNAPPY = 0x02
    CODEC_LZ4 = 0x03
    CODEC_ZSTD = 0x04
    TIMESTAMP_TYPE_MASK = 0x08
    HEADER_SIZE = 22  # crc(4), magic(1), attributes(1), timestamp(8), key+value size(4*2)

    def __init__(self, value, key=None, magic=0, attributes=0, crc=0,
                 timestamp=None):
        assert value is None or isinstance(value, bytes), 'value must be bytes'
        assert key is None or isinstance(key, bytes), 'key must be bytes'
        assert magic > 0 or timestamp is None, 'timestamp not supported in v0'

        # Default timestamp to now for v1 messages
        if magic > 0 and timestamp is None:
            timestamp = int(time.time() * 1000)
        self.timestamp = timestamp
        self.crc = crc
        self._validated_crc = None
        self.magic = magic
        self.attributes = attributes
        self.key = key
        self.value = value
        self.encode = WeakMethod(self._encode_self)

    @property
    def timestamp_type(self):
        """0 for CreateTime; 1 for LogAppendTime; None if unsupported.

        Value is determined by broker; produced messages should always set to 0
        Requires Kafka >= 0.10 / message version >= 1
        """
        if self.magic == 0:
            return None
        elif self.attributes & self.TIMESTAMP_TYPE_MASK:
            return 1
        else:
            return 0

    def _encode_self(self, recalc_crc=True):
        version = self.magic
        if version == 1:
            fields = (self.crc, self.magic, self.attributes, self.timestamp, self.key, self.value)
        elif version == 0:
            fields = (self.crc, self.magic, self.attributes, self.key, self.value)
        else:
            raise ValueError('Unrecognized message version: %s' % (version,))
        message = Message.SCHEMAS[version].encode(fields)
        if not recalc_crc:
            return message
        self.crc = crc32(message[4:])
        crc_field = self.SCHEMAS[version].fields[0]
        return crc_field.encode(self.crc) + message[4:]

    @classmethod
    def decode(cls, data):
        _validated_crc = None
        if isinstance(data, bytes):
            _validated_crc = crc32(data[4:])
            data = io.BytesIO(data)
        # Partial decode required to determine message version
        base_fields = cls.SCHEMAS[0].fields[0:3]
        crc, magic, attributes = [field.decode(data) for field in base_fields]
        remaining = cls.SCHEMAS[magic].fields[3:]
        fields = [field.decode(data) for field in remaining]
        if magic == 1:
            timestamp = fields[0]
        else:
            timestamp = None
        msg = cls(fields[-1], key=fields[-2],
                  magic=magic, attributes=attributes, crc=crc,
                  timestamp=timestamp)
        msg._validated_crc = _validated_crc
        return msg

    def validate_crc(self):
        if self._validated_crc is None:
            raw_msg = self._encode_self(recalc_crc=False)
            self._validated_crc = crc32(raw_msg[4:])
        if self.crc == self._validated_crc:
            return True
        return False

    def is_compressed(self):
        return self.attributes & self.CODEC_MASK != 0

    def decompress(self):
        codec = self.attributes & self.CODEC_MASK
        assert codec in (self.CODEC_GZIP, self.CODEC_SNAPPY, self.CODEC_LZ4, self.CODEC_ZSTD)
        if codec == self.CODEC_GZIP:
            assert has_gzip(), 'Gzip decompression unsupported'
            raw_bytes = gzip_decode(self.value)
        elif codec == self.CODEC_SNAPPY:
            assert has_snappy(), 'Snappy decompression unsupported'
            raw_bytes = snappy_decode(self.value)
        elif codec == self.CODEC_LZ4:
            assert has_lz4(), 'LZ4 decompression unsupported'
            if self.magic == 0:
                raw_bytes = lz4_decode_old_kafka(self.value)
            else:
                raw_bytes = lz4_decode(self.value)
        elif codec == self.CODEC_ZSTD:
            assert has_zstd(), "ZSTD decompression unsupported"
            raw_bytes = zstd_decode(self.value)
        else:
            raise Exception('This should be impossible')

        return MessageSet.decode(raw_bytes, bytes_to_read=len(raw_bytes))

    def __hash__(self):
        return hash(self._encode_self(recalc_crc=False))


class PartialMessage(bytes):
    def __repr__(self):
        return 'PartialMessage(%s)' % (self,)


class MessageSet(AbstractType):
    ITEM = Schema(
        ('offset', Int64),
        ('message', Bytes)
    )
    HEADER_SIZE = 12  # offset + message_size

    @classmethod
    def encode(cls, items, prepend_size=True):
        # RecordAccumulator encodes messagesets internally
        if isinstance(items, (io.BytesIO, KafkaBytes)):
            size = Int32.decode(items)
            if prepend_size:
                # rewind and return all the bytes
                items.seek(items.tell() - 4)
                size += 4
            return items.read(size)

        encoded_values = []
        for (offset, message) in items:
            encoded_values.append(Int64.encode(offset))
            encoded_values.append(Bytes.encode(message))
        encoded = b''.join(encoded_values)
        if prepend_size:
            return Bytes.encode(encoded)
        else:
            return encoded

    @classmethod
    def decode(cls, data, bytes_to_read=None):
        """Compressed messages should pass in bytes_to_read (via message size)
        otherwise, we decode from data as Int32
        """
        if isinstance(data, bytes):
            data = io.BytesIO(data)
        if bytes_to_read is None:
            bytes_to_read = Int32.decode(data)

        # if FetchRequest max_bytes is smaller than the available message set
        # the server returns partial data for the final message
        # So create an internal buffer to avoid over-reading
        raw = io.BytesIO(data.read(bytes_to_read))

        items = []
        while bytes_to_read:
            try:
                offset = Int64.decode(raw)
                msg_bytes = Bytes.decode(raw)
                bytes_to_read -= 8 + 4 + len(msg_bytes)
                items.append((offset, len(msg_bytes), Message.decode(msg_bytes)))
            except ValueError:
                # PartialMessage to signal that max_bytes may be too small
                items.append((None, None, PartialMessage()))
                break
        return items

    @classmethod
    def repr(cls, messages):
        if isinstance(messages, (KafkaBytes, io.BytesIO)):
            offset = messages.tell()
            decoded = cls.decode(messages)
            messages.seek(offset)
            messages = decoded
        return str([cls.ITEM.repr(m) for m in messages])