Enforce context usage of Container within migrator

.github/workflows/ci-code.yml

@@ -77,7 +77,8 @@ jobs:
       # NOTE1: Python 3.12 has a performance regression when running with code coverage
       # so run code coverage only for python 3.9.
       run: |
-        pytest -n auto --db-backend ${{ matrix.database-backend }} -m 'not nightly' tests/ ${{ matrix.python-version == '3.9' && '--cov aiida' || '' }}
+        #pytest -n auto --db-backend ${{ matrix.database-backend }} -m 'not nightly' tests/ ${{ matrix.python-version == '3.9' && '--cov aiida' || '' }}
+        pytest -n auto --db-backend ${{ matrix.database-backend }} -m 'not nightly' tests/storage/sqlite/test_container.py ${{ matrix.python-version == '3.9' && '--cov aiida' || '' }}
 
     - name: Upload coverage report
       if: matrix.python-version == 3.9 && github.repository == 'aiidateam/aiida-core'

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
   'circus~=0.18.0',
   'click-spinner~=0.1.8',
   'click~=8.1',
-  'disk-objectstore~=1.2',
+  'disk-objectstore',
   'docstring-parser',
   'get-annotations~=0.1;python_version<"3.10"',
   'graphviz~=0.19',
@@ -521,3 +521,6 @@ commands = molecule {posargs:test}
 # .github/actions/install-aiida-core/action.yml
 # .readthedocs.yml
 required-version = ">=0.5.21"
+
+[tool.uv.sources]
+disk-objectstore = {git = "https://github.com/aiidateam/disk-objectstore", branch = "fix-open-session"}

src/aiida/manage/configuration/config.py

@@ -16,8 +16,6 @@
 from __future__ import annotations
 
 import codecs
-import contextlib
-import io
 import json
 import os
 import uuid
@@ -527,11 +525,12 @@ def create_profile(
 
         LOGGER.report('Initialising the storage backend.')
         try:
-            with contextlib.redirect_stdout(io.StringIO()):
-                profile.storage_cls.initialise(profile)
+            # not sure if scope is good, at least put it in log
+            # with contextlib.redirect_stdout(io.StringIO()):
+            profile.storage_cls.initialise(profile)
         except Exception as exception:
             raise StorageMigrationError(
-                f'Storage backend initialisation failed, probably because the configuration is incorrect:\n{exception}'
+                f'During initialisation of storage backend following excepion was raised: {exception}.\n Please check above for full traceback.'
             )
         LOGGER.report('Storage initialisation completed.')
 

src/aiida/storage/psql_dos/backend.py

@@ -9,7 +9,6 @@
 """SqlAlchemy implementation of `aiida.orm.implementation.backends.Backend`."""
 
 import functools
-import gc
 import pathlib
 from contextlib import contextmanager, nullcontext
 from typing import TYPE_CHECKING, Iterator, List, Optional, Sequence, Set, Union
@@ -186,16 +185,18 @@ def close(self) -> None:
             return  # the instance is already closed, and so this is a no-op
         # close the connection
 
+        # PR QUESTION: can we have multiple sessons? Then we should close all of them
+        # from sqlalchemy.orm import close_all_sessions
+        # close_all_sessions()
+
         engine = self._session_factory.bind
-        if engine is not None:
-            engine.dispose()  # type: ignore[union-attr]
         self._session_factory.expunge_all()
         self._session_factory.close()
         self._session_factory = None
 
-        # Without this, sqlalchemy keeps a weakref to a session
-        # in sqlalchemy.orm.session._sessions
-        gc.collect()
+        # IMPORTANT: Dispose engine only after sessions have been closed
+        if engine is not None:
+            engine.dispose()  # type: ignore[union-attr]
 
     def _clear(self) -> None:
         from aiida.storage.psql_dos.models.settings import DbSetting

src/aiida/storage/psql_dos/migrator.py

@@ -18,7 +18,7 @@
 
 import contextlib
 import pathlib
-from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional
+from typing import TYPE_CHECKING, Any, Dict, Generator, Iterator, Optional
 
 from alembic.command import downgrade, upgrade
 from alembic.config import Config
@@ -170,7 +170,12 @@
                 f"but the disk-objectstore's is {repository_uuid}."
             )
 
-    def get_container(self) -> 'Container':
+    def get_container(self):
+        # TODO do we need to deprecate it?
+        raise NotImplementedError()
+
+    @contextlib.contextmanager
+    def container_context(self) -> Generator['Container', None, None]:
         """Return the disk-object store container.
 
         :returns: The disk-object store container configured for the repository path of the current profile.
@@ -179,7 +184,8 @@
 
         from .backend import get_filepath_container
 
-        return Container(get_filepath_container(self.profile))
+        with Container(get_filepath_container(self.profile)) as container:
+            yield container
 
     def get_repository_uuid(self) -> str:
         """Return the UUID of the repository.
@@ -188,12 +194,18 @@
         :raises: :class:`~aiida.common.exceptions.UnreachableStorage` if the UUID cannot be retrieved, which probably
             means that the repository is not initialised.
         """
+        container = None
         try:
-            return self.get_container().container_id
+            with self.container_context() as container:
+                return container.container_id
         except Exception as exception:
-            raise exceptions.UnreachableStorage(
-                f'Could not access disk-objectstore {self.get_container()}: {exception}'
-            ) from exception
+            # PR COMMENT removed the printing of exception in message since we use from exception
+            if container is None:
+                msg = 'During creation of the container context for the disk-objectstore the following error was raised'
+            else:
+                msg = f'During access of disk-objectstore {container} error was raised.'
+
+            raise exceptions.UnreachableStorage(msg) from exception
 
     def initialise(self, reset: bool = False) -> bool:
         """Initialise the storage backend.
@@ -239,7 +251,8 @@
 
         :returns: ``True`` if the repository is initialised, ``False`` otherwise.
         """
-        return self.get_container().is_initialised
+        with self.container_context() as container:
+            return container.is_initialised
 
     @property
     def is_database_initialised(self) -> bool:
@@ -261,7 +274,8 @@
         import shutil
 
         try:
-            shutil.rmtree(self.get_container().get_folder())
+            with self.container_context() as container:
+                shutil.rmtree(container.get_folder())
         except FileNotFoundError:
             pass
 
@@ -276,8 +290,8 @@
         """Initialise the repository."""
         from aiida.storage.psql_dos.backend import CONTAINER_DEFAULTS
 
-        container = self.get_container()
-        container.init_container(clear=True, **CONTAINER_DEFAULTS)
+        with self.container_context() as container:
+            container.init_container(clear=True, **CONTAINER_DEFAULTS)
 
     def initialise_database(self) -> None:
         """Initialise the database.

src/aiida/storage/sqlite_dos/backend.py

@@ -10,11 +10,12 @@
 
 from __future__ import annotations
 
+import contextlib
 import pathlib
 from functools import cached_property, lru_cache
 from pathlib import Path
 from shutil import rmtree
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Generator, Optional
 from uuid import uuid4
 
 from alembic.config import Config
@@ -56,6 +57,7 @@
 REPOSITORY_UUID_KEY = 'repository|uuid'
 
 
+# PR COMMENT why is this not in the migrator.py?
 class SqliteDosMigrator(PsqlDosMigrator):
     """Class for validating and migrating `sqlite_dos` storage instances.
 
@@ -79,13 +81,19 @@
         self._engine = create_sqla_engine(filepath_database)
         self._connection = None
 
-    def get_container(self) -> Container:
+    def get_container(self):
+        # TODO do we need to deprecate it?
+        raise NotImplementedError()
+
+    @contextlib.contextmanager
+    def container_context(self) -> Generator['Container', None, None]:
         """Return the disk-object store container.
 
         :returns: The disk-object store container configured for the repository path of the current profile.
         """
         filepath_container = Path(self.profile.storage_config['filepath']) / FILENAME_CONTAINER
-        return Container(str(filepath_container))
+        with Container(str(filepath_container)) as container:
+            yield container
 
     def initialise_database(self) -> None:
         """Initialise the database.
@@ -271,13 +279,20 @@
             rmtree(self.filepath_root)
             LOGGER.report(f'Deleted storage directory at `{self.filepath_root}`.')
 
