Add size_in_bytes to ABCRecordBatch and implement for Legacy and Default

Author: dpkp

kafka/record/abc.py

@@ -9,6 +9,11 @@
 class ABCRecord(object):
     __slots__ = ()
 
+    @abc.abstractproperty
+    def size_in_bytes(self):
+        """ Number of total bytes in record
+        """
+
     @abc.abstractproperty
     def offset(self):
         """ Absolute offset of record

kafka/record/default_records.py

@@ -275,10 +275,10 @@ def _read_msg(
 
         if self.is_control_batch:
             return ControlRecord(
-                offset, timestamp, self.timestamp_type, key, value, headers)
+                length, offset, timestamp, self.timestamp_type, key, value, headers)
         else:
             return DefaultRecord(
-                offset, timestamp, self.timestamp_type, key, value, headers)
+                length, offset, timestamp, self.timestamp_type, key, value, headers)
 
     def __iter__(self):
         self._maybe_uncompress()
@@ -314,17 +314,22 @@ def validate_crc(self):
 
 class DefaultRecord(ABCRecord):
 
-    __slots__ = ("_offset", "_timestamp", "_timestamp_type", "_key", "_value",
+    __slots__ = ("_size_in_bytes", "_offset", "_timestamp", "_timestamp_type", "_key", "_value",
                  "_headers")
 
-    def __init__(self, offset, timestamp, timestamp_type, key, value, headers):
+    def __init__(self, size_in_bytes, offset, timestamp, timestamp_type, key, value, headers):
+        self._size_in_bytes = size_in_bytes
         self._offset = offset
         self._timestamp = timestamp
         self._timestamp_type = timestamp_type
         self._key = key
         self._value = value
         self._headers = headers
 
+    @property
+    def size_in_bytes(self):
+        return self._size_in_bytes
+
     @property
     def offset(self):
         return self._offset
@@ -371,16 +376,16 @@ def __repr__(self):
 
 
 class ControlRecord(DefaultRecord):
-    __slots__ = ("_offset", "_timestamp", "_timestamp_type", "_key", "_value",
+    __slots__ = ("_size_in_bytes", "_offset", "_timestamp", "_timestamp_type", "_key", "_value",
                  "_headers", "_version", "_type")
 
     KEY_STRUCT = struct.Struct(
         ">h"  # Current Version => Int16
         "h"  # Type => Int16 (0 indicates an abort marker, 1 indicates a commit)
     )
 
-    def __init__(self, offset, timestamp, timestamp_type, key, value, headers):
-        super(ControlRecord, self).__init__(offset, timestamp, timestamp_type, key, value, headers)
+    def __init__(self, size_in_bytes, offset, timestamp, timestamp_type, key, value, headers):
+        super(ControlRecord, self).__init__(size_in_bytes, offset, timestamp, timestamp_type, key, value, headers)
         (self._version, self._type) = self.KEY_STRUCT.unpack(self._key)
 
     # see https://kafka.apache.org/documentation/#controlbatch

kafka/record/legacy_records.py

@@ -339,6 +339,10 @@ def headers(self):
     def checksum(self):
         return self._crc
 
+    @property
+    def size_in_bytes(self):
+        return LegacyRecordBatchBuilder.estimate_size_in_bytes(self._magic, None, self._key, self._value)
+
     def __repr__(self):
         return (
             "LegacyRecord(magic={!r} offset={!r}, timestamp={!r}, timestamp_type={!r},"