From e5e29065c80e125d362ececaf1debca3d1c56aa7 Mon Sep 17 00:00:00 2001 From: Alexander Schepanovski Date: Tue, 6 Aug 2019 11:10:58 +0400 Subject: [PATCH] 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 | 22 ++++++- dvc/remote/base.py | 3 + dvc/repo/__init__.py | 23 +++++--- dvc/utils/http.py | 106 ++++++++++++++++++++++++++++++++++ setup.py | 1 + tests/unit/utils/test_http.py | 37 ++++++++++++ tests/utils/httpd.py | 40 ++++++++----- 7 files changed, 208 insertions(+), 24 deletions(-) create mode 100644 dvc/utils/http.py create mode 100644 tests/unit/utils/test_http.py 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: