kafka-python 3.0.0__py3-none-any.whl

kafka/protocol/api_header.py

@@ -0,0 +1,55 @@
+from .api_data import JsonSchemaData
+from .data_container import DataContainer, SlotsBuilder
+from .schemas import BaseField, StructField, load_json
+
+
+class ApiHeader(DataContainer, metaclass=JsonSchemaData, init=False):
+    __slots__ = ()
+
+    def __init_subclass__(cls, **kw):
+        super().__init_subclass__(**kw)
+        if kw.get('init', True):
+            # pylint: disable=E1101
+            assert cls._json['type'] == 'header'
+            cls._flexible_versions = BaseField.parse_versions(cls._json['flexibleVersions'])
+            cls._valid_versions = BaseField.parse_versions(cls._json['validVersions'])
+
+    def encode(self, flexible=False):
+        # Request versions are 1-2, Response versions are 0-1
+        version = self._flexible_versions[0] if flexible else self._valid_versions[0] # pylint: disable=E1136
+        # compact=False is probably wrong,
+        # but it works to make sure that the client_id request header field
+        # is never encoded as compact (required to support ApiVersionsRequest for unsupported version)
+        return super().encode(version=version, compact=False, tagged=flexible)
+
+    def encode_into(self, out, flexible=False):
+        version = self._flexible_versions[0] if flexible else self._valid_versions[0] # pylint: disable=E1136
+        super().encode_into(out, version=version, compact=False, tagged=flexible)
+
+    @classmethod
+    def decode(cls, data, flexible=False):
+        # Request versions are 1-2, Response versions are 0-1
+        version = cls._flexible_versions[0] if flexible else cls._valid_versions[0] # pylint: disable=E1136
+        return cls._struct.decode(data, version=version, compact=False, tagged=flexible, data_class=cls)
+
+
+class ResponseClassRegistry:
+    _response_class_registry = {}
+
+    @classmethod
+    def register_response_class(cls, response_class):
+        cls._response_class_registry[response_class.API_KEY] = response_class
+
+    @classmethod
+    def get_response_class(cls, request_header):
+        response_class = cls._response_class_registry.get(request_header.request_api_key)
+        if response_class is not None:
+            return response_class[request_header.request_api_version]
+
+
+class RequestHeader(ApiHeader):
+    def get_response_class(self):
+        return ResponseClassRegistry.get_response_class(self)
+
+
+class ResponseHeader(ApiHeader): pass

kafka/protocol/api_key.py

