diff --git a/pillar/api/utils/storage.py b/pillar/api/utils/storage.py index 4bef8979..9e251842 100644 --- a/pillar/api/utils/storage.py +++ b/pillar/api/utils/storage.py @@ -1,5 +1,6 @@ """Utility for managing storage backends and files.""" +import abc import logging import os import shutil @@ -8,6 +9,18 @@ from flask import current_app log = logging.getLogger(__name__) +# Mapping from backend name to backend class +backends = {} + + +def register_backend(backend_name): + def wrapper(cls): + assert backend_name not in backends + backends[backend_name] = cls + return cls + + return wrapper + class StorageBackend(object): """Can be a GCS bucket or simply a project folder in Pillar @@ -17,9 +30,16 @@ class StorageBackend(object): """ + __metaclass__ = abc.ABCMeta + def __init__(self, backend): self.backend = backend + @abc.abstractmethod + def upload_file(self, param1, param2, param3): + """docstuff""" + pass + class FileInStorage(object): """A wrapper for file or blob objects. @@ -35,6 +55,7 @@ class FileInStorage(object): self.size = None +@register_backend('local') class PillarStorage(StorageBackend): def __init__(self, project_id): super(PillarStorage, self).__init__(backend='local') @@ -58,3 +79,10 @@ class PillarStorageFile(FileInStorage): shutil.copyfileobj(uploaded_file, outfile) self.size = file_size + + +def default_storage_backend(): + from flask import current_app + + backend_cls = backends[current_app.config['STORAGE_BACKEND']] + return backend_cls()