object_store.py

""" Store packages in S3 """
import logging
from binascii import hexlify
from hashlib import md5

from pyramid.httpexceptions import HTTPFound
from pyramid.settings import asbool
from smart_open import open as _open

from pypicloud.models import Package

from .base import IStorage

LOG = logging.getLogger(__name__)


class ObjectStoreStorage(IStorage):

    """Storage backend base class containing code that is common between
    supported object stores (S3 / GCS)
    """

    test = False

    def __init__(
        self,
        request=None,
        expire_after=None,
        bucket_prefix=None,
        prepend_hash=None,
        redirect_urls=None,
        sse=None,
        object_acl=None,
        storage_class=None,
        region_name=None,
        public_url=False,
        **kwargs
    ):
        super(ObjectStoreStorage, self).__init__(request, **kwargs)
        self.expire_after = expire_after
        self.bucket_prefix = bucket_prefix
        self.prepend_hash = prepend_hash
        self.redirect_urls = redirect_urls
        self.sse = sse
        self.object_acl = object_acl
        self.storage_class = storage_class
        self.region_name = region_name
        self.public_url = public_url

    def _generate_url(self, package: Package) -> str:
        """Subclasses must implement a method for generating signed URLs to
        the package in the object store
        """
        raise NotImplementedError

    @classmethod
    def package_from_object(cls, obj, factory):
        """Subclasses must implement a method for constructing a Package
        instance from the backend's storage object format
        """
        raise NotImplementedError

    @classmethod
    def _subclass_specific_config(cls, settings, common_config):
        """Method to allow subclasses to extract configuration parameters
        specific to them and not covered in the common configuration
        in this class.
        """
        return {}

    @classmethod
    def configure(cls, settings):
        kwargs = super(ObjectStoreStorage, cls).configure(settings)
        kwargs["expire_after"] = int(settings.get("storage.expire_after", 60 * 60 * 24))
        kwargs["bucket_prefix"] = settings.get("storage.prefix", "")
        kwargs["prepend_hash"] = asbool(settings.get("storage.prepend_hash", True))
        kwargs["object_acl"] = settings.get("storage.object_acl", None)
        kwargs["storage_class"] = settings.get("storage.storage_class")
        kwargs["redirect_urls"] = asbool(settings.get("storage.redirect_urls", True))
        kwargs["region_name"] = settings.get("storage.region_name")
        kwargs["public_url"] = asbool(settings.get("storage.public_url"))

        kwargs.update(cls._subclass_specific_config(settings, kwargs))
        return kwargs

    def calculate_path(self, package):
        """Calculates the path of a package"""
        path = package.name + "/" + package.filename
        if self.prepend_hash:
            m = md5()
            m.update(package.filename.encode("utf-8"))
            prefix = hexlify(m.digest()).decode("utf-8")[:4]
            path = prefix + "/" + path
        return path

    def get_path(self, package):
        """Get the fully-qualified bucket path for a package"""
        if "path" not in package.data:
            filename = self.calculate_path(package)
            package.data["path"] = self.bucket_prefix + filename
        return package.data["path"]

    def get_url(self, package):
        if self.redirect_urls:
            return super(ObjectStoreStorage, self).get_url(package)
        else:
            return self._generate_url(package)

    def download_response(self, package):
        return HTTPFound(location=self._generate_url(package))

    def open(self, package):
        url = self._generate_url(package)
        return _open(url, "rb", compression="disable")