@@ -0,0 +1,97 @@
+from enum import IntEnum
+
+
+class ApiKey(IntEnum):
+    Produce = 0
+    Fetch = 1
+    ListOffsets = 2
+    Metadata = 3
+    LeaderAndIsr = 4
+    StopReplica = 5
+    UpdateMetadata = 6
+    ControlledShutdown = 7
+    OffsetCommit = 8
+    OffsetFetch = 9
+    FindCoordinator = 10
+    JoinGroup = 11
+    Heartbeat = 12
+    LeaveGroup = 13
+    SyncGroup = 14
+    DescribeGroups = 15
+    ListGroups = 16
+    SaslHandshake = 17
+    ApiVersions = 18
+    CreateTopics = 19
+    DeleteTopics = 20
+    DeleteRecords = 21
+    InitProducerId = 22
+    OffsetForLeaderEpoch = 23
+    AddPartitionsToTxn = 24
+    AddOffsetsToTxn = 25
+    EndTxn = 26
+    WriteTxnMarkers = 27
+    TxnOffsetCommit = 28
+    DescribeAcls = 29
+    CreateAcls = 30
+    DeleteAcls = 31
+    DescribeConfigs = 32
+    AlterConfigs = 33
+    AlterReplicaLogDirs = 34
+    DescribeLogDirs = 35
+    SaslAuthenticate = 36
+    CreatePartitions = 37
+    CreateDelegationToken = 38
+    RenewDelegationToken = 39
+    ExpireDelegationToken = 40
+    DescribeDelegationToken = 41
+    DeleteGroups = 42
+    ElectLeaders = 43
+    IncrementalAlterConfigs = 44
+    AlterPartitionReassignments = 45
+    ListPartitionReassignments = 46
+    OffsetDelete = 47
+    DescribeClientQuotas = 48
+    AlterClientQuotas = 49
+    DescribeUserScramCredentials = 50
+    AlterUserScramCredentials = 51
+    Vote = 52
+    BeginQuorumEpoch = 53
+    EndQuorumEpoch = 54
+    DescribeQuorum = 55
+    AlterPartition = 56
+    UpdateFeatures = 57
+    Envelope = 58
+    FetchSnapshot = 59
+    DescribeCluster = 60
+    DescribeProducers = 61
+    BrokerRegistration = 62
+    BrokerHeartbeat = 63
+    UnregisterBroker = 64
+    DescribeTransactions = 65
+    ListTransactions = 66
+    AllocateProducerIds = 67
+    ConsumerGroupHeartbeat = 68
+    ConsumerGroupDescribe = 69
+    ControllerRegistration = 70
+    GetTelemetrySubscriptions = 71
+    PushTelemetry = 72
+    AssignReplicasToDirs = 73
+    ListConfigResources = 74
+    DescribeTopicPartitions = 75
+    ShareGroupHeartbeat = 76
+    ShareGroupDescribe = 77
+    ShareFetch = 78
+    ShareAcknowledge = 79
+    AddRaftVoter = 80
+    RemoveRaftVoter = 81
+    UpdateRaftVoter = 82
+    InitializeShareGroupState = 83
+    ReadShareGroupState = 84
+    WriteShareGroupState = 85
+    DeleteShareGroupState = 86
+    ReadShareGroupStateSummary = 87
+    StreamsGroupHeartbeat = 88
+    StreamsGroupDescribe = 89
+    DescribeShareGroupOffsets = 90
+    AlterShareGroupOffsets = 91
+    DeleteShareGroupOffsets = 92

kafka/protocol/api_message.py

