From 4e76e515bc19ff819f824d6ebe4bc8c0f7f62aa5 Mon Sep 17 00:00:00 2001
From: ramirezfranciscof <ramirezfranciscof@gmail.com>
Date: Mon, 24 Jan 2022 11:10:23 +0000
Subject: [PATCH 1/4] Improve storage maintain related methods

Moved the general purpose functions into methods of the storage class

Use the lock mechanism for the maintain operation
---
 aiida/cmdline/commands/cmd_storage.py         |  32 +++---
 aiida/orm/implementation/storage_backend.py   |  23 ++++
 aiida/repository/backend/disk_object_store.py |   7 +-
 aiida/storage/__init__.py                     |   2 -
 aiida/storage/control.py                      | 101 ------------------
 aiida/storage/log.py                          |  14 +++
 aiida/storage/psql_dos/backend.py             |  56 +++++++++-
 .../archive/implementations/sqlite/backend.py |   6 ++
 tests/cmdline/commands/test_storage.py        |  25 +++--
 .../backend/test_disk_object_store.py         |   2 +-
 .../test_backend.py}                          |  68 ++++--------
 11 files changed, 164 insertions(+), 172 deletions(-)
 delete mode 100644 aiida/storage/control.py
 create mode 100644 aiida/storage/log.py
 rename tests/storage/{test_control.py => psql_dos/test_backend.py} (70%)

diff --git a/aiida/cmdline/commands/cmd_storage.py b/aiida/cmdline/commands/cmd_storage.py
index 574c58fb99..8fc97badb3 100644
--- a/aiida/cmdline/commands/cmd_storage.py
+++ b/aiida/cmdline/commands/cmd_storage.py
@@ -93,13 +93,16 @@ def storage_integrity():
 def storage_info(statistics):
     """Summarise the contents of the storage."""
     from aiida.cmdline.utils.common import get_database_summary
+    from aiida.manage.manager import get_manager
     from aiida.orm import QueryBuilder
-    from aiida.storage.control import get_repository_info
+
+    manager = get_manager()
+    storage = manager.get_profile_storage()
 
     with spinner():
         data = {
             'database': get_database_summary(QueryBuilder, statistics),
-            'repository': get_repository_info(statistics=statistics),
+            'repository': storage.get_info(statistics=statistics),
         }
 
     echo.echo_dictionary(data, sort_keys=False, fmt='yaml')
@@ -119,17 +122,22 @@ def storage_info(statistics):
 )
 def storage_maintain(full, dry_run):
     """Performs maintenance tasks on the repository."""
-    from aiida.storage.control import repository_maintain
+    from aiida.manage.manager import get_manager
+
+    manager = get_manager()
+    profile = manager.get_profile()
+    storage = manager.get_profile_storage()
 
     if full:
         echo.echo_warning(
-            '\nIn order to safely perform the full maintenance operations on the internal storage, no other '
-            'process should be using the AiiDA profile being maintained. '
-            'This includes daemon workers, verdi shells, scripts with the profile loaded, etc). '
-            'Please make sure there is nothing like this currently running and that none is started until '
-            'these procedures conclude. '
-            'For performing maintanance operations that are safe to run while actively using AiiDA, just run '
-            '`verdi storage maintain`, without the `--full` flag.\n'
+            '\nIn order to safely perform the full maintenance operations on the internal storage, the profile '
+            f'{profile.name} needs to be locked. '
+            'This means that no other process will be able to access it and will fail instead. '
+            'Moreover, if any process is already using the profile, the locking attempt will fail and you will '
+            'have to either look for these processes and kill them or wait for them to stop by themselves. '
+            'Note that this includes verdi shells, daemon workers, scripts that manually load it, etc.\n'
+            'For performing maintenance operations that are safe to run while actively using AiiDA, just run '
+            '`verdi storage maintain` without the `--full` flag.\n'
         )
 
     else:
@@ -137,7 +145,7 @@ def storage_maintain(full, dry_run):
             '\nThis command will perform all maintenance operations on the internal storage that can be safely '
             'executed while still running AiiDA. '
             'However, not all operations that are required to fully optimize disk usage and future performance '
-            'can be done in this way. '
+            'can be done in this way.\n'
             'Whenever you find the time or opportunity, please consider running `verdi repository maintenance '
             '--full` for a more complete optimization.\n'
         )
