From 641eb57b27ccc544266edd07997bd41f12bfa225 Mon Sep 17 00:00:00 2001
From: Alexander Schepanovski <suor.web@gmail.com>
Date: Mon, 12 Aug 2019 20:21:51 +0700
Subject: [PATCH] api: stream data in api.open() for S3, GS and OSS

Employ the same sign url and use resumable download approach tested for
Azure.
---
 dvc/remote/azure.py | 5 -----
 dvc/remote/base.py  | 6 ++++++
 dvc/remote/gs.py    | 8 ++++++++
 dvc/remote/oss.py   | 5 +++++
 dvc/remote/s3.py    | 6 ++++++
 5 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py
index 6e37545ed..aadc86a48 100644
--- a/dvc/remote/azure.py
+++ b/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)
 
diff --git a/dvc/remote/base.py b/dvc/remote/base.py
index 20f8b7f16..050c16659 100644
--- a/dvc/remote/base.py
+++ b/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):
diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py
index 4d07226cd..e64a41e42 100644
--- a/dvc/remote/gs.py
+++ b/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)
diff --git a/dvc/remote/oss.py b/dvc/remote/oss.py
index cbf1fe3a7..f90fa6100 100644
--- a/dvc/remote/oss.py
+++ b/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)
diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py
index 757db8980..3d79f3eed 100644
--- a/dvc/remote/s3.py
+++ b/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)
+        )