1
0
mirror of https://github.com/quay/quay.git synced 2026-01-26 06:21:37 +03:00
Files
quay/storage/distributedstorage.py
2020-02-05 19:55:07 -08:00

100 lines
3.5 KiB
Python

import random
import logging
from functools import wraps
from storage.basestorage import StoragePaths, BaseStorage, BaseStorageV2
logger = logging.getLogger(__name__)
def _location_aware(unbound_func, requires_write=False):
@wraps(unbound_func)
def wrapper(self, locations, *args, **kwargs):
if requires_write:
assert not self.readonly_mode
storage = None
for preferred in self.preferred_locations:
if preferred in locations:
storage = self._storages[preferred]
break
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=None,
default_locations=None,
proxy=None,
readonly_mode=False,
):
self._storages = dict(storages)
self.preferred_locations = list(preferred_locations or [])
self.default_locations = list(default_locations or [])
self.proxy = proxy
self.readonly_mode = readonly_mode
@property
def locations(self):
"""
Returns the names of the locations supported.
"""
return list(self._storages.keys())
_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, requires_write=True)
stream_read = _location_aware(BaseStorage.stream_read)
stream_read_file = _location_aware(BaseStorage.stream_read_file)
stream_write = _location_aware(BaseStorage.stream_write, requires_write=True)
exists = _location_aware(BaseStorage.exists)
remove = _location_aware(BaseStorage.remove, requires_write=True)
validate = _location_aware(BaseStorage.validate, requires_write=True)
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, requires_write=True
)
stream_upload_chunk = _location_aware(BaseStorageV2.stream_upload_chunk, requires_write=True)
complete_chunked_upload = _location_aware(
BaseStorageV2.complete_chunked_upload, requires_write=True
)
cancel_chunked_upload = _location_aware(
BaseStorageV2.cancel_chunked_upload, requires_write=True
)
def get_direct_download_url(
self, locations, path, request_ip=None, expires_in=600, requires_cors=False, head=False
):
download_url = self._get_direct_download_url(
locations, path, request_ip, expires_in, requires_cors, head
)
if download_url is None:
return None
if self.proxy is None:
return download_url
return self.proxy.proxy_download_url(download_url)
def copy_between(self, path, source_location, destination_location):
"""
Copies a file between the source location and the destination location.
"""
source_storage = self._storages[source_location]
destination_storage = self._storages[destination_location]
source_storage.copy_to(destination_storage, path)