@@ -146,5 +154,5 @@ def storage_maintain(full, dry_run):
         if not click.confirm('Are you sure you want continue in this mode?'):
             return
 
-    repository_maintain(full=full, dry_run=dry_run)
+    storage.maintain(full=full, dry_run=dry_run)
     echo.echo_success('Requested maintenance procedures finished.')
diff --git a/aiida/orm/implementation/storage_backend.py b/aiida/orm/implementation/storage_backend.py
index 2541925d51..ec5e4999fd 100644
--- a/aiida/orm/implementation/storage_backend.py
+++ b/aiida/orm/implementation/storage_backend.py
@@ -243,3 +243,26 @@ def get_global_variable(self, key: str) -> Union[None, str, int, float]:
 
         :raises: `KeyError` if the setting does not exist
         """
+
+    @abc.abstractmethod
+    def maintain(self, full: bool = False, dry_run: bool = False, **kwargs) -> None:
+        """Performs maintenance tasks on the storage.
+
+        If `full == True`, then this method may attempt to block the profile associated with the
+        storage to guarantee the safety of its procedures. This will not only prevent any other
+        subsequent process from accessing that profile, but will also first check if there is
+        already any process using it and raise if that is the case. The user will have to manually
+        stop any processes that is currently accessing the profile themselves or wait for it to
+        finish on its own.
+
+        :param full:
+            flag to perform operations that require to stop using the profile to be maintained.
+
+        :param dry_run:
+            flag to only print the actions that would be taken without actually executing them.
+
+        """
+
+    @abc.abstractmethod
+    def get_info(self, statistics: bool = False) -> dict:
+        """Returns general information on the storage."""
diff --git a/aiida/repository/backend/disk_object_store.py b/aiida/repository/backend/disk_object_store.py
index ee0d9c3b60..6dd0ae19cc 100644
--- a/aiida/repository/backend/disk_object_store.py
+++ b/aiida/repository/backend/disk_object_store.py
@@ -7,6 +7,7 @@
 from disk_objectstore import Container
 
 from aiida.common.lang import type_check
+from aiida.storage.log import STORAGE_LOGGER
 
 from .abstract import AbstractRepositoryBackend
 
@@ -14,6 +15,8 @@
 
 BYTES_TO_MB = 1 / 1024**2
 
+logger = STORAGE_LOGGER.getChild('disk_object_store')
+
 
 class DiskObjectStoreRepositoryBackend(AbstractRepositoryBackend):
     """Implementation of the ``AbstractRepositoryBackend`` using the ``disk-object-store`` as the backend.
