Merge pull request #2389 from Suor/streaming-clouds

api: stream data in api.open() for S3, GS and OSS

dvc/remote/azure.py

@@ -21,7 +21,6 @@ from dvc.progress import progress
 from dvc.config import Config
 from dvc.remote.base import RemoteBASE
 from dvc.path_info import CloudURLInfo
-from dvc.utils.http import open_url
 
 
 logger = logging.getLogger(__name__)
@@ -141,10 +140,6 @@ class RemoteAZURE(RemoteBASE):
         paths = self._list_paths(path_info.bucket, path_info.path)
         return any(path_info.path == path for path in paths)
 
-    def open(self, path_info, mode="r", encoding=None):
-        get_url = lambda: self._generate_download_url(path_info)  # noqa: E731
-        return open_url(get_url, mode=mode, encoding=encoding)
-
     def _generate_download_url(self, path_info, expires=3600):
         expires_at = datetime.utcnow() + timedelta(seconds=expires)
 

dvc/remote/base.py

@@ -8,6 +8,7 @@ import json
 import logging
 import tempfile
 import itertools
+from functools import partial
 from operator import itemgetter
 from multiprocessing import cpu_count
 from concurrent.futures import ThreadPoolExecutor
@@ -23,6 +24,7 @@ from dvc.progress import progress, ProgressCallback
 from dvc.utils import LARGE_DIR_SIZE, tmp_fname, move, relpath, makedirs
 from dvc.state import StateNoop
 from dvc.path_info import PathInfo, URLInfo
+from dvc.utils.http import open_url
 
 
 logger = logging.getLogger(__name__)
@@ -508,6 +510,10 @@ class RemoteBASE(object):
         return 0
 
     def open(self, path_info, mode="r", encoding=None):
+        if hasattr(self, "_generate_download_url"):
+            get_url = partial(self._generate_download_url, path_info)
+            return open_url(get_url, mode=mode, encoding=encoding)
+
         raise RemoteActionNotImplemented("open", self.scheme)
 
     def remove(self, path_info):

dvc/remote/gs.py

@@ -1,6 +1,7 @@
 from __future__ import unicode_literals
 
 import logging
+from datetime import timedelta
 from funcy import cached_property
 
 try:
@@ -97,3 +98,10 @@ class RemoteGS(RemoteBASE):
         bucket = self.gs.bucket(from_info.bucket)
         blob = bucket.get_blob(from_info.path)
         blob.download_to_filename(to_file)
+
+    def _generate_download_url(self, path_info, expires=3600):
+        expiration = timedelta(seconds=int(expires))
+
+        bucket = self.gs.bucket(path_info.bucket)
+        blob = bucket.get_blob(path_info.path)
+        return blob.generate_signed_url(expiration=expiration)

dvc/remote/oss.py

@@ -120,6 +120,11 @@ class RemoteOSS(RemoteBASE):
             from_info.path, to_file, progress_callback=cb
         )
 
+    def _generate_download_url(self, path_info, expires=3600):
+        assert path_info.bucket == self.path_info.bucket
+
+        return self.oss_service.sign_url("GET", path_info.path, expires)
+
     def exists(self, path_info):
         paths = self._list_paths(path_info.path)
         return any(path_info.path == path for path in paths)

dvc/remote/s3.py

@@ -238,3 +238,9 @@ class RemoteS3(RemoteBASE):
         self.s3.download_file(
             from_info.bucket, from_info.path, to_file, Callback=cb
         )
+
+    def _generate_download_url(self, path_info, expires=3600):
+        params = {"Bucket": path_info.bucket, "Key": path_info.path}
+        return self.s3.generate_presigned_url(
+            ClientMethod="get_object", Params=params, ExpiresIn=int(expires)
+        )