api: make .open() stream data for Azure

This employs generating temporary signed public url and then streaming
from it. The bulk of the code might be reused for other remotes.

Url streaming does resume on connection error using "Range" header and
regenerates public url too.

If remote class doesn't provide .open() then this falls back to pulling
a file and opening it locally.

dvc/remote/azure.py

@@ -4,11 +4,12 @@ from __future__ import unicode_literals
 import os
 import re
 import logging
+from datetime import datetime, timedelta
 
 from dvc.scheme import Schemes
 
 try:
-    from azure.storage.blob import BlockBlobService
+    from azure.storage.blob import BlockBlobService, BlobPermissions
     from azure.common import AzureMissingResourceHttpError
 except ImportError:
     BlockBlobService = None
@@ -18,6 +19,7 @@ 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__)
@@ -134,3 +136,21 @@ class RemoteAZURE(RemoteBASE):
         self.blob_service.get_blob_to_path(
             from_info.bucket, from_info.path, to_file, progress_callback=cb
         )
+
+    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)
+
+        sas_token = self.blob_service.generate_blob_shared_access_signature(
+            path_info.bucket,
+            path_info.path,
+            permission=BlobPermissions.READ,
+            expiry=expires_at,
+        )
+        download_url = self.blob_service.make_blob_url(
+            path_info.bucket, path_info.path, sas_token=sas_token
+        )
+        return download_url

dvc/remote/base.py

@@ -498,6 +498,9 @@ class RemoteBASE(object):
 
         return 0
 
+    def open(self, path_info, mode="r", encoding=None):
+        raise RemoteActionNotImplemented("open", self.scheme)
+
     def remove(self, path_info):
         raise RemoteActionNotImplemented("remove", self.scheme)
 

dvc/repo/__init__.py

@@ -16,6 +16,7 @@ from dvc.exceptions import (
 )
 from dvc.ignore import DvcIgnoreFilter
 from dvc.path_info import PathInfo
+from dvc.remote.base import RemoteActionNotImplemented
 from dvc.utils.compat import open as _open, fspath_py35
 from dvc.utils import relpath
 
@@ -464,15 +465,23 @@ class Repo(object):
         cache_file = self.cache.local.checksum_to_path_info(out.checksum)
         cache_file = fspath_py35(cache_file)
 
-        with self.state:
-            cache_info = out.get_used_cache(remote=remote)
-            self.cloud.pull(cache_info, remote=remote)
+        if os.path.exists(cache_file):
+            return _open(cache_file, mode=mode, encoding=encoding)
 
-        # Since pull may just skip with a warning, we need to check it here
-        if not os.path.exists(cache_file):
-            raise OutputFileMissingError(relpath(path, self.root_dir))
+        try:
+            remote_obj = self.cloud.get_remote(remote)
+            remote_info = remote_obj.checksum_to_path_info(out.checksum)
+            return remote_obj.open(remote_info, mode=mode, encoding=encoding)
+        except RemoteActionNotImplemented:
+            with self.state:
+                cache_info = out.get_used_cache(remote=remote)
+                self.cloud.pull(cache_info, remote=remote)
 
-        return _open(cache_file, mode=mode, encoding=encoding)
+            # Since pull may just skip with a warning, we need to check it here
+            if not os.path.exists(cache_file):
+                raise OutputFileMissingError(relpath(path, self.root_dir))
+
+            return _open(cache_file, mode=mode, encoding=encoding)
 
     @cached_property
     def dvcignore(self):

dvc/utils/http.py

@@ -0,0 +1,106 @@
+from contextlib import contextmanager
+import io
+
+import requests
+
+
+@contextmanager
+def open_url(url, mode="r", encoding=None):
+    """Opens an url as a readable stream.
+
+    Resumes on connection error.
+    Url could be a string or a callable returning a string.
+    """
+    assert mode in {"r", "rt", "rb"}
+
+    with iter_url(url) as (response, it):
+        bytes_stream = IterStream(it)
+
+        if mode == "rb":
+            yield bytes_stream
+        else:
+            encoding = encoding or response.encoding
+            yield io.TextIOWrapper(bytes_stream, encoding=encoding)
+
+
+@contextmanager
+def iter_url(url, chunk_size=io.DEFAULT_BUFFER_SIZE):
+    """Iterate over chunks requested from url."""
+
+    def request(headers=None):
+        the_url = url() if callable(url) else url
+        response = requests.get(the_url, stream=True, headers=headers)
+        response.raise_for_status()
+        return response
+
+    def gen(response):
+        try:
+            pos = 0
+            while True:
+                try:
+                    for chunk in response.iter_content(chunk_size):
+                        pos += len(chunk)
+                        yield chunk
+                    break
+                except requests.ConnectionError:
+                    response.close()
+                    if response.headers.get("Accept-Ranges") != "bytes":
+                        raise
+
+                    # Reopen request from where we stopped
+                    headers = {"Range": "bytes={}-".format(pos)}
+                    response = request(headers)
+        finally:
+            response.close()
+
+    response = request()
+    it = gen(response)
+    try:
+        yield response, it
+    finally:
+        # Ensure connection is closed
+        it.close()
+
+
+class IterStream(io.RawIOBase):
+    """Wraps an iterator yielding bytes as a file object"""
+
+    def __init__(self, iterator):
+        self.iterator = iterator
+        self.leftover = None
+
+    def readable(self):
+        return True
+
+    # Python 3 requires only .readinto() method, it still uses other ones
+    # under some circumstances and falls back if those are absent. Since
+    # iterator already constructs byte strings for us, .readinto() is not the
+    # most optimal, so we provide .read1() too.
+
+    def readinto(self, b):
+        try:
+            n = len(b)  # We're supposed to return at most this much
+            chunk = self.leftover or next(self.iterator)
+            output, self.leftover = chunk[:n], chunk[n:]
+
+            n_out = len(output)
+            b[:n_out] = output
+            return n_out
+        except StopIteration:
+            return 0  # indicate EOF
+
+    readinto1 = readinto
+
+    def read1(self, n=-1):
+        try:
+            chunk = self.leftover or next(self.iterator)
+        except StopIteration:
+            return b""
+
+        # Return an arbitrary number or bytes
+        if n <= 0:
+            self.leftover = None
+            return chunk
+
+        output, self.leftover = chunk[:n], chunk[n:]
+        return output

