kafka-python 3.0.0__py3-none-any.whl

kafka/producer/producer_batch.py

@@ -0,0 +1,192 @@
+from enum import IntEnum
+import logging
+import time
+
+import kafka.errors as Errors
+from kafka.producer.future import FutureRecordMetadata, FutureProduceResult
+
+
+log = logging.getLogger(__name__)
+
+
+class FinalState(IntEnum):
+    ABORTED = 0
+    FAILED = 1
+    SUCCEEDED = 2
+
+
+class ProducerBatch:
+    def __init__(self, tp, records, now=None):
+        now = time.monotonic() if now is None else now
+        self.max_record_size = 0
+        self.created = now
+        self.drained = None
+        self.attempts = 0
+        self.last_attempt = now
+        self.last_append = now
+        self.records = records
+        self.topic_partition = tp
+        self.produce_future = FutureProduceResult(tp)
+        self._record_futures = []
+        self._retry = False
+        self._final_state = None
+
+    @property
+    def final_state(self):
+        return self._final_state
+
+    @property
+    def record_count(self):
+        return self.records.next_offset()
+
+    @property
+    def producer_id(self):
+        return self.records.producer_id if self.records else None
+
+    @property
+    def producer_epoch(self):
+        return self.records.producer_epoch if self.records else None
+
+    @property
+    def base_sequence(self):
+        return self.records.base_sequence if self.records else None
+
+    @property
+    def has_sequence(self):
+        base_seq = self.base_sequence
+        return base_seq is not None and base_seq != -1
+
+    def try_append(self, timestamp_ms, key, value, headers, now=None):
+        metadata = self.records.append(timestamp_ms, key, value, headers)
+        if metadata is None:
+            return None
+
+        now = time.monotonic() if now is None else now
+        self.max_record_size = max(self.max_record_size, metadata.size)
+        self.last_append = now
+        future = FutureRecordMetadata(
+            self.produce_future,
+            metadata.offset,
+            metadata.timestamp,
+            metadata.crc,
+            len(key) if key is not None else -1,
+            len(value) if value is not None else -1,
+            sum(len(h_key.encode("utf-8")) + len(h_val) for h_key, h_val in headers) if headers else -1)
+        self._record_futures.append(future)
+        return future
+
+    def abort(self, exception):
+        """Abort the batch and complete the future and callbacks."""
+        if self._final_state is not None:
+            raise Errors.IllegalStateError("Batch has already been completed in final state: %s" % self._final_state)
+        self._final_state = FinalState.ABORTED
+
+        log.debug("Aborting batch for partition %s: %s", self.topic_partition, exception)
+        self._complete_future(-1, -1, lambda _: exception)
+
+    def complete(self, base_offset, log_append_time):
+        """Complete the batch successfully.
+
+        Arguments:
+            base_offset (int): The base offset of the messages assigned by the server
+            log_append_time (int): The log append time or -1 if CreateTime is being used
+
+        Returns: True if the batch was completed as a result of this call, and False
+            if it had been completed previously.
+        """
+        return self.done(base_offset=base_offset, timestamp_ms=log_append_time)
+
+    def complete_exceptionally(self, top_level_exception, record_exceptions_fn):
+        """
+        Complete the batch exceptionally. The provided top-level exception will be used
+        for each record future contained in the batch.
+
+        Arguments:
+            top_level_exception (Exception): top-level partition error.
+            record_exceptions_fn (callable int -> Exception): Record exception function mapping
+                batch_index to the respective record exception.
+
+        Raises:
+            TypeError: if top_level_exception is not Exception,
+                or record_exceptions_fn is not callable.
+
+        Returns:
+            True if the batch was completed as a result of this call,
+            or False if it had been completed previously.
+        """
+        if not isinstance(top_level_exception, Exception):
+            raise TypeError('top_level_exception must be type Exception')
+        if not callable(record_exceptions_fn):
+            raise TypeError('record_exceptions_fn must be callable')
+        return self.done(top_level_exception=top_level_exception, record_exceptions_fn=record_exceptions_fn)
+
+    def done(self, base_offset=None, timestamp_ms=None, top_level_exception=None, record_exceptions_fn=None):
+        """
+        Finalize the state of a batch. Final state, once set, is immutable. This function may be called
+        once or twice on a batch. It may be called twice if
+            1. An inflight batch expires before a response from the broker is received. The batch's final
+            state is set to FAILED. But it could succeed on the broker and second time around batch.done() may
+            try to set SUCCEEDED final state.
+
+            2. If a transaction abortion happens or if the producer is closed forcefully, the final state is
+            ABORTED but again it could succeed if broker responds with a success.
+
+        Attempted transitions from [FAILED | ABORTED] --> SUCCEEDED are logged.
+        Attempted transitions from one failure state to the same or a different failed state are ignored.
+        Attempted transitions from SUCCEEDED to the same or a failed state throw an exception.
+        """
+        final_state = FinalState.SUCCEEDED if top_level_exception is None else FinalState.FAILED
+        if self._final_state is None:
+            self._final_state = final_state
+            if final_state is FinalState.SUCCEEDED:
+                log.debug("Successfully produced messages to %s with base offset %s", self.topic_partition, base_offset)
+            else:
+                log.warning("Failed to produce messages to topic-partition %s with base offset %s: %s",
+                            self.topic_partition, base_offset, top_level_exception)
+            self._complete_future(base_offset, timestamp_ms, record_exceptions_fn)
+            return True
+
+        elif self._final_state is not FinalState.SUCCEEDED:
+            if final_state is FinalState.SUCCEEDED:
+                # Log if a previously unsuccessful batch succeeded later on.
+                log.debug("ProduceResponse returned %s for %s after batch with base offset %s had already been %s.",
+                          final_state, self.topic_partition, base_offset, self._final_state)
+            else:
+                # FAILED --> FAILED and ABORTED --> FAILED transitions are ignored.
+                log.debug("Ignored state transition %s -> %s for %s batch with base offset %s",
+                          self._final_state, final_state, self.topic_partition, base_offset)
+        else:
+            # A SUCCESSFUL batch must not attempt another state change.
+            raise Errors.IllegalStateError("A %s batch must not attempt another state change to %s" % (self._final_state, final_state))
+        return False
+
+    def _complete_future(self, base_offset, timestamp_ms, record_exceptions_fn):
+        if self.produce_future.is_done:
+            raise Errors.IllegalStateError('Batch is already closed!')
+        self.produce_future.success((base_offset, timestamp_ms, record_exceptions_fn))
+
+    def has_reached_delivery_timeout(self, delivery_timeout_ms, now=None):
+        now = time.monotonic() if now is None else now
+        return delivery_timeout_ms / 1000 <= now - self.created
+
+    def in_retry(self):
+        return self._retry
+
+    def retry(self, now=None):
+        now = time.monotonic() if now is None else now
+        self._retry = True
+        self.attempts += 1
+        self.last_attempt = now
+        self.last_append = now
+
+    @property
+    def is_done(self):
+        return self.produce_future.is_done
+
+    def __str__(self):
+        return 'ProducerBatch(topic_partition=%s, record_count=%d)' % (
+            self.topic_partition, self.records.next_offset())
+
+    # for heapq
+    def __lt__(self, other):
+        return self.created < other.created