def test_written_bytes_equals_size_in_bytes(magic): key = b"test" value = b"Super" builder = LegacyRecordBatchBuilder( magic=magic, compression_type=0, batch_size=1024 * 1024) size_in_bytes = builder.size_in_bytes( 0, timestamp=9999999, key=key, value=value) pos = builder.size() builder.append(0, timestamp=9999999, key=key, value=value) assert builder.size() - pos == size_in_bytes
def test_written_bytes_equals_size_in_bytes(magic): key = b"test" value = b"Super" builder = LegacyRecordBatchBuilder( magic=magic, compression_type=0, batch_size=1024 * 1024) size_in_bytes = builder.size_in_bytes( 0, timestamp=9999999, key=key, value=value) pos = builder.size() builder.append(0, timestamp=9999999, key=key, value=value) assert builder.size() - pos == size_in_bytes
class BatchBuilder: def __init__(self, magic, batch_size, compression_type): self._builder = LegacyRecordBatchBuilder(magic, compression_type, batch_size) self._relative_offset = 0 self._buffer = None self._closed = False def append(self, *, timestamp, key, value): """Add a message to the batch. Arguments: timestamp (float or None): epoch timestamp in seconds. If None, the timestamp will be set to the current time. If submitting to an 0.8.x or 0.9.x broker, the timestamp will be ignored. key (bytes or None): the message key. `key` and `value` may not both be None. value (bytes or None): the message value. `key` and `value` may not both be None. Returns: If the message was successfully added, returns a metadata object with crc, offset, size, and timestamp fields. If the batch is full or closed, returns None. """ if self._closed: return None metadata = self._builder.append(self._relative_offset, timestamp, key, value) # Check if we could add the message if metadata is None: return None self._relative_offset += 1 return metadata def close(self): """Close the batch to further updates. Closing the batch before submitting to the producer ensures that no messages are added via the ``producer.send()`` interface. To gracefully support both the batch and individual message interfaces, leave the batch open. For complete control over the batch's contents, close before submission. Closing a batch has no effect on when it's sent to the broker. A batch may not be reopened after it's closed. """ if self._closed: return self._closed = True data = self._builder.build() self._buffer = io.BytesIO(Int32.encode(len(data)) + data) del self._builder def _build(self): if not self._closed: self.close() return self._buffer def size(self): """Get the size of batch in bytes.""" if self._buffer: return self._buffer.getbuffer().nbytes else: return self._builder.size() def record_count(self): """Get the number of records in the batch.""" return self._relative_offset