@@ -156,10 +159,6 @@ def maintain( # type: ignore[override] # pylint: disable=arguments-differ,too-ma
         :param do_vacuum:flag for forcing the vacuuming of the internal database when cleaning the repository.
         :return:a dictionary with information on the operations performed.
         """
-        from aiida.storage.control import MAINTAIN_LOGGER
-
-        logger = MAINTAIN_LOGGER.getChild('disk_object_store')
-
         if live and (do_repack or clean_storage or do_vacuum):
             overrides = {'do_repack': do_repack, 'clean_storage': clean_storage, 'do_vacuum': do_vacuum}
             keys = ', '.join([key for key, override in overrides if override is True])  # type: ignore
diff --git a/aiida/storage/__init__.py b/aiida/storage/__init__.py
index c792ca6718..d10f4a7799 100644
--- a/aiida/storage/__init__.py
+++ b/aiida/storage/__init__.py
@@ -14,10 +14,8 @@
 # yapf: disable
 # pylint: disable=wildcard-import
 
-from .control import *
 
 __all__ = (
-    'MAINTAIN_LOGGER',
 )
 
 # yapf: enable
diff --git a/aiida/storage/control.py b/aiida/storage/control.py
deleted file mode 100644
index 96d439f571..0000000000
--- a/aiida/storage/control.py
+++ /dev/null
@@ -1,101 +0,0 @@
-# -*- coding: utf-8 -*-
-###########################################################################
-# Copyright (c), The AiiDA team. All rights reserved.                     #
-# This file is part of the AiiDA code.                                    #
-#                                                                         #
-# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
-# For further information on the license, see the LICENSE.txt file        #
-# For further information please visit http://www.aiida.net               #
-###########################################################################
-"""Module for overall repository control commands."""
-# Note: these functions are not methods of `AbstractRepositoryBackend` because they need access to the orm.
-# This is because they have to go through all the nodes to gather the list of keys that AiiDA is keeping
-# track of (since they are descentralized in each node entry).
-# See the get_unreferenced_keyset function
-from typing import TYPE_CHECKING, Optional, Set
-
-from aiida.common.log import AIIDA_LOGGER
-from aiida.manage import get_manager
-
-if TYPE_CHECKING:
-    from aiida.orm.implementation import StorageBackend
-
-__all__ = ('MAINTAIN_LOGGER',)
-
-MAINTAIN_LOGGER = AIIDA_LOGGER.getChild('maintain')
-
-
-def repository_maintain(
-    full: bool = False,
-    dry_run: bool = False,
-    backend: Optional['StorageBackend'] = None,
-    **kwargs,
-) -> None:
-    """Performs maintenance tasks on the repository.
-
-        :param full:
-            flag to perform operations that require to stop using the maintained profile.
-
-        :param dry_run:
-            flag to only print the actions that would be taken without actually executing them.
-
-        :param backend:
-            specific backend in which to apply the maintenance (defaults to current profile).
-    """
-
-    if backend is None:
-        backend = get_manager().get_profile_storage()
-    repository = backend.get_repository()
-
-    unreferenced_objects = get_unreferenced_keyset(aiida_backend=backend)
-    MAINTAIN_LOGGER.info(f'Deleting {len(unreferenced_objects)} unreferenced objects ...')
-    if not dry_run:
-        repository.delete_objects(list(unreferenced_objects))
-
-    MAINTAIN_LOGGER.info('Starting repository-specific operations ...')
-    repository.maintain(live=not full, dry_run=dry_run, **kwargs)
-
-
-def get_unreferenced_keyset(check_consistency: bool = True,
-                            aiida_backend: Optional['StorageBackend'] = None) -> Set[str]:
-    """Returns the keyset of objects that exist in the repository but are not tracked by AiiDA.
-
-    This should be all the soft-deleted files.
-
-        :param check_consistency:
-            toggle for a check that raises if there are references in the database with no actual object in the
-            underlying repository.
-
-        :param aiida_backend:
-            specific backend in which to apply the operation (defaults to current profile).
-
-        :return:
-            a set with all the objects in the underlying repository that are not referenced in the database.
-    """
-    from aiida import orm
-    MAINTAIN_LOGGER.info('Obtaining unreferenced object keys ...')
-
-    if aiida_backend is None:
-        aiida_backend = get_manager().get_profile_storage()
-
-    repository = aiida_backend.get_repository()
-
-    keyset_repository = set(repository.list_objects())
-    keyset_database = set(orm.Node.objects(aiida_backend).iter_repo_keys())
-
-    if check_consistency:
-        keyset_missing = keyset_database - keyset_repository
-        if len(keyset_missing) > 0:
-            raise RuntimeError(
-                'There are objects referenced in the database that are not present in the repository. Aborting!'
-            )
-
-    return keyset_repository - keyset_database
-
-
-def get_repository_info(statistics: bool = False, backend: Optional['StorageBackend'] = None) -> dict:
-    """Returns general information on the repository."""
-    if backend is None:
-        backend = get_manager().get_profile_storage()
-    repository = backend.get_repository()
-    return repository.get_info(statistics)
diff --git a/aiida/storage/log.py b/aiida/storage/log.py
new file mode 100644
index 0000000000..11ef376b36
--- /dev/null
+++ b/aiida/storage/log.py
@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+###########################################################################
+# Copyright (c), The AiiDA team. All rights reserved.                     #
+# This file is part of the AiiDA code.                                    #
+#                                                                         #
+# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
+# For further information on the license, see the LICENSE.txt file        #
+# For further information please visit http://www.aiida.net               #
+###########################################################################
+"""Initialize the storage logger."""
+
+from aiida.common.log import AIIDA_LOGGER
+
+STORAGE_LOGGER = AIIDA_LOGGER.getChild('storage')
diff --git a/aiida/storage/psql_dos/backend.py b/aiida/storage/psql_dos/backend.py
index 0ffa7a3035..983f0ae2e5 100644
--- a/aiida/storage/psql_dos/backend.py
+++ b/aiida/storage/psql_dos/backend.py
@@ -11,7 +11,7 @@
 # pylint: disable=missing-function-docstring
 from contextlib import contextmanager, nullcontext
 import functools
-from typing import TYPE_CHECKING, Iterator, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Iterator, List, Optional, Sequence, Set, Union
 
 from disk_objectstore import Container
 from sqlalchemy import table
@@ -22,6 +22,7 @@
 from aiida.orm import User
 from aiida.orm.entities import EntityTypes
 from aiida.orm.implementation import BackendEntity, StorageBackend
+from aiida.storage.log import STORAGE_LOGGER
 from aiida.storage.psql_dos.migrator import REPOSITORY_UUID_KEY, PsqlDostoreMigrator
 from aiida.storage.psql_dos.models import base
 
@@ -346,3 +347,56 @@ def get_global_variable(self, key: str) -> Union[None, str, int, float]:
             if setting is None:
                 raise KeyError(f'No setting found with key {key}')
             return setting.val
+
+    def maintain(self, full: bool = False, dry_run: bool = False, **kwargs) -> None:
+        from aiida.manage.profile_access import ProfileAccessManager
+
+        repository = self.get_repository()
+
+        if full:
+            maintenance_context = ProfileAccessManager(self._profile).lock
+        else:
+            maintenance_context = nullcontext
+
+        with maintenance_context():
+            unreferenced_objects = self.get_unreferenced_keyset()
+            STORAGE_LOGGER.info(f'Deleting {len(unreferenced_objects)} unreferenced objects ...')
+            if not dry_run:
+                repository.delete_objects(list(unreferenced_objects))
+
+            STORAGE_LOGGER.info('Starting repository-specific operations ...')
+            repository.maintain(live=not full, dry_run=dry_run, **kwargs)
+
+    def get_unreferenced_keyset(self, check_consistency: bool = True) -> Set[str]:
+        """Returns the keyset of objects that exist in the repository but are not tracked by AiiDA.
+
+        This should be all the soft-deleted files.
+
+        :param check_consistency:
+            toggle for a check that raises if there are references in the database with no actual object in the
+            underlying repository.
+
+        :return:
+            a set with all the objects in the underlying repository that are not referenced in the database.
+        """
+        from aiida import orm
+
+        STORAGE_LOGGER.info('Obtaining unreferenced object keys ...')
+
+        repository = self.get_repository()
+
+        keyset_repository = set(repository.list_objects())
+        keyset_database = set(orm.Node.objects(self).iter_repo_keys())
+
+        if check_consistency:
+            keyset_missing = keyset_database - keyset_repository
+            if len(keyset_missing) > 0:
+                raise RuntimeError(
+                    'There are objects referenced in the database that are not present in the repository. Aborting!'
+                )
+
+        return keyset_repository - keyset_database
+
+    def get_info(self, statistics: bool = False) -> dict:
+        repository = self.get_repository()
+        return repository.get_info(statistics)
diff --git a/aiida/tools/archive/implementations/sqlite/backend.py b/aiida/tools/archive/implementations/sqlite/backend.py
index 06afc51907..18752fc5b9 100644
--- a/aiida/tools/archive/implementations/sqlite/backend.py
+++ b/aiida/tools/archive/implementations/sqlite/backend.py
@@ -376,6 +376,12 @@ def get_global_variable(self, key: str):
     def set_global_variable(self, key: str, value, description: Optional[str] = None, overwrite=True) -> None:
         raise ReadOnlyError()
 
+    def maintain(self, full: bool = False, dry_run: bool = False, **kwargs) -> None:
+        raise NotImplementedError
+
+    def get_info(self, statistics: bool = False) -> dict:
+        raise NotImplementedError
+
 
 def create_backend_cls(base_class, model_cls):
     """Create an archive backend class for the given model class."""
diff --git a/tests/cmdline/commands/test_storage.py b/tests/cmdline/commands/test_storage.py
index c8f4e6f2f9..e3cf329a06 100644
--- a/tests/cmdline/commands/test_storage.py
+++ b/tests/cmdline/commands/test_storage.py
@@ -118,19 +118,32 @@ def tests_storage_maintain_logging(run_cli_command, monkeypatch, caplog):
     """Test all the information and cases of the storage maintain command."""
     import logging
 
-    from aiida.storage import control
+    from aiida.manage import get_manager
+    storage = get_manager().get_profile_storage()
+
+    def mock_maintain(*args, **kwargs):
+        """Mocks for `maintain` method of `storage`, logging the inputs passed."""
+        log_message = ''
+
+        log_message += 'Provided args:\n'
+        for arg in args:
+            log_message += f' > {arg}\n'
 
-    def mock_maintain(**kwargs):
-        logmsg = 'Provided kwargs:\n'
+        log_message += 'Provided kwargs:\n'
         for key, val in kwargs.items():
-            logmsg += f' > {key}: {val}\n'
-        logging.info(logmsg)
+            log_message += f' > {key}: {val}\n'
 
-    monkeypatch.setattr(control, 'repository_maintain', mock_maintain)
+        logging.info(log_message)
+
+    monkeypatch.setattr(storage, 'maintain', mock_maintain)
 
     with caplog.at_level(logging.INFO):
         _ = run_cli_command(cmd_storage.storage_maintain, user_input='Y')
 
+    for record in caplog.records:
+        print(record.msg.splitlines())
+        print('-----')
+
     message_list = caplog.records[0].msg.splitlines()
     assert ' > full: False' in message_list
     assert ' > dry_run: False' in message_list
diff --git a/tests/repository/backend/test_disk_object_store.py b/tests/repository/backend/test_disk_object_store.py
index ac4e0e4b0a..a0b8e5406f 100644
--- a/tests/repository/backend/test_disk_object_store.py
+++ b/tests/repository/backend/test_disk_object_store.py
@@ -281,7 +281,7 @@ def test_maintain_logging(caplog, populated_repository, do_vacuum):
     list_of_logmsg = []
     for record in caplog.records:
         assert record.levelname == 'REPORT'
-        assert record.name == 'aiida.maintain.disk_object_store'
+        assert record.name.endswith('.disk_object_store')
         list_of_logmsg.append(record.msg)
 
     assert 'packing' in list_of_logmsg[0].lower()
diff --git a/tests/storage/test_control.py b/tests/storage/psql_dos/test_backend.py
similarity index 70%
rename from tests/storage/test_control.py
rename to tests/storage/psql_dos/test_backend.py
index ebcfb6b012..1f71d017f6 100644
--- a/tests/storage/test_control.py
+++ b/tests/storage/psql_dos/test_backend.py
@@ -7,40 +7,13 @@
 # For further information on the license, see the LICENSE.txt file        #
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
-"""Tests for the :mod:`aiida.storage.control` module."""
+# pylint: disable=import-error,no-name-in-module,no-member,protected-access
+"""Testing the general methods of the psql_dos backend."""
 import pytest
 
 from aiida.manage import get_manager
 
 
-class MockRepositoryBackend():
-    """Mock of the RepositoryBackend for testing purposes."""
-
-    # pylint: disable=no-self-use
-
-    def get_info(self, *args, **kwargs):  # pylint: disable=unused-argument
-        """Method to return information."""
-        return 'this is information about the repo'
-
-    def delete_objects(self, *args, **kwargs):  # pylint: disable=unused-argument
-        """Method to delete objects."""
-
-    def maintain(self, live=True, dry_run=False, **kwargs):
-        """Method to perform maintainance operations."""
-
-        if live:
-            raise ValueError('Signal that live=True')
-
-        elif dry_run:
-            raise ValueError('Signal that dry_run=True')
-
-        elif len(kwargs) > 0:
-            raise ValueError('Signal that kwargs were passed')
-
-        else:
-            raise ValueError('Signal that live and dry_run are False')
-
-
 @pytest.fixture(scope='function')
 def clear_storage_before_test(aiida_profile_clean):  # pylint: disable=unused-argument
     """Clears the storage before a test."""
@@ -65,10 +38,11 @@ def test_get_unreferenced_keyset():
     from io import BytesIO
 
     from aiida import orm
-    from aiida.storage.control import get_unreferenced_keyset
+
+    storage_backend = get_manager().get_profile_storage()
 
     # Coverage code pass
-    unreferenced_keyset = get_unreferenced_keyset()
+    unreferenced_keyset = storage_backend.get_unreferenced_keyset()
     assert unreferenced_keyset == set()
 
     # Error catching: put a file, get the keys from the aiida db, manually delete the keys
@@ -83,16 +57,19 @@ def test_get_unreferenced_keyset():
     repository_backend = aiida_backend.get_repository()
     repository_backend.delete_objects(keys)
 
-    with pytest.raises(
-        RuntimeError, match='There are objects referenced in the database that are not present in the repository'
-    ) as exc:
-        get_unreferenced_keyset()
+    expected_error = 'There are objects referenced in the database that are not present in the repository'
+    with pytest.raises(RuntimeError, match=expected_error) as exc:
+        storage_backend.get_unreferenced_keyset()
     assert 'aborting' in str(exc.value).lower()
 
 
 #yapf: disable
-@pytest.mark.parametrize(('kwargs', 'logged_texts'), (
+@pytest.mark.parametrize(
     (
+        'kwargs',
+        'logged_texts'
+    ),
+    ((
         {},
         [' > live: True', ' > dry_run: False']
     ),
@@ -107,11 +84,11 @@ def test_get_unreferenced_keyset():
 ))
 # yapf: enable
 @pytest.mark.usefixtures('clear_storage_before_test')
-def test_repository_maintain(caplog, monkeypatch, kwargs, logged_texts):
-    """Test the ``repository_maintain`` method."""
+def test_maintain(caplog, monkeypatch, kwargs, logged_texts):
+    """Test the ``maintain`` method."""
     import logging
 
-    from aiida.storage.control import repository_maintain
+    storage_backend = get_manager().get_profile_storage()
 
     def mock_maintain(self, live=True, dry_run=False, **kwargs):  # pylint: disable=unused-argument
         logmsg = 'keywords provided:\n'
@@ -125,16 +102,17 @@ def mock_maintain(self, live=True, dry_run=False, **kwargs):  # pylint: disable=
     monkeypatch.setattr(RepoBackendClass, 'maintain', mock_maintain)
 
     with caplog.at_level(logging.INFO):
-        repository_maintain(**kwargs)
+        storage_backend.maintain(**kwargs)
 
     message_list = caplog.records[0].msg.splitlines()
     for text in logged_texts:
         assert text in message_list
 
 
-def test_repository_info(monkeypatch):
-    """Test the ``repository_info`` method."""
-    from aiida.storage.control import get_repository_info
+def test_get_info(monkeypatch):
+    """Test the ``get_info`` method."""
+
+    storage_backend = get_manager().get_profile_storage()
 
     def mock_get_info(self, statistics=False, **kwargs):  # pylint: disable=unused-argument
         output = {'value': 42}
@@ -145,12 +123,12 @@ def mock_get_info(self, statistics=False, **kwargs):  # pylint: disable=unused-a
     RepoBackendClass = get_manager().get_profile_storage().get_repository().__class__  # pylint: disable=invalid-name
     monkeypatch.setattr(RepoBackendClass, 'get_info', mock_get_info)
 
-    repository_info_out = get_repository_info()
+    repository_info_out = storage_backend.get_info()
     assert 'value' in repository_info_out
     assert 'extra_value' not in repository_info_out
     assert repository_info_out['value'] == 42
 
-    repository_info_out = get_repository_info(statistics=True)
+    repository_info_out = storage_backend.get_info(statistics=True)
     assert 'value' in repository_info_out
     assert 'extra_value' in repository_info_out
     assert repository_info_out['value'] == 42

From c98fb5c629fdd41ff5932cbf4154c22778803999 Mon Sep 17 00:00:00 2001
From: ramirezfranciscof <ramirezfranciscof@gmail.com>
Date: Wed, 23 Feb 2022 08:51:26 +0000
Subject: [PATCH 2/4] Apply feedback from PR review

---
 aiida/cmdline/commands/cmd_storage.py       |  5 +++--
 aiida/orm/implementation/storage_backend.py | 12 ++++++++++--
 tests/cmdline/commands/test_storage.py      |  4 ----
 tests/storage/psql_dos/test_backend.py      |  2 +-
 4 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/aiida/cmdline/commands/cmd_storage.py b/aiida/cmdline/commands/cmd_storage.py
index 8fc97badb3..40a1fa5a17 100644
--- a/aiida/cmdline/commands/cmd_storage.py
+++ b/aiida/cmdline/commands/cmd_storage.py
@@ -120,12 +120,13 @@ def storage_info(statistics):
     help=
     'Run the maintenance in dry-run mode which will print actions that would be taken without actually executing them.'
 )
-def storage_maintain(full, dry_run):
+@click.pass_context
+def storage_maintain(ctx, full, dry_run):
     """Performs maintenance tasks on the repository."""
     from aiida.manage.manager import get_manager
 
     manager = get_manager()
-    profile = manager.get_profile()
+    profile = ctx.obj.profile
     storage = manager.get_profile_storage()
 
     if full:
diff --git a/aiida/orm/implementation/storage_backend.py b/aiida/orm/implementation/storage_backend.py
index ec5e4999fd..8adbae5e6d 100644
--- a/aiida/orm/implementation/storage_backend.py
+++ b/aiida/orm/implementation/storage_backend.py
@@ -246,7 +246,7 @@ def get_global_variable(self, key: str) -> Union[None, str, int, float]:
 
     @abc.abstractmethod
     def maintain(self, full: bool = False, dry_run: bool = False, **kwargs) -> None:
-        """Performs maintenance tasks on the storage.
+        """Perform maintenance tasks on the storage.
 
         If `full == True`, then this method may attempt to block the profile associated with the
         storage to guarantee the safety of its procedures. This will not only prevent any other
@@ -265,4 +265,12 @@ def maintain(self, full: bool = False, dry_run: bool = False, **kwargs) -> None:
 
     @abc.abstractmethod
     def get_info(self, statistics: bool = False) -> dict:
-        """Returns general information on the storage."""
+        """Return general information on the storage.
+
+        :param statistics:
+            flag to request more detailed information about the content of the storage.
+
+        :returns:
+            a nested dict with the relevant information.
+
+        """
diff --git a/tests/cmdline/commands/test_storage.py b/tests/cmdline/commands/test_storage.py
index e3cf329a06..d172a291b9 100644
--- a/tests/cmdline/commands/test_storage.py
+++ b/tests/cmdline/commands/test_storage.py
@@ -140,10 +140,6 @@ def mock_maintain(*args, **kwargs):
     with caplog.at_level(logging.INFO):
         _ = run_cli_command(cmd_storage.storage_maintain, user_input='Y')
 
-    for record in caplog.records:
-        print(record.msg.splitlines())
-        print('-----')
-
     message_list = caplog.records[0].msg.splitlines()
     assert ' > full: False' in message_list
     assert ' > dry_run: False' in message_list
diff --git a/tests/storage/psql_dos/test_backend.py b/tests/storage/psql_dos/test_backend.py
index 1f71d017f6..8cc812db0a 100644
--- a/tests/storage/psql_dos/test_backend.py
+++ b/tests/storage/psql_dos/test_backend.py
@@ -63,7 +63,7 @@ def test_get_unreferenced_keyset():
     assert 'aborting' in str(exc.value).lower()
 
 
-#yapf: disable
+# yapf: disable
 @pytest.mark.parametrize(
     (
         'kwargs',

From d98e3014eb96dbb2674bc30c6cfbf9b5fa9136d4 Mon Sep 17 00:00:00 2001
From: ramirezfranciscof <ramirezfranciscof@gmail.com>
Date: Wed, 23 Feb 2022 11:12:38 +0000
Subject: [PATCH 3/4] Generalize the get_info method.

---
 aiida/cmdline/commands/cmd_storage.py       | 6 ++----
 aiida/orm/implementation/storage_backend.py | 3 ++-
 aiida/storage/psql_dos/backend.py           | 6 +++++-
 tests/storage/psql_dos/test_backend.py      | 9 +++++++--
 4 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/aiida/cmdline/commands/cmd_storage.py b/aiida/cmdline/commands/cmd_storage.py
index 40a1fa5a17..40e078207c 100644
--- a/aiida/cmdline/commands/cmd_storage.py
+++ b/aiida/cmdline/commands/cmd_storage.py
@@ -100,10 +100,8 @@ def storage_info(statistics):
     storage = manager.get_profile_storage()
 
     with spinner():
-        data = {
-            'database': get_database_summary(QueryBuilder, statistics),
-            'repository': storage.get_info(statistics=statistics),
-        }
+        data = storage.get_info(statistics=statistics)
+        data['database']['summary'] = get_database_summary(QueryBuilder, statistics)
 
     echo.echo_dictionary(data, sort_keys=False, fmt='yaml')
 
diff --git a/aiida/orm/implementation/storage_backend.py b/aiida/orm/implementation/storage_backend.py
index 8adbae5e6d..e0768fda10 100644
--- a/aiida/orm/implementation/storage_backend.py
+++ b/aiida/orm/implementation/storage_backend.py
@@ -271,6 +271,7 @@ def get_info(self, statistics: bool = False) -> dict:
             flag to request more detailed information about the content of the storage.
 
         :returns:
-            a nested dict with the relevant information.
+            a nested dict with the relevant information (with at least one key for `database`
+            and one for `repository`).
 
         """
diff --git a/aiida/storage/psql_dos/backend.py b/aiida/storage/psql_dos/backend.py
index 983f0ae2e5..379f963af3 100644
--- a/aiida/storage/psql_dos/backend.py
+++ b/aiida/storage/psql_dos/backend.py
@@ -399,4 +399,8 @@ def get_unreferenced_keyset(self, check_consistency: bool = True) -> Set[str]:
 
     def get_info(self, statistics: bool = False) -> dict:
         repository = self.get_repository()
-        return repository.get_info(statistics)
+        output_dict = {
+            'database': {},
+            'repository': repository.get_info(statistics),
+        }
+        return output_dict
diff --git a/tests/storage/psql_dos/test_backend.py b/tests/storage/psql_dos/test_backend.py
index 8cc812db0a..5c8e325056 100644
--- a/tests/storage/psql_dos/test_backend.py
+++ b/tests/storage/psql_dos/test_backend.py
@@ -123,12 +123,17 @@ def mock_get_info(self, statistics=False, **kwargs):  # pylint: disable=unused-a
     RepoBackendClass = get_manager().get_profile_storage().get_repository().__class__  # pylint: disable=invalid-name
     monkeypatch.setattr(RepoBackendClass, 'get_info', mock_get_info)
 
-    repository_info_out = storage_backend.get_info()
+    storage_info_out = storage_backend.get_info()
+    assert 'database' in storage_info_out
+    assert 'repository' in storage_info_out
+
+    repository_info_out = storage_info_out['repository']
     assert 'value' in repository_info_out
     assert 'extra_value' not in repository_info_out
     assert repository_info_out['value'] == 42
 
-    repository_info_out = storage_backend.get_info(statistics=True)
+    storage_info_out = storage_backend.get_info(statistics=True)
+    repository_info_out = storage_info_out['repository']
     assert 'value' in repository_info_out
     assert 'extra_value' in repository_info_out
     assert repository_info_out['value'] == 42

From 1b0c41315ed8b74f7cbe8c5696be74becbd2d03d Mon Sep 17 00:00:00 2001
From: ramirezfranciscof <ramirezfranciscof@gmail.com>
Date: Wed, 23 Feb 2022 11:37:46 +0000
Subject: [PATCH 4/4] Final details

---
 aiida/cmdline/commands/cmd_storage.py       | 1 +
 aiida/orm/implementation/storage_backend.py | 3 +--
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/aiida/cmdline/commands/cmd_storage.py b/aiida/cmdline/commands/cmd_storage.py
index 40e078207c..3ea75ea6e5 100644
--- a/aiida/cmdline/commands/cmd_storage.py
+++ b/aiida/cmdline/commands/cmd_storage.py
@@ -101,6 +101,7 @@ def storage_info(statistics):
 
     with spinner():
         data = storage.get_info(statistics=statistics)
+        data.setdefault('database', {})
         data['database']['summary'] = get_database_summary(QueryBuilder, statistics)
 
     echo.echo_dictionary(data, sort_keys=False, fmt='yaml')
diff --git a/aiida/orm/implementation/storage_backend.py b/aiida/orm/implementation/storage_backend.py
index e0768fda10..8adbae5e6d 100644
--- a/aiida/orm/implementation/storage_backend.py
+++ b/aiida/orm/implementation/storage_backend.py
@@ -271,7 +271,6 @@ def get_info(self, statistics: bool = False) -> dict:
             flag to request more detailed information about the content of the storage.
 
         :returns:
-            a nested dict with the relevant information (with at least one key for `database`
-            and one for `repository`).
+            a nested dict with the relevant information.
 
         """