Changes to SQLite-based attribute container store (#58)

acstore/helpers/json_serializer.py

@@ -45,10 +45,10 @@ def ConvertAttributeContainerToJSON(cls, attribute_container):
 
     for attribute_name, attribute_value in attribute_container.GetAttributes():
       data_type = schema.get(attribute_name, None)
-      if data_type:
-        serializer = schema_helper.SchemaHelper.GetAttributeSerializer(
-            data_type, 'json')
+      serializer = schema_helper.SchemaHelper.GetAttributeSerializer(
+          data_type, 'json')
 
+      if serializer:
         attribute_value = serializer.SerializeValue(attribute_value)
 
       # JSON will not serialize certain runtime types like set, therefore

acstore/interface.py

@@ -108,16 +108,16 @@ def _SetAttributeContainerNextSequenceNumber(
         container_type] = next_sequence_number
 
   @abc.abstractmethod
-  def _WriteNewAttributeContainer(self, container):
-    """Writes a new attribute container to the store.
+  def _WriteExistingAttributeContainer(self, container):
+    """Writes an existing attribute container to the store.
 
     Args:
       container (AttributeContainer): attribute container.
     """
 
   @abc.abstractmethod
-  def _WriteExistingAttributeContainer(self, container):
-    """Writes an existing attribute container to the store.
+  def _WriteNewAttributeContainer(self, container):
+    """Writes a new attribute container to the store.
 
     Args:
       container (AttributeContainer): attribute container.
@@ -150,12 +150,6 @@ def GetAttributeContainerByIdentifier(self, container_type, identifier):
 
     Returns:
       AttributeContainer: attribute container or None if not available.
-
-    Raises:
-      IOError: when the store is closed or if an unsupported identifier is
-          provided.
-      OSError: when the store is closed or if an unsupported identifier is
-          provided.
     """
 
   @abc.abstractmethod
@@ -168,10 +162,6 @@ def GetAttributeContainerByIndex(self, container_type, index):
 
     Returns:
       AttributeContainer: attribute container or None if not available.
-
-    Raises:
-      IOError: when the store is closed.
-      OSError: when the store is closed.
     """
 
   @abc.abstractmethod
@@ -185,10 +175,6 @@ def GetAttributeContainers(self, container_type, filter_expression=None):
 
     Returns:
       generator(AttributeContainer): attribute container generator.
-
-    Raises:
-      IOError: when the store is closed.
-      OSError: when the store is closed.
     """
 
   @abc.abstractmethod
@@ -280,10 +266,6 @@ def _GetCachedAttributeContainer(self, container_type, index):
 
     Returns:
       AttributeContainer: attribute container or None if not available.
-
-    Raises:
-      IOError: when there is an error querying the attribute container store.
-      OSError: when there is an error querying the attribute container store.
     """
     lookup_key = f'{container_type:s}.{index:d}'
     attribute_container = self._attribute_container_cache.get(lookup_key, None)

acstore/sqlite_store.py

@@ -500,6 +500,42 @@ def _GetAttributeContainersWithFilter(
           if self._storage_profiler:
             self._storage_profiler.StopTiming('get_containers')
 
+  def _GetNumberOfAttributeContainerRows(self, container_type):
+    """Retrieves the number of attribute container rows.
+
+    Args:
+      container_type (str): attribute container type.
+
+    Returns:
+      int: the number of rows of a specified attribute container type.
+
+    Raises:
+      IOError: when there is an error querying the attribute container store.
+      OSError: when there is an error querying the attribute container store.
+    """
+    self._CommitWriteCache(container_type)
+
+    if not self._HasTable(container_type):
+      return 0
+
+    # Note that this is SQLite specific, and will give inaccurate results if
+    # there are DELETE commands run on the table. acstore does not run any
+    # DELETE commands.
+    query = f'SELECT MAX(_ROWID_) FROM {container_type:s} LIMIT 1'
+
+    try:
+      self._cursor.execute(query)
+    except (sqlite3.InterfaceError, sqlite3.OperationalError) as exception:
+      raise IOError((
+          f'Unable to query attribute container store with error: '
+          f'{exception!s}'))
+
+    row = self._cursor.fetchone()
+    if not row:
+      return 0
+
+    return row[0] or 0
+
   def _HasTable(self, table_name):
     """Determines if a specific table exists.
 
@@ -938,33 +974,8 @@ def GetNumberOfAttributeContainers(self, container_type):
 
     Returns:
       int: the number of containers of a specified type.
-
-    Raises:
-      IOError: when there is an error querying the attribute container store.
-      OSError: when there is an error querying the attribute container store.
     """
-    self._CommitWriteCache(container_type)
-
-    if not self._HasTable(container_type):
-      return 0
-
-    # Note that this is SQLite specific, and will give inaccurate results if
-    # there are DELETE commands run on the table. acstore does not run any
-    # DELETE commands.
-    query = f'SELECT MAX(_ROWID_) FROM {container_type:s} LIMIT 1'
-
-    try:
-      self._cursor.execute(query)
-    except (sqlite3.InterfaceError, sqlite3.OperationalError) as exception:
-      raise IOError((
-          f'Unable to query attribute container store with error: '
-          f'{exception!s}'))
-
-    row = self._cursor.fetchone()
-    if not row:
-      return 0
-
-    return row[0] or 0
+    return self._attribute_container_sequence_numbers[container_type]
 
   def HasAttributeContainers(self, container_type):
     """Determines if store contains a specific type of attribute containers.
@@ -975,13 +986,8 @@ def HasAttributeContainers(self, container_type):
     Returns:
       bool: True if the store contains the specified type of attribute
           containers.
-
-    Raises:
-      IOError: when there is an error querying the attribute container store.
-      OSError: when there is an error querying the attribute container store.
     """
-    count = self.GetNumberOfAttributeContainers(container_type)
-    return count > 0
+    return self._attribute_container_sequence_numbers[container_type] > 0
 
   def Open(self, path=None, read_only=True, **unused_kwargs):  # pylint: disable=arguments-differ
     """Opens the store.
@@ -1062,6 +1068,7 @@ def Open(self, path=None, read_only=True, **unused_kwargs):  # pylint: disable=a
     # Initialize next_sequence_number based on the file contents so that
     # AttributeContainerIdentifier points to the correct attribute container.
     for container_type in self._containers_manager.GetContainerTypes():
-      next_sequence_number = self.GetNumberOfAttributeContainers(container_type)
+      next_sequence_number = self._GetNumberOfAttributeContainerRows(
+          container_type)
       self._SetAttributeContainerNextSequenceNumber(
           container_type, next_sequence_number)

tests/interface.py

@@ -59,5 +59,40 @@ def testSetStorageProfiler(self):
     test_store.SetStorageProfiler(None)
 
 
+class AttributeContainerStoreWithReadCacheTest(test_lib.BaseTestCase):
+  """Tests for the attribute container store with read cache."""
+
+  # pylint: disable=protected-access
+
+  def testCacheAttributeContainerByIndex(self):
+    """Tests the _CacheAttributeContainerByIndex function."""
+    attribute_container = test_lib.TestAttributeContainer()
+
+    with test_lib.TempDirectory():
+      test_store = interface.AttributeContainerStoreWithReadCache()
+
+      self.assertEqual(len(test_store._attribute_container_cache), 0)
+
+      test_store._CacheAttributeContainerByIndex(attribute_container, 0)
+      self.assertEqual(len(test_store._attribute_container_cache), 1)
+
+  def testGetCachedAttributeContainer(self):
+    """Tests the _GetCachedAttributeContainer function."""
+    attribute_container = test_lib.TestAttributeContainer()
+
+    with test_lib.TempDirectory():
+      test_store = interface.AttributeContainerStoreWithReadCache()
+
+      cached_container = test_store._GetCachedAttributeContainer(
+          attribute_container.CONTAINER_TYPE, 1)
+      self.assertIsNone(cached_container)
+
+      test_store._CacheAttributeContainerByIndex(attribute_container, 1)
+
+      cached_container = test_store._GetCachedAttributeContainer(
+          attribute_container.CONTAINER_TYPE, 1)
+      self.assertIsNotNone(cached_container)
+
+
 if __name__ == '__main__':
   unittest.main()

tests/sqlite_store.py

@@ -116,18 +116,6 @@ def tearDown(self):
     containers_manager.AttributeContainersManager.DeregisterAttributeContainer(
         test_lib.TestAttributeContainer)
 
-  def testCacheAttributeContainerByIndex(self):
-    """Tests the _CacheAttributeContainerByIndex function."""
-    attribute_container = test_lib.TestAttributeContainer()
-
-    with test_lib.TempDirectory():
-      test_store = sqlite_store.SQLiteAttributeContainerStore()
-
-      self.assertEqual(len(test_store._attribute_container_cache), 0)
-
-      test_store._CacheAttributeContainerByIndex(attribute_container, 0)
-      self.assertEqual(len(test_store._attribute_container_cache), 1)
-
   def testCheckStorageMetadata(self):
     """Tests the _CheckStorageMetadata function."""
     with test_lib.TempDirectory():
@@ -219,22 +207,36 @@ def testGetAttributeContainersWithFilter(self):
       finally:
         test_store.Close()
 
-  def testGetCachedAttributeContainer(self):
-    """Tests the _GetCachedAttributeContainer function."""
+  def testGetNumberOfAttributeContainerRows(self):
+    """Tests the _GetNumberOfAttributeContainerRows function."""
     attribute_container = test_lib.TestAttributeContainer()
 
-    with test_lib.TempDirectory():
+    with test_lib.TempDirectory() as temp_directory:
+      test_path = os.path.join(temp_directory, 'acstore.sqlite')
       test_store = sqlite_store.SQLiteAttributeContainerStore()
+      test_store.Open(path=test_path, read_only=False)
 
-      cached_container = test_store._GetCachedAttributeContainer(
-          attribute_container.CONTAINER_TYPE, 1)
-      self.assertIsNone(cached_container)
+      try:
+        number_of_containers = test_store._GetNumberOfAttributeContainerRows(
+            attribute_container.CONTAINER_TYPE)
+        self.assertEqual(number_of_containers, 0)
 
-      test_store._CacheAttributeContainerByIndex(attribute_container, 1)
+        test_store.AddAttributeContainer(attribute_container)
 
-      cached_container = test_store._GetCachedAttributeContainer(
-          attribute_container.CONTAINER_TYPE, 1)
-      self.assertIsNotNone(cached_container)
+        number_of_containers = test_store._GetNumberOfAttributeContainerRows(
+            attribute_container.CONTAINER_TYPE)
+        self.assertEqual(number_of_containers, 1)
+
+        # Test for a supported container type that does not have a table
+        # present in the storage file.
+        query = f'DROP TABLE {attribute_container.CONTAINER_TYPE:s}'
+        test_store._cursor.execute(query)
+        number_of_containers = test_store._GetNumberOfAttributeContainerRows(
+            attribute_container.CONTAINER_TYPE)
+        self.assertEqual(number_of_containers, 0)
+
+      finally:
+        test_store.Close()
 
   def testHasTable(self):
     """Tests the _HasTable function."""
@@ -283,19 +285,19 @@ def testWriteExistingAttributeContainer(self):
       test_store.Open(path=test_path, read_only=False)
 
       try:
-        number_of_containers = test_store.GetNumberOfAttributeContainers(
+        number_of_containers = test_store._GetNumberOfAttributeContainerRows(
             attribute_container.CONTAINER_TYPE)
         self.assertEqual(number_of_containers, 0)
 
         test_store._WriteNewAttributeContainer(attribute_container)
 
-        number_of_containers = test_store.GetNumberOfAttributeContainers(
+        number_of_containers = test_store._GetNumberOfAttributeContainerRows(
             attribute_container.CONTAINER_TYPE)
         self.assertEqual(number_of_containers, 1)
 
         test_store._WriteExistingAttributeContainer(attribute_container)
 
-        number_of_containers = test_store.GetNumberOfAttributeContainers(
+        number_of_containers = test_store._GetNumberOfAttributeContainerRows(
             attribute_container.CONTAINER_TYPE)
         self.assertEqual(number_of_containers, 1)
 
@@ -315,13 +317,13 @@ def testWriteNewAttributeContainer(self):
       test_store.Open(path=test_path, read_only=False)
 
       try:
-        number_of_containers = test_store.GetNumberOfAttributeContainers(
+        number_of_containers = test_store._GetNumberOfAttributeContainerRows(
             attribute_container.CONTAINER_TYPE)
         self.assertEqual(number_of_containers, 0)
 
         test_store._WriteNewAttributeContainer(attribute_container)
 
-        number_of_containers = test_store.GetNumberOfAttributeContainers(
+        number_of_containers = test_store._GetNumberOfAttributeContainerRows(
             attribute_container.CONTAINER_TYPE)
         self.assertEqual(number_of_containers, 1)
 
@@ -468,12 +470,8 @@ def testGetNumberOfAttributeContainers(self):
             attribute_container.CONTAINER_TYPE)
         self.assertEqual(number_of_containers, 1)
 
-        # Test for a supported container type that does not have a table
-        # present in the storage file.
-        query = f'DROP TABLE {attribute_container.CONTAINER_TYPE:s}'
-        test_store._cursor.execute(query)
         number_of_containers = test_store.GetNumberOfAttributeContainers(
-            attribute_container.CONTAINER_TYPE)
+            'bogus')
         self.assertEqual(number_of_containers, 0)
 
       finally: