initial import for Open Source 🎉

storage/distributedstorage.py

@@ -0,0 +1,83 @@
+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)