quay/storage/distributedstorage.py

import random
import logging

from functools import wraps

from storage.basestorage import StoragePaths, BaseStorage


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)