Raw File
objstorage.py
# Copyright (C) 2020 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 Dict, Iterable, Optional

from swh.storage.interface import Sha1
from swh.model.model import Content, MissingData
from swh.objstorage import get_objstorage
from swh.objstorage.exc import ObjNotFoundError

from .exc import StorageArgumentException


class ObjStorage:
    """Objstorage collaborator in charge of adding objects to
    the objstorage.

    """

    def __init__(self, objstorage_config: Dict):
        self.objstorage = get_objstorage(**objstorage_config)

    def __getattr__(self, key):
        if key == "objstorage":
            raise AttributeError(key)
        return getattr(self.objstorage, key)

    def content_get(self, obj_id: Sha1) -> Optional[bytes]:
        """Retrieve data associated to the content from the objstorage

        Args:
            content: content identitier

        Returns:
            associated content's data if any, None otherwise.

        """
        try:
            data = self.objstorage.get(obj_id)
        except ObjNotFoundError:
            data = None

        return data

    def content_add(self, contents: Iterable[Content]) -> Dict:
        """Add contents to the objstorage.

        Args:
            contents: List of contents to add1

        Returns:
            The summary dict of content and content bytes added to the
            objstorage.

        """
        try:
            contents = [c.with_data() for c in contents]
        except MissingData:
            raise StorageArgumentException("Missing data") from None
        summary = self.objstorage.add_batch({cont.sha1: cont.data for cont in contents})
        return {
            "content:add": summary["object:add"],
            "content:add:bytes": summary["object:add:bytes"],
        }
back to top