-    def get_container(self) -> 'Container':
-        return Container(str(self.filepath_container))
+    @contextlib.contextmanager
+    def container_context(self) -> Generator['Container', None, None]:
+        """Return the disk-object store container.
+
+        :returns: The disk-object store container configured for the repository path of the current profile.
+        """
+        with Container(str(self.filepath_container)) as container:
+            yield container
 
     def get_repository(self) -> 'DiskObjectStoreRepositoryBackend':
         from aiida.repository.backend import DiskObjectStoreRepositoryBackend
 
-        return DiskObjectStoreRepositoryBackend(container=self.get_container())
+        with self.container_context() as container:
+            return DiskObjectStoreRepositoryBackend(container=container)
 
     @classmethod
     def version_profile(cls, profile: Profile) -> Optional[str]:

tests/manage/configuration/test_config.py

@@ -445,8 +445,8 @@ def test_create_profile_raises(config_with_profile, monkeypatch, entry_points):
     config = config_with_profile
     profile_name = uuid.uuid4().hex
 
-    def raise_storage_migration_error(*args, **kwargs):
-        raise exceptions.StorageMigrationError()
+    def raise_storage_migration_error(*_, **__):
+        raise exceptions.StorageMigrationError('Monkey patchted error')
 
     monkeypatch.setattr(SqliteTempBackend, 'initialise', raise_storage_migration_error)
     entry_points.add(InvalidBaseStorage, 'aiida.storage:core.invalid_base')
@@ -460,7 +460,10 @@ def raise_storage_migration_error(*args, **kwargs):
     with pytest.raises(ValueError, match=r'The entry point `.*` could not be loaded'):
         config.create_profile(profile_name, 'core.non_existant', {})
 
-    with pytest.raises(exceptions.StorageMigrationError, match='Storage backend initialisation failed.*'):
+    with pytest.raises(
+        exceptions.StorageMigrationError,
+        match=r'During initialisation of storage backend following excepion was raised: Monkey patchted error.*',
+    ):
         config.create_profile(profile_name, 'core.sqlite_temp', {})
 
 

tests/storage/sqlite/test_container.py

@@ -0,0 +1,52 @@
+"""Test container initialization."""
+
+import os
+
+import psutil
+
+
+def test_file_descriptor_closed(aiida_profile):
+    """Checks if the number of open file descriptors change during a reset."""
+
+    def list_open_file_descriptors():
+        process = psutil.Process(os.getpid())
+        return process.open_files()
+
+    # We have some connections active due to aiida profile during the first
+    # reset these are destroyed. We check the second time changes the number of
+    # open file descriptors.
+    from aiida.manage import get_manager
+
+    storage_backend = get_manager().get_profile_storage()
+    migrator_context = storage_backend.migrator_context
+    open_file_descriptors_before = list_open_file_descriptors()
+    with migrator_context(aiida_profile) as migrator:
+        migrator.initialise_repository()
+        migrator.reset_repository()
+    open_file_descriptors_after = list_open_file_descriptors()
+    assert (
+        len(open_file_descriptors_before) == len(open_file_descriptors_after)
+    ), f'Before these file descriptor were open:\n{open_file_descriptors_before}\n Now these are open:\n{open_file_descriptors_after}'
+
+
+# PR COMMENT this is just a sanity check for me, I don' think that the test should be included in final PR
+def test_reset_storage_file_descriptor_closed(aiida_profile):
+    """Checks if the number of open file descriptors change during a reset."""
+
+    def list_open_file_descriptors():
+        process = psutil.Process(os.getpid())
+        return process.open_files()
+
+    # We have some connections active due to aiida profile during the first
+    # reset these are destroyed. We check the second time changes the number of
+    # open file descriptors.
+    # TODO The fix should keep the existing connections alive and just reuse them
+    #      then one does not need to do two resets.
+
+    aiida_profile.reset_storage()
+    open_file_descriptors_before = list_open_file_descriptors()
+    aiida_profile.reset_storage()
+    open_file_descriptors_after = list_open_file_descriptors()
+    assert (
+        len(open_file_descriptors_before) == len(open_file_descriptors_after)
+    ), f'Before these file descriptor were open:\n{open_file_descriptors_before}\n Now these are open:\n{open_file_descriptors_after}'