Refactored Redis store

log2timeline · Apr 8, 2024 · 42b2745 · 42b2745
1 parent e5fe4dd
commit 42b2745
Showing 1 changed file with 34 additions and 19 deletions.
diff --git a/plaso/storage/redis/redis_store.py b/plaso/storage/redis/redis_store.py
@@ -112,7 +112,18 @@ def _WriteExistingAttributeContainer(self, container):
     redis_key = identifier.CopyToString()
 
     json_dict = self._json_serializer.ConvertAttributeContainerToJSON(container)
-    json_string = json.dumps(json_dict)
+
+    try:
+      json_string = json.dumps(json_dict)
+    except TypeError as exception:
+      raise IOError((
+          f'Unable to serialize attribute container: '
+          f'{container.CONTAINER_TYPE:s} with error: {exception!s}.'))
+
+    if not json_string:
+      raise IOError((
+          f'Unable to serialize attribute container: '
+          f'{container.CONTAINER_TYPE:s}'))
 
     self._redis_client.hset(redis_hash_name, key=redis_key, value=json_string)
 
@@ -133,7 +144,18 @@ def _WriteNewAttributeContainer(self, container):
     redis_key = identifier.CopyToString()
 
     json_dict = self._json_serializer.ConvertAttributeContainerToJSON(container)
-    json_string = json.dumps(json_dict)
+
+    try:
+      json_string = json.dumps(json_dict)
+    except TypeError as exception:
+      raise IOError((
+          f'Unable to serialize attribute container: '
+          f'{container.CONTAINER_TYPE:s} with error: {exception!s}.'))
+
+    if not json_string:
+      raise IOError((
+          f'Unable to serialize attribute container: '
+          f'{container.CONTAINER_TYPE:s}'))
 
     self._redis_client.hsetnx(redis_hash_name, redis_key, json_string)
 
@@ -449,19 +471,18 @@ def GetAttributeContainerByIndex(self, container_type, index):
           RedisAttributeContainerStore, self).GetAttributeContainerByIndex(
               container_type, index)
 
-    sequence_number = index + 1
+    identifier = containers_interface.AttributeContainerIdentifier(
+        name=container_type, sequence_number=index + 1)
+
     redis_hash_name = self._GetRedisHashName(container_type)
+    redis_key = identifier.CopyToString()
 
-    serialized_data = self._redis_client.hget(
-        redis_hash_name, f'{container_type:s}.{sequence_number:d}')
+    serialized_data = self._redis_client.hget(redis_hash_name, redis_key)
     if not serialized_data:
       return None
 
     container = self._DeserializeAttributeContainer(
         container_type, serialized_data)
-
-    identifier = containers_interface.AttributeContainerIdentifier(
-        name=container_type, sequence_number=sequence_number)
     container.SetIdentifier(identifier)
 
     if container.CONTAINER_TYPE == self._CONTAINER_TYPE_EVENT_DATA:
@@ -495,16 +516,12 @@ def GetAttributeContainers(self, container_type, filter_expression=None):
       redis_hash_name = self._GetRedisHashName(container_type)
       for redis_key, serialized_data in self._redis_client.hscan_iter(
           redis_hash_name):
-        redis_key = redis_key.decode('utf-8')
-
         container = self._DeserializeAttributeContainer(
             container_type, serialized_data)
 
-        _, sequence_number = redis_key.split('.')
-        sequence_number = int(sequence_number, 10)
-        identifier = containers_interface.AttributeContainerIdentifier(
-            name=container_type, sequence_number=sequence_number)
-        container.SetIdentifier(identifier)
+        redis_key = redis_key.decode('utf-8')
+        identifier = containers_interface.AttributeContainerIdentifier()
+        identifier.CopyFromString(redis_key)
 
         if container.CONTAINER_TYPE == self._CONTAINER_TYPE_EVENT_DATA:
           identifier = getattr(container, '_event_data_stream_identifier', None)
@@ -538,11 +555,9 @@ def GetSortedEvents(self, time_range=None):
 
     for redis_key, _ in self._redis_client.zscan_iter(event_index_name):
       redis_key = redis_key.decode('utf-8')
+      identifier = containers_interface.AttributeContainerIdentifier()
+      identifier.CopyFromString(redis_key)
 
-      container_type, sequence_number = redis_key.split('.')
-      sequence_number = int(sequence_number, 10)
-      identifier = containers_interface.AttributeContainerIdentifier(
-          name=container_type, sequence_number=sequence_number)
       yield self.GetAttributeContainerByIdentifier(
           self._CONTAINER_TYPE_EVENT, identifier)