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] 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