@@ -0,0 +1,277 @@
+import io
+import weakref
+
+from .api_data import JsonSchemaData
+from .api_header import RequestHeader, ResponseHeader, ResponseClassRegistry
+from .data_container import DataContainer
+from .schemas import BaseField, StructField, load_json
+from .schemas.fields.codecs import Int32, EncodeBufferPool
+from struct import pack_into
+
+from kafka.util import classproperty
+
+
+class VersionSubscriptable(type):
+    def __init__(cls, name, bases, attrs, **kw):
+        super().__init__(name, bases, attrs, **kw)
+        if kw.get('init', True):
+            # The primary message class has _version = None
+            # and a _VERSIONS dict that provides access to version-specific wrappers
+            # We also include cls[None] -> primary class to "exit" a version class
+            if getattr(cls, '_class_version', None) is None:
+                cls._class_version = None
+                cls._VERSIONS = {}
+
+    def __getitem__(cls, version):
+        # Use [] lookups to move from primary class to "versioned" classes
+        # which are simple wrappers around the primary class but with a _version attr
+        if cls._class_version is not None:
+            primary_cls = cls.mro()[1]
+            if version is None:
+                return primary_cls
+            return primary_cls[version]
+        elif version is None:
+            return cls
+        if cls._valid_versions is not None:
+            if version < 0:
+                version += 1 + cls.max_version # support negative index, e.g., [-1]
+            if not cls.min_version <= version <= cls.max_version:
+                raise ValueError('Invalid version! min=%d, max=%d' % (cls.min_version, cls.max_version))
+        if version in cls._VERSIONS:
+            return cls._VERSIONS[version]
+        klass_name = cls.__name__ + '_v' + str(version)
+        cls._VERSIONS[version] = type(klass_name, tuple(cls.mro()), {'_class_version': version}, init=False)
+        return cls._VERSIONS[version]
+
+    def __len__(cls):
+        # Maintain compatibility
+        if cls._valid_versions is None:
+            raise RuntimeError('Unable to calculate __len__ for class without valid_versions')
+        elif cls._class_version is not None:
+            raise TypeError('len() only supported on primary message class (not versioned)')
+        return cls._valid_versions[1] + 1
+
+
+class ApiMessageData(VersionSubscriptable, JsonSchemaData):
+    def __init__(cls, name, bases, attrs, **kw):
+        super().__init__(name, bases, attrs, **kw)
+        if kw.get('init', True):
+            # Ignore min valid version on request/response schemas
+            # We'll get the brokers supported versions via ApiVersionsRequest
+            if cls._struct._versions[0] > 0:
+                cls._struct._versions = (0, cls._struct._versions[1])
+
+
+class ApiMessage(DataContainer, metaclass=ApiMessageData, init=False):
+    __slots__ = ('_header', '_min_version', '_max_version')
+
+    def __init_subclass__(cls, **kw):
+        super().__init_subclass__(**kw)
+        if kw.get('init', True):
+            # pylint: disable=E1101
+            assert cls._json is not None
+            assert cls._json['type'] in ('request', 'response')
+            cls._flexible_versions = BaseField.parse_versions(cls._json['flexibleVersions'])
+            cls._valid_versions = BaseField.parse_versions(cls._json['validVersions'])
+            if not cls.is_request():
+                ResponseClassRegistry.register_response_class(weakref.proxy(cls))
+
+    def __new__(cls, *args, **kwargs):
+        # Translate "versioned" classes back to primary w/ version= kwarg on construction
+        if cls._class_version is not None:
+            if kwargs.get('version', cls._class_version) != cls._class_version: # pylint: disable=E1101
+                raise ValueError("Version has already been set by class")
+            kwargs['version'] = cls._class_version
+            instance = super().__new__(cls[None])
+            instance.__init__(*args, **kwargs)
+            return instance
+        return super().__new__(cls)
+
+    def __init__(self, *args, **kwargs):
+        self._header = None
+        self._min_version = kwargs.pop('min_version', None)
+        self._max_version = kwargs.pop('max_version', None)
+        super().__init__(*args, **kwargs)
+
+    @classproperty
+    def name(cls): # pylint: disable=E0213
+        return cls._json['name'] # pylint: disable=E1101
+
+    @classproperty
+    def type(cls): # pylint: disable=E0213
+        return cls._json['type'] # pylint: disable=E1101
+
+    @classproperty
+    def API_KEY(cls): # pylint: disable=E0213
+        return cls._json['apiKey'] # pylint: disable=E1101
+
+    @classproperty
+    def json(cls): # pylint: disable=E0213
+        return cls._json # pylint: disable=E1101
+
+    @classproperty
+    def valid_versions(cls): # pylint: disable=E0213
+        return cls._valid_versions
+
+    @classproperty
+    def min_version(cls): # pylint: disable=E0213
+        return 0
+
+    @classproperty
+    def max_version(cls): # pylint: disable=E0213
+        if cls._valid_versions is not None:
+            return cls._valid_versions[1] # pylint: disable=E1136
+        return None
+
+    @classmethod
+    def flexible_version_q(cls, version):
+        if cls._flexible_versions is not None:
+            if cls._flexible_versions[0] <= version <= cls._flexible_versions[1]: # pylint: disable=E1136
+                return True
+        return False
+
+    @classmethod
+    def is_request(cls):
+        return cls.type == 'request'
+
+    # allow override by api-specific classes (e.g., ProduceRequest)
+    def expect_response(self):
+        return True
+
+    @property
+    def API_VERSION(self):
+        return self._version if self._version is not None else self._class_version # pylint: disable=E1101
+
+    @API_VERSION.setter
+    def API_VERSION(self, version):
+        if not 0 <= version <= self.max_version:
+            raise ValueError('Invalid version %s (max version is %s).' % (version, self.max_version))
+        self._version = version
+        if self._header is not None:
+            self._header.request_api_version = version
+
+    @property
+    def header(self):
+        return self._header
+
+    @classproperty
+    def header_class(cls): # pylint: disable=E0213
+        if cls.type == 'response':
+            return ResponseHeader
+        elif cls.type == 'request':
+            return RequestHeader
+        elif cls.type is None:
+            return None
+        else:
+            raise ValueError('Expected request or response type: %s' % cls.type)
+
+    def with_header(self, correlation_id=0, client_id='kafka-python'):
+        if self.is_request():
+            kwargs = {
+                'request_api_key': self.API_KEY,
+                'request_api_version': self.API_VERSION,
+                'correlation_id': correlation_id,
+                'client_id': client_id,
+            }
+        else:
+            kwargs = {
+                'correlation_id': correlation_id,
+            }
+        self._header = self.header_class(**kwargs)
+
+    # allow override by api-specific classes (e.g., ApiVersionsResponse)
+    def encode_header(self, flexible=False):
+        return self._header.encode(flexible=flexible) # pylint: disable=E1120
+
+    def encode_header_into(self, out, flexible=False):
+        # Subclasses may override encode_header to change flexible flag
+        # (e.g., ApiVersionsResponse forces flexible=False).
+        # Fall back to bytes-based encode_header if not overridden at this level.
+        header_bytes = self.encode_header(flexible=flexible)
+        n = len(header_bytes)
+        out.buf[out.pos:out.pos+n] = header_bytes
+        out.pos += n
+
+    @classmethod
+    def parse_header(cls, data, version=None):
+        version = cls._class_version if version is None else version
+        if version is None:
+            raise ValueError('Version required to decode data')
+        elif not 0 <= version <= cls.max_version:
+            raise ValueError('Invalid version %s (max version is %s).' % (version, cls.max_version))
+        flexible = cls.flexible_version_q(version)
+        return cls.header_class.decode(data, flexible=flexible) # pylint: disable=E1101
+
+    @classmethod
+    def encode_into__optimized_context(cls, version=None):
+        if version is None:
+            version = cls._class_version
+        if version is None:
+            raise ValueError('Version required to encode data')
+        flexible = cls.flexible_version_q(version)
+        return cls._struct.encode_into__optimized_context(version, compact=flexible, tagged=flexible)
+
+    def encode(self, version=None, header=False, framed=False):
+        if version is not None:
+            self.API_VERSION = version
+        if self.API_VERSION is None:
+            raise ValueError('Version required to encode data')
+        if header and self._header is None:
+            raise ValueError('No header found')
+
+        flexible = self.flexible_version_q(self.API_VERSION)
+        with EncodeBufferPool.acquire() as out:
+            if framed:
+                out.pos += 4  # reserve space for frame size
+            if header:
+                self.encode_header_into(out, flexible=flexible)
+            fast_encode = self._struct.compiled_encode_into(self.API_VERSION, compact=flexible, tagged=flexible)
+            fast_encode(self, out)
+            if framed:
+                payload_size = out.pos - 4
+                pack_into('>i', out.buf, 0, payload_size)
+            return out.result()
+
+    @classmethod
+    def decode(cls, data, version=None, header=False, framed=False):
+        version = cls._class_version if version is None else version
+        if version is None:
+            raise ValueError('Version required to decode data')
+        elif not 0 <= version <= cls.max_version:
+            raise ValueError('Invalid version %s (max version is %s).' % (version, cls.max_version))
+
+        # Normalize input to memoryview
+        if isinstance(data, bytes):
+            data = memoryview(data)
+        elif not isinstance(data, memoryview):
+            if hasattr(data, 'read'):
+                data = memoryview(data.read())
+            else:
+                # else: assume buffer protocol compatible
+                data = memoryview(data)
+        pos = 0
+
+        if framed:
+            size, pos = Int32.decode_from(data, pos)
+        if header:
+            # Header decode uses BytesIO (variable-length fields)
+            hdr_data = io.BytesIO(bytes(data[pos:]))
+            hdr = cls.parse_header(hdr_data, version=version)
+            pos += hdr_data.tell()
+        else:
+            hdr = None
+
+        # Resolve versioned class to primary class (e.g., Response_v0 -> Response)
+        data_class = cls[None]
+
+        flexible = cls.flexible_version_q(version)
+        fast_decode = cls._struct.compiled_decode_from(version, compact=flexible, tagged=flexible, data_class=data_class)
+        ret, pos = fast_decode(data, pos)
+        if hdr is not None:
+            ret._header = hdr
+        return ret
+
+    def __eq__(self, other):
+        if self.version != other.version:
+            return False
+        return super().__eq__(other)