diff --git a/plaso/storage/sqlite/sqlite_file.py b/plaso/storage/sqlite/sqlite_file.py index 45e65a6cc0..dfecd71778 100644 --- a/plaso/storage/sqlite/sqlite_file.py +++ b/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 diff --git a/tests/storage/sqlite/sqlite_file.py b/tests/storage/sqlite/sqlite_file.py index 7565137bcb..14f38c0c0e 100644 --- a/tests/storage/sqlite/sqlite_file.py +++ b/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()