diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index 316211d4c..0b8330ae3 100644 --- a/dvc/remote/azure.py +++ b/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 diff --git a/dvc/remote/base.py b/dvc/remote/base.py index edacef724..b2caf7a91 100644 --- a/dvc/remote/base.py +++ b/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) diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 2efdaa4da..9b842b5f5 100644 --- a/dvc/repo/__init__.py +++ b/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): diff --git a/dvc/utils/http.py b/dvc/utils/http.py new file mode 100644 index 000000000..b6c0af596 --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index 05ff967f3..3313b0e8d 100644 --- a/setup.py +++ b/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): diff --git a/tests/unit/utils/test_http.py b/tests/unit/utils/test_http.py new file mode 100644 index 000000000..44e88f108 --- /dev/null +++ b/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 diff --git a/tests/utils/httpd.py b/tests/utils/httpd.py index 7369d90d3..6e90598c3 100644 --- a/tests/utils/httpd.py +++ b/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: