Correction to storage profiler write sample data size #5191

plaso/storage/sqlite/sqlite_file.py

@@ -302,9 +302,10 @@ def _WriteNewAttributeContainer(self, container):
 
             if self.compression_format == definitions.COMPRESSION_FORMAT_ZLIB:
                 compressed_data = zlib.compress(serialized_data)
-                serialized_data = sqlite3.Binary(compressed_data)
+                column_value = sqlite3.Binary(compressed_data)
             else:
-                compressed_data = ""
+                compressed_data = b""
+                column_value = serialized_data
 
             if self._storage_profiler:
                 self._storage_profiler.Sample(
@@ -316,7 +317,7 @@ def _WriteNewAttributeContainer(self, container):
                 )
 
             column_names = ["_data"]
-            values = [serialized_data]
+            values = [column_value]
 
             self._CacheAttributeContainerForWrite(
                 container.CONTAINER_TYPE, column_names, values

tests/storage/sqlite/sqlite_file.py

@@ -1,10 +1,14 @@
 #!/usr/bin/env python3
 """Tests for the SQLite-based storage."""
 
+import csv
+import gzip
 import os
 import unittest
 
 from plaso.containers import events
+from plaso.engine import configurations
+from plaso.engine import profilers
 from plaso.lib import definitions
 from plaso.storage.sqlite import sqlite_file
 
@@ -283,6 +287,43 @@ def testWriteNewAttributeContainer(self):
             finally:
                 test_store.Close()
 
+    def testWriteNewAttributeContainerWithStorageProfiler(self):
+        """Tests the _WriteNewAttributeContainer function with a storage profiler."""
+        event_data = events.EventData()
+        event_data.data_type = "test:event_data"
+        event_data.body = "test" * 256
+
+        profiling_configuration = configurations.ProfilingConfiguration()
+
+        with shared_test_lib.TempDirectory() as temp_directory:
+            profiling_configuration.directory = temp_directory
+
+            test_profiler = profilers.StorageProfiler("test", profiling_configuration)
+            test_profiler.Start()
+
+            test_path = os.path.join(temp_directory, "plaso.sqlite")
+            test_store = sqlite_file.SQLiteStorageFile()
+            test_store.SetStorageProfiler(test_profiler)
+            test_store.Open(path=test_path, read_only=False)
+
+            try:
+                test_store._WriteNewAttributeContainer(event_data)
+            finally:
+                test_store.Close()
+                test_profiler.Stop()
+
+            sample_path = os.path.join(temp_directory, "storage-test.csv.gz")
+            with gzip.open(sample_path, "rt", encoding="utf-8") as file_object:
+                rows = list(csv.reader(file_object, delimiter="\t"))
+
+        write_rows = [row for row in rows if row[1] == "write_new"]
+        self.assertEqual(len(write_rows), 1)
+
+        data_size = int(write_rows[0][5])
+        compressed_data_size = int(write_rows[0][6])
+
+        self.assertGreater(data_size, compressed_data_size)
+
     def testAddAttributeContainer(self):
         """Tests the AddAttributeContainer function."""
         event_data_stream = events.EventDataStream()