diff --git a/swh/vault/cookers/base.py b/swh/vault/cookers/base.py index 9e28f19..a4ffdac 100644 --- a/swh/vault/cookers/base.py +++ b/swh/vault/cookers/base.py @@ -1,128 +1,128 @@ # Copyright (C) 2016-2018 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 import abc import io import logging from swh.core import config from swh.model import hashutil from swh.storage import get_storage from swh.vault.api.client import RemoteVaultClient DEFAULT_CONFIG_PATH = 'vault/cooker' DEFAULT_CONFIG = { 'storage': ('dict', { 'cls': 'remote', 'args': { 'url': 'http://localhost:5002/', }, }), 'vault_url': ('str', 'http://localhost:5005/'), 'max_bundle_size': ('int', 2 ** 29), # 512 MiB } class PolicyError(Exception): """Raised when the bundle violates the cooking policy.""" pass class BundleTooLargeError(PolicyError): """Raised when the bundle is too large to be cooked.""" pass class BytesIOBundleSizeLimit(io.BytesIO): def __init__(self, *args, size_limit=None, **kwargs): - super().__init__(self, *args, **kwargs) + super().__init__(*args, **kwargs) self.size_limit = size_limit def write(self, chunk): if ((self.size_limit is not None and self.getbuffer().nbytes + len(chunk) > self.size_limit)): raise BundleTooLargeError( "The requested bundle exceeds the maximum allowed " "size of {} bytes.".format(self.size_limit)) return super().write(chunk) class BaseVaultCooker(metaclass=abc.ABCMeta): """Abstract base class for the vault's bundle creators This class describes a common API for the cookers. To define a new cooker, inherit from this class and override: - CACHE_TYPE_KEY: key to use for the bundle to reference in cache - def cook(): cook the object into a bundle """ CACHE_TYPE_KEY = None def __init__(self, obj_type, obj_id): """Initialize the cooker. The type of the object represented by the id depends on the concrete class. Very likely, each type of bundle will have its own cooker class. Args: storage: the storage object cache: the cache where to store the bundle obj_id: id of the object to be cooked into a bundle. """ self.config = config.load_named_config(DEFAULT_CONFIG_PATH, DEFAULT_CONFIG) self.obj_type = obj_type self.obj_id = hashutil.hash_to_bytes(obj_id) self.backend = RemoteVaultClient(self.config['vault_url']) self.storage = get_storage(**self.config['storage']) self.max_bundle_size = self.config['max_bundle_size'] @abc.abstractmethod def check_exists(self): """Checks that the requested object exists and can be cooked. Override this in the cooker implementation. """ raise NotImplemented @abc.abstractmethod def prepare_bundle(self): """Implementation of the cooker. Yields chunks of the bundle bytes. Override this with the cooker implementation. """ raise NotImplemented def write(self, chunk): self.fileobj.write(chunk) def cook(self): """Cook the requested object into a bundle """ self.backend.set_status(self.obj_type, self.obj_id, 'pending') self.backend.set_progress(self.obj_type, self.obj_id, 'Processing...') self.fileobj = BytesIOBundleSizeLimit(size_limit=self.max_bundle_size) try: self.prepare_bundle() bundle = self.fileobj.getvalue() except PolicyError as e: self.backend.set_status(self.obj_type, self.obj_id, 'failed') self.backend.set_progress(self.obj_type, self.obj_id, str(e)) except Exception as e: self.backend.set_status(self.obj_type, self.obj_id, 'failed') self.backend.set_progress( self.obj_type, self.obj_id, "Internal Server Error. This incident will be reported.") logging.exception("Bundle cooking failed.") else: # TODO: use proper content streaming instead of put_bundle() self.backend.put_bundle(self.CACHE_TYPE_KEY, self.obj_id, bundle) self.backend.set_status(self.obj_type, self.obj_id, 'done') self.backend.set_progress(self.obj_type, self.obj_id, None) finally: self.backend.send_notif(self.obj_type, self.obj_id)