Merge pull request #2371 from Suor/streaming

api: make .open() stream data for Azure

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

@@ -515,6 +515,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: