import random
import logging
from functools import wraps
from storage.basestorage import StoragePaths, BaseStorage, BaseStorageV2
logger = logging.getLogger(__name__)
def _location_aware(unbound_func):
@wraps(unbound_func)
def wrapper(self, locations, *args, **kwargs):
storage = None
for preferred in self.preferred_locations:
if preferred in locations:
storage = self._storages[preferred]
if not storage:
storage = self._storages[random.sample(locations, 1)[0]]
storage_func = getattr(storage, unbound_func.__name__)
return storage_func(*args, **kwargs)
return wrapper
class DistributedStorage(StoragePaths):
def __init__(self, storages, preferred_locations=[]):
self._storages = dict(storages)
self.preferred_locations = list(preferred_locations)
get_direct_download_url = _location_aware(BaseStorage.get_direct_download_url)
get_direct_upload_url = _location_aware(BaseStorage.get_direct_upload_url)
get_content = _location_aware(BaseStorage.get_content)
put_content = _location_aware(BaseStorage.put_content)
stream_read = _location_aware(BaseStorage.stream_read)
stream_read_file = _location_aware(BaseStorage.stream_read_file)
stream_write = _location_aware(BaseStorage.stream_write)
list_directory = _location_aware(BaseStorage.list_directory)
exists = _location_aware(BaseStorage.exists)
remove = _location_aware(BaseStorage.remove)
get_checksum = _location_aware(BaseStorage.get_checksum)
get_supports_resumable_downloads = _location_aware(BaseStorage.get_supports_resumable_downloads)
initiate_chunked_upload = _location_aware(BaseStorageV2.initiate_chunked_upload)
stream_upload_chunk = _location_aware(BaseStorageV2.stream_upload_chunk)
complete_chunked_upload = _location_aware(BaseStorageV2.complete_chunked_upload)
cancel_chunked_upload = _location_aware(BaseStorageV2.cancel_chunked_upload)