setup.py

@@ -107,6 +107,7 @@ tests_requirements = [
     "jaraco.windows==3.9.2",
     "mock-ssh-server>=0.5.0",
     "moto>=1.3.14.dev55",
+    "rangehttpserver==1.2.0",
 ]
 
 if (sys.version_info) >= (3, 6):

tests/unit/utils/test_http.py

@@ -0,0 +1,37 @@
+from __future__ import unicode_literals
+
+import io
+import requests
+
+from dvc.utils.http import open_url
+from tests.utils.httpd import StaticFileServer
+
+
+def test_open_url(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+
+    # Simulate bad connection
+    original_iter_content = requests.Response.iter_content
+
+    def bad_iter_content(self, *args, **kwargs):
+        it = original_iter_content(self, *args, **kwargs)
+        for i, chunk in enumerate(it):
+            # Drop connection error on second chunk if there is one
+            if i > 0:
+                raise requests.ConnectionError("Simulated connection drop")
+            yield chunk
+
+    monkeypatch.setattr(requests.Response, "iter_content", bad_iter_content)
+
+    # Text should be longer than default chunk to test resume,
+    # using twice of that plus something tests second resume,
+    # this is important because second response is different
+    text = "0123456789" * (io.DEFAULT_BUFFER_SIZE // 10 + 1)
+    (tmp_path / "sample.txt").write_text(text * 2)
+
+    with StaticFileServer() as httpd:
+        url = "http://localhost:{}/sample.txt".format(httpd.server_port)
+        with open_url(url) as fd:
+            # Test various .read() variants
+            assert fd.read(len(text)) == text
+            assert fd.read() == text

tests/utils/httpd.py

@@ -2,31 +2,39 @@ import hashlib
 import os
 import threading
 
-from dvc.utils.compat import HTTPServer, SimpleHTTPRequestHandler
+from RangeHTTPServer import RangeRequestHandler
+
+from dvc.utils.compat import HTTPServer
 
 
-class ETagHandler(SimpleHTTPRequestHandler):
+class TestRequestHandler(RangeRequestHandler):
+    checksum_header = None
+
     def end_headers(self):
-        file = self.translate_path(self.path)
+        # RangeRequestHandler only sends Accept-Ranges header if Range header
+        # is present, see https://github.com/danvk/RangeHTTPServer/issues/23
+        if not self.headers.get("Range"):
+            self.send_header("Accept-Ranges", "bytes")
 
-        if not os.path.isdir(file) and os.path.exists(file):
-            with open(file, "r") as fd:
-                etag = hashlib.md5(fd.read().encode("utf8")).hexdigest()
-                self.send_header("ETag", etag)
+        # Add a checksum header
+        if self.checksum_header:
+            file = self.translate_path(self.path)
 
-        SimpleHTTPRequestHandler.end_headers(self)
+            if not os.path.isdir(file) and os.path.exists(file):
+                with open(file, "r") as fd:
+                    encoded_text = fd.read().encode("utf8")
+                    checksum = hashlib.md5(encoded_text).hexdigest()
+                    self.send_header(self.checksum_header, checksum)
+
+        RangeRequestHandler.end_headers(self)
 
 
-class ContentMD5Handler(SimpleHTTPRequestHandler):
-    def end_headers(self):
-        file = self.translate_path(self.path)
+class ETagHandler(TestRequestHandler):
+    checksum_header = "ETag"
 
-        if not os.path.isdir(file) and os.path.exists(file):
-            with open(file, "r") as fd:
-                md5 = hashlib.md5(fd.read().encode("utf8")).hexdigest()
-                self.send_header("Content-MD5", md5)
 
-        SimpleHTTPRequestHandler.end_headers(self)
+class ContentMD5Handler(TestRequestHandler):
+    checksum_header = "Content-MD5"
 
 
 class StaticFileServer: