diff --git a/swh/storage/algos/directory.py b/swh/storage/algos/directory.py new file mode 100644 --- /dev/null +++ b/swh/storage/algos/directory.py @@ -0,0 +1,38 @@ +# Copyright (C) 2018-2022 The Software Heritage developers +# See the AUTHORS file at the top-level directory of this distribution +# License: GNU General Public License version 3, or any later version +# See top-level LICENSE file for more information + +from typing import Iterable, Optional + +from swh.core.api.classes import stream_results_optional +from swh.model.model import ( + Sha1Git, + Directory, + DirectoryEntry, +) +from swh.storage.interface import StorageInterface + + +def directory_get_all_entries( + storage: StorageInterface, directory_id: Sha1Git +) -> Optional[Directory]: + """Get all the entries for a given directory + + Args: + storage (swh.storage.interface.StorageInterface): the storage instance + directory_id (bytes): the directory's identifier + Returns: + The directory if it could be properly put back together. + """ + entries: Optional[Iterable[DirectoryEntry]] = stream_results_optional( + storage.directory_get_entries, + directory_id=directory_id, + ) + if entries is None: + return None + return Directory( + id=directory_id, + entries=tuple(entries), + raw_manifest=storage.directory_get_raw_manifest([directory_id])[directory_id], + ) diff --git a/swh/storage/tests/storage_tests.py b/swh/storage/tests/storage_tests.py --- a/swh/storage/tests/storage_tests.py +++ b/swh/storage/tests/storage_tests.py @@ -41,6 +41,7 @@ ) from swh.model.swhids import CoreSWHID, ObjectType from swh.storage import get_storage +from swh.storage.algos.directory import directory_get_all_entries from swh.storage.cassandra.storage import CassandraStorage from swh.storage.common import origin_url_to_sha1 as sha1 from swh.storage.exc import HashCollision, StorageArgumentException @@ -777,15 +778,11 @@ swh_storage.directory_add(directories) for directory in directories: - assert directory == Directory( - id=directory.id, - entries=tuple( - stream_results(swh_storage.directory_get_entries, directory.id) - ), - raw_manifest=swh_storage.directory_get_raw_manifest([directory.id])[ - directory.id - ], + rebuilt_directory = directory_get_all_entries( + swh_storage.directory_get_entries, directory.id ) + assert rebuilt_directory is not None + assert directory == directory def test_directory_add_twice(self, swh_storage, sample_data): directory = sample_data.directories[1]