From 8da0390c57267710db1d08488a45f02338c9081b Mon Sep 17 00:00:00 2001 From: Alexander Schepanovski Date: Fri, 22 Nov 2019 12:58:08 +0100 Subject: [PATCH 01/16] remote: small .save_info()/.get_checksum() cleanup --- dvc/output/base.py | 7 +------ dvc/remote/base.py | 5 +++-- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/dvc/output/base.py b/dvc/output/base.py index 120489d19..8344f1f3d 100644 --- a/dvc/output/base.py +++ b/dvc/output/base.py @@ -170,12 +170,7 @@ class OutputBase(object): return self.remote.exists(self.path_info) def changed_checksum(self): - return ( - self.checksum - != self.remote.save_info(self.path_info)[ - self.remote.PARAM_CHECKSUM - ] - ) + return self.checksum != self.remote.get_checksum(self.path_info) def changed_cache(self): if not self.use_cache or not self.checksum: diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 011b0af4a..10969250a 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -293,6 +293,8 @@ class RemoteBASE(object): return checksum.endswith(cls.CHECKSUM_DIR_SUFFIX) def get_checksum(self, path_info): + assert path_info.scheme == self.scheme + if not self.exists(path_info): return None @@ -322,7 +324,6 @@ class RemoteBASE(object): return checksum def save_info(self, path_info): - assert path_info.scheme == self.scheme return {self.PARAM_CHECKSUM: self.get_checksum(path_info)} def changed(self, path_info, checksum_info): @@ -363,7 +364,7 @@ class RemoteBASE(object): ) return True - actual = self.save_info(path_info)[self.PARAM_CHECKSUM] + actual = self.get_checksum(path_info) if checksum != actual: logger.debug( "checksum '{}'(actual '{}') for '{}' has changed.".format( From 8681cc3cd1b50c13ea6c90f7eca5c7f9788d78c7 Mon Sep 17 00:00:00 2001 From: Alexander Schepanovski Date: Fri, 22 Nov 2019 14:39:56 +0100 Subject: [PATCH 02/16] perf: optimize cache listing for local, ssh and hdfs - make all three lazy - simplify local one - use deque instead of list for hdfs --- dvc/remote/hdfs.py | 33 +++++++++++++-------------------- dvc/remote/local.py | 12 +----------- dvc/remote/ssh/__init__.py | 4 +++- 3 files changed, 17 insertions(+), 32 deletions(-) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index 75469ed18..0301fce3c 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -1,24 +1,20 @@ from __future__ import unicode_literals +from dvc.utils.compat import FileNotFoundError, urlparse import io import logging import os import posixpath import re -from contextlib import closing -from contextlib import contextmanager -from subprocess import PIPE -from subprocess import Popen +from collections import deque +from contextlib import closing, contextmanager +import subprocess -from .base import RemoteBASE -from .base import RemoteCmdError +from .base import RemoteBASE, RemoteCmdError from .pool import get_connection from dvc.config import Config from dvc.scheme import Schemes -from dvc.utils import fix_env -from dvc.utils import tmp_fname -from dvc.utils.compat import FileNotFoundError -from dvc.utils.compat import urlparse +from dvc.utils import fix_env, tmp_fname logger = logging.getLogger(__name__) @@ -68,15 +64,15 @@ class RemoteHDFS(RemoteBASE): close_fds = os.name != "nt" executable = os.getenv("SHELL") if os.name != "nt" else None - p = Popen( + p = subprocess.Popen( cmd, shell=True, close_fds=close_fds, executable=executable, env=fix_env(os.environ), - stdin=PIPE, - stdout=PIPE, - stderr=PIPE, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, ) out, err = p.communicate() if p.returncode != 0: @@ -156,10 +152,9 @@ class RemoteHDFS(RemoteBASE): def list_cache_paths(self): if not self.exists(self.path_info): - return [] + return - files = [] - dirs = [self.path_info.path] + dirs = deque([self.path_info.path]) with self.hdfs(self.path_info) as hdfs: while dirs: @@ -167,6 +162,4 @@ class RemoteHDFS(RemoteBASE): if entry["kind"] == "directory": dirs.append(urlparse(entry["name"]).path) elif entry["kind"] == "file": - files.append(urlparse(entry["name"]).path) - - return files + yield urlparse(entry["name"]).path diff --git a/dvc/remote/local.py b/dvc/remote/local.py index bd20e91b5..6e7900fd5 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -85,17 +85,7 @@ class RemoteLOCAL(RemoteBASE): def list_cache_paths(self): assert self.path_info is not None - - clist = [] - for entry in os.listdir(fspath_py35(self.path_info)): - subdir = self.path_info / entry - if not os.path.isdir(fspath_py35(subdir)): - continue - clist.extend( - subdir / cache for cache in os.listdir(fspath_py35(subdir)) - ) - - return clist + return walk_files(self.path_info, None) def get(self, md5): if not md5: diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index e8469f968..589bea8d1 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -264,7 +264,9 @@ class RemoteSSH(RemoteBASE): def list_cache_paths(self): with self.ssh(self.path_info) as ssh: - return list(ssh.walk_files(self.path_info.path)) + # If we simply return an iterator then with above closes instantly + for path in ssh.walk_files(self.path_info.path): + yield path def walk_files(self, path_info): with self.ssh(path_info) as ssh: From 31aed9b27b8b14243ce2e8b3d33c3d7ab865e964 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= Date: Tue, 26 Nov 2019 14:33:48 +0100 Subject: [PATCH 03/16] NoRemoteInExternalRepoError: dont pass cause of exception --- dvc/exceptions.py | 9 ++++----- dvc/external_repo.py | 6 +++--- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/dvc/exceptions.py b/dvc/exceptions.py index 92f46a882..1a167117a 100644 --- a/dvc/exceptions.py +++ b/dvc/exceptions.py @@ -332,13 +332,12 @@ class CollectCacheError(DvcException): pass -class RemoteNotSpecifiedInExternalRepoError(DvcException): - def __init__(self, url, cause=None): - super(RemoteNotSpecifiedInExternalRepoError, self).__init__( +class NoRemoteInExternalRepoError(DvcException): + def __init__(self, url): + super(NoRemoteInExternalRepoError, self).__init__( "No DVC remote is specified in the target repository '{}'".format( url - ), - cause=cause, + ) ) diff --git a/dvc/external_repo.py b/dvc/external_repo.py index 0555f2d75..9d9096300 100644 --- a/dvc/external_repo.py +++ b/dvc/external_repo.py @@ -8,7 +8,7 @@ from distutils.dir_util import copy_tree from funcy import retry from dvc.config import NoRemoteError -from dvc.exceptions import RemoteNotSpecifiedInExternalRepoError +from dvc.exceptions import NoRemoteInExternalRepoError from dvc.exceptions import NoOutputInExternalRepoError from dvc.exceptions import OutputNotFoundError from dvc.utils.fs import remove @@ -25,8 +25,8 @@ def external_repo(url=None, rev=None, rev_lock=None, cache_dir=None): repo = Repo(path) try: yield repo - except NoRemoteError as exc: - raise RemoteNotSpecifiedInExternalRepoError(url, cause=exc) + except NoRemoteError: + raise NoRemoteInExternalRepoError(url) except OutputNotFoundError as exc: if exc.repo is repo: raise NoOutputInExternalRepoError(exc.output, repo.root_dir, url) From 78c505f6cb7a98962ad5fecb3b773be5b2cc39d3 Mon Sep 17 00:00:00 2001 From: Dani Hodovic Date: Thu, 28 Nov 2019 15:53:52 +0100 Subject: [PATCH 04/16] test: refactor & remove redundant test fixtures https://github.com/iterative/dvc/pull/2837#discussion_r351318487 --- tests/func/test_api.py | 6 +++--- tests/func/test_get.py | 16 ++++++++-------- tests/func/test_import.py | 16 ++++++++-------- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/tests/func/test_api.py b/tests/func/test_api.py index 3ba71ea46..ddf101479 100644 --- a/tests/func/test_api.py +++ b/tests/func/test_api.py @@ -84,13 +84,13 @@ def test_get_url(repo_dir, dvc_repo, remote_url): @pytest.mark.parametrize("remote_url", remote_params, indirect=True) -def test_get_url_external(repo_dir, dvc_repo, erepo, remote_url): +def test_get_url_external(dvc_repo, erepo, remote_url): _set_remote_url_and_commit(erepo.dvc, remote_url) # Using file url to force clone to tmp repo repo_url = "file://" + erepo.dvc.root_dir expected_url = URLInfo(remote_url) / "ac/bd18db4cc2f85cedef654fccc4a4d8" - assert api.get_url(repo_dir.FOO, repo=repo_url) == expected_url + assert api.get_url(erepo.FOO, repo=repo_url) == expected_url @pytest.mark.parametrize("remote_url", all_remote_params, indirect=True) @@ -107,7 +107,7 @@ def test_open(repo_dir, dvc_repo, remote_url): @pytest.mark.parametrize("remote_url", all_remote_params, indirect=True) -def test_open_external(repo_dir, dvc_repo, erepo, remote_url): +def test_open_external(dvc_repo, erepo, remote_url): erepo.dvc.scm.checkout("branch") _set_remote_url_and_commit(erepo.dvc, remote_url) erepo.dvc.scm.checkout("master") diff --git a/tests/func/test_get.py b/tests/func/test_get.py index 18afef5b5..2b65d5136 100644 --- a/tests/func/test_get.py +++ b/tests/func/test_get.py @@ -14,7 +14,7 @@ from dvc.utils import makedirs from tests.utils import trees_equal -def test_get_repo_file(repo_dir, erepo): +def test_get_repo_file(erepo): src = erepo.FOO dst = erepo.FOO + "_imported" @@ -22,10 +22,10 @@ def test_get_repo_file(repo_dir, erepo): assert os.path.exists(dst) assert os.path.isfile(dst) - assert filecmp.cmp(repo_dir.FOO, dst, shallow=False) + assert filecmp.cmp(erepo.FOO, dst, shallow=False) -def test_get_repo_dir(repo_dir, erepo): +def test_get_repo_dir(erepo): src = erepo.DATA_DIR dst = erepo.DATA_DIR + "_imported" @@ -36,7 +36,7 @@ def test_get_repo_dir(repo_dir, erepo): trees_equal(src, dst) -def test_cache_type_is_properly_overridden(repo_dir, erepo): +def test_cache_type_is_properly_overridden(erepo): erepo.dvc.config.set( Config.SECTION_CACHE, Config.SECTION_CACHE_TYPE, "symlink" ) @@ -53,7 +53,7 @@ def test_cache_type_is_properly_overridden(repo_dir, erepo): assert os.path.isfile(dst) -def test_get_repo_rev(repo_dir, erepo): +def test_get_repo_rev(erepo): src = "version" dst = src @@ -70,13 +70,13 @@ def test_get_from_non_dvc_repo(git_erepo): Repo.get(git_erepo.root_dir, "some_file.zip") -def test_get_a_dvc_file(repo_dir, erepo): +def test_get_a_dvc_file(erepo): with pytest.raises(GetDVCFileError): Repo.get(erepo.root_dir, "some_file.dvc") @pytest.mark.parametrize("dname", [".", "dir", "dir/subdir"]) -def test_get_to_dir(dname, repo_dir, erepo): +def test_get_to_dir(dname, erepo): src = erepo.FOO makedirs(dname, exist_ok=True) @@ -86,4 +86,4 @@ def test_get_to_dir(dname, repo_dir, erepo): dst = os.path.join(dname, os.path.basename(src)) assert os.path.isdir(dname) - assert filecmp.cmp(repo_dir.FOO, dst, shallow=False) + assert filecmp.cmp(erepo.FOO, dst, shallow=False) diff --git a/tests/func/test_import.py b/tests/func/test_import.py index 067d5be8e..6c9a956c5 100644 --- a/tests/func/test_import.py +++ b/tests/func/test_import.py @@ -15,7 +15,7 @@ from dvc.utils import makedirs from tests.utils import trees_equal -def test_import(repo_dir, git, dvc_repo, erepo): +def test_import(git, dvc_repo, erepo): src = erepo.FOO dst = erepo.FOO + "_imported" @@ -23,11 +23,11 @@ def test_import(repo_dir, git, dvc_repo, erepo): assert os.path.exists(dst) assert os.path.isfile(dst) - assert filecmp.cmp(repo_dir.FOO, dst, shallow=False) + assert filecmp.cmp(erepo.FOO, dst, shallow=False) assert git.git.check_ignore(dst) -def test_import_dir(repo_dir, git, dvc_repo, erepo): +def test_import_dir(git, dvc_repo, erepo): src = erepo.DATA_DIR dst = erepo.DATA_DIR + "_imported" @@ -39,7 +39,7 @@ def test_import_dir(repo_dir, git, dvc_repo, erepo): assert git.git.check_ignore(dst) -def test_import_rev(repo_dir, git, dvc_repo, erepo): +def test_import_rev(git, dvc_repo, erepo): src = "version" dst = src @@ -70,7 +70,7 @@ def test_pull_imported_stage(dvc_repo, erepo): assert os.path.isfile(dst_cache) -def test_cache_type_is_properly_overridden(repo_dir, git, dvc_repo, erepo): +def test_cache_type_is_properly_overridden(git, dvc_repo, erepo): erepo.dvc.config.set( Config.SECTION_CACHE, Config.SECTION_CACHE_TYPE, "symlink" ) @@ -85,7 +85,7 @@ def test_cache_type_is_properly_overridden(repo_dir, git, dvc_repo, erepo): assert not System.is_symlink(dst) assert os.path.exists(dst) assert os.path.isfile(dst) - assert filecmp.cmp(repo_dir.FOO, dst, shallow=False) + assert filecmp.cmp(erepo.FOO, dst, shallow=False) assert git.git.check_ignore(dst) @@ -125,7 +125,7 @@ def test_download_error_pulling_imported_stage(dvc_repo, erepo): @pytest.mark.parametrize("dname", [".", "dir", "dir/subdir"]) -def test_import_to_dir(dname, repo_dir, dvc_repo, erepo): +def test_import_to_dir(dname, dvc_repo, erepo): src = erepo.FOO makedirs(dname, exist_ok=True) @@ -136,7 +136,7 @@ def test_import_to_dir(dname, repo_dir, dvc_repo, erepo): assert stage.outs[0].fspath == os.path.abspath(dst) assert os.path.isdir(dname) - assert filecmp.cmp(repo_dir.FOO, dst, shallow=False) + assert filecmp.cmp(erepo.FOO, dst, shallow=False) def test_pull_non_workspace(git, dvc_repo, erepo): From 73eac3b6cc0ac8ebd3cca488e1d8d41868591753 Mon Sep 17 00:00:00 2001 From: Alexander Schepanovski Date: Thu, 28 Nov 2019 21:05:20 +0100 Subject: [PATCH 05/16] remote: protect all remote client/session creation code with locks With @cached_property last write wins, so all threads will use the same session. However, several might be created concurrently wasting resources and even potentially asking something from the user. This change guarantees only one client instance will be created for a remote. --- dvc/remote/azure.py | 7 ++++--- dvc/remote/gs.py | 4 +++- dvc/remote/http.py | 4 +++- dvc/remote/oss.py | 42 ++++++++++++++++++++++-------------------- dvc/remote/pool.py | 4 +++- dvc/remote/s3.py | 4 +++- setup.py | 2 +- 7 files changed, 39 insertions(+), 28 deletions(-) diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index 217ce91fc..b991195df 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -4,10 +4,10 @@ from __future__ import unicode_literals import logging import os import re -from datetime import datetime -from datetime import timedelta +from datetime import datetime, timedelta +import threading -from funcy import cached_property +from funcy import cached_property, wrap_prop from dvc.config import Config from dvc.path_info import CloudURLInfo @@ -64,6 +64,7 @@ class RemoteAZURE(RemoteBASE): else self.path_cls.from_parts(scheme=self.scheme, netloc=bucket) ) + @wrap_prop(threading.Lock()) @cached_property def blob_service(self): from azure.storage.blob import BlockBlobService diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index 8ba47c292..6f389ec4f 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -5,8 +5,9 @@ from datetime import timedelta from functools import wraps import io import os.path +import threading -from funcy import cached_property +from funcy import cached_property, wrap_prop from dvc.config import Config from dvc.exceptions import DvcException @@ -91,6 +92,7 @@ class RemoteGS(RemoteBASE): self.projectname = config.get(Config.SECTION_GCP_PROJECTNAME, None) self.credentialpath = config.get(Config.SECTION_GCP_CREDENTIALPATH) + @wrap_prop(threading.Lock()) @cached_property def gs(self): from google.cloud.storage import Client diff --git a/dvc/remote/http.py b/dvc/remote/http.py index 7f96e0b8e..c22416625 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -1,8 +1,9 @@ from __future__ import unicode_literals import logging +import threading -from funcy import cached_property +from funcy import cached_property, wrap_prop from dvc.config import Config from dvc.config import ConfigError @@ -81,6 +82,7 @@ class RemoteHTTP(RemoteBASE): return etag + @wrap_prop(threading.Lock()) @cached_property def _session(self): import requests diff --git a/dvc/remote/oss.py b/dvc/remote/oss.py index ca41e68db..9445fd0b2 100644 --- a/dvc/remote/oss.py +++ b/dvc/remote/oss.py @@ -3,6 +3,9 @@ from __future__ import unicode_literals import logging import os +import threading + +from funcy import cached_property, wrap_prop from dvc.config import Config from dvc.path_info import CloudURLInfo @@ -61,30 +64,29 @@ class RemoteOSS(RemoteBASE): or "defaultSecret" ) - self._bucket = None - - @property + @wrap_prop(threading.Lock()) + @cached_property def oss_service(self): import oss2 - if self._bucket is None: - logger.debug("URL {}".format(self.path_info)) - logger.debug("key id {}".format(self.key_id)) - logger.debug("key secret {}".format(self.key_secret)) - auth = oss2.Auth(self.key_id, self.key_secret) - self._bucket = oss2.Bucket( - auth, self.endpoint, self.path_info.bucket + logger.debug("URL {}".format(self.path_info)) + logger.debug("key id {}".format(self.key_id)) + logger.debug("key secret {}".format(self.key_secret)) + + auth = oss2.Auth(self.key_id, self.key_secret) + bucket = oss2.Bucket(auth, self.endpoint, self.path_info.bucket) + + # Ensure bucket exists + try: + bucket.get_bucket_info() + except oss2.exceptions.NoSuchBucket: + bucket.create_bucket( + oss2.BUCKET_ACL_PUBLIC_READ, + oss2.models.BucketCreateConfig( + oss2.BUCKET_STORAGE_CLASS_STANDARD + ), ) - try: # verify that bucket exists - self._bucket.get_bucket_info() - except oss2.exceptions.NoSuchBucket: - self._bucket.create_bucket( - oss2.BUCKET_ACL_PUBLIC_READ, - oss2.models.BucketCreateConfig( - oss2.BUCKET_STORAGE_CLASS_STANDARD - ), - ) - return self._bucket + return bucket def remove(self, path_info): if path_info.scheme != self.scheme: diff --git a/dvc/remote/pool.py b/dvc/remote/pool.py index bec579cc4..59f16b01b 100644 --- a/dvc/remote/pool.py +++ b/dvc/remote/pool.py @@ -1,7 +1,8 @@ from collections import deque from contextlib import contextmanager +import threading -from funcy import memoize +from funcy import memoize, wrap_with @contextmanager @@ -17,6 +18,7 @@ def get_connection(conn_func, *args, **kwargs): pool.release(conn) +@wrap_with(threading.Lock()) @memoize def get_pool(conn_func, *args, **kwargs): return Pool(conn_func, *args, **kwargs) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 952db8fa4..5594806f0 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -3,8 +3,9 @@ from __future__ import unicode_literals import logging import os +import threading -from funcy import cached_property +from funcy import cached_property, wrap_prop from dvc.config import Config from dvc.exceptions import DvcException @@ -56,6 +57,7 @@ class RemoteS3(RemoteBASE): if shared_creds: os.environ.setdefault("AWS_SHARED_CREDENTIALS_FILE", shared_creds) + @wrap_prop(threading.Lock()) @cached_property def s3(self): import boto3 diff --git a/setup.py b/setup.py index ae460aa3c..d029ee4d5 100644 --- a/setup.py +++ b/setup.py @@ -74,7 +74,7 @@ install_requires = [ "humanize>=0.5.1", "PyYAML>=5.1.2", "ruamel.yaml>=0.16.1", - "funcy>=1.12", + "funcy>=1.14", "pathspec>=0.6.0", "shortuuid>=0.5.0", "tqdm>=4.38.0,<5", From 8af1aeefb4b51e956e32dfabf58fdd6a44f56dd5 Mon Sep 17 00:00:00 2001 From: Alexander Schepanovski Date: Thu, 28 Nov 2019 21:10:35 +0100 Subject: [PATCH 06/16] remote: refactor ssh ask password code --- dvc/remote/ssh/__init__.py | 32 +++++++++++++++----------------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index 589bea8d1..7509d7504 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -8,8 +8,9 @@ import logging import os import threading from concurrent.futures import ThreadPoolExecutor -from contextlib import closing -from contextlib import contextmanager +from contextlib import closing, contextmanager + +from funcy import memoize, wrap_with import dvc.prompt as prompt from dvc.config import Config @@ -24,8 +25,15 @@ from dvc.utils.compat import urlparse logger = logging.getLogger(__name__) -saved_passwords = {} -saved_passwords_lock = threading.Lock() +@wrap_with(threading.Lock()) +@memoize +def ask_password(host, user, port): + return prompt.password( + "Enter a private key passphrase or a password for " + "host '{host}' port '{port}' user '{user}'".format( + host=host, port=port, user=user + ) + ) class RemoteSSH(RemoteBASE): @@ -120,21 +128,11 @@ class RemoteSSH(RemoteBASE): def ensure_credentials(self, path_info=None): if path_info is None: path_info = self.path_info - host, user, port = path_info.host, path_info.user, path_info.port + # NOTE: we use the same password regardless of the server :( if self.ask_password and self.password is None: - with saved_passwords_lock: - server_key = (host, user, port) - password = saved_passwords.get(server_key) - - if password is None: - saved_passwords[server_key] = password = prompt.password( - "Enter a private key passphrase or a password for " - "host '{host}' port '{port}' user '{user}'".format( - host=host, port=port, user=user - ) - ) - self.password = password + host, user, port = path_info.host, path_info.user, path_info.port + self.password = ask_password(host, user, port) def ssh(self, path_info): self.ensure_credentials(path_info) From 593ab17d2539cd7d58aee3b809a4d9b50213ab63 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Fri, 29 Nov 2019 02:14:25 +0200 Subject: [PATCH 07/16] setup: don't forget contextlib2 --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index ae460aa3c..b131559fb 100644 --- a/setup.py +++ b/setup.py @@ -161,6 +161,7 @@ setup( ":python_version=='2.7'": [ "futures", "pathlib2", + "contextlib2", "zc.lockfile>=1.2.1", ], ":python_version>='3.0'": ["flufl.lock>=3.2"], From e7227898e4e10a149d8255edf8c00a04c02f571b Mon Sep 17 00:00:00 2001 From: Max Risuhin Date: Fri, 29 Nov 2019 01:18:27 -0800 Subject: [PATCH 08/16] Support GDrive as remote --- dvc/config.py | 14 +- dvc/remote/__init__.py | 2 + dvc/remote/gdrive/__init__.py | 301 ++++++++++++++++++++++++ dvc/remote/gdrive/utils.py | 25 ++ dvc/scheme.py | 1 + setup.py | 7 +- tests/func/test_data_cloud.py | 76 ++++++ tests/unit/remote/gdrive/__init__.py | 0 tests/unit/remote/gdrive/conftest.py | 9 + tests/unit/remote/gdrive/test_gdrive.py | 9 + 10 files changed, 441 insertions(+), 3 deletions(-) create mode 100644 dvc/remote/gdrive/__init__.py create mode 100644 dvc/remote/gdrive/utils.py create mode 100644 tests/unit/remote/gdrive/__init__.py create mode 100644 tests/unit/remote/gdrive/conftest.py create mode 100644 tests/unit/remote/gdrive/test_gdrive.py diff --git a/dvc/config.py b/dvc/config.py index 90094b3a4..bee6fb8ef 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -110,6 +110,8 @@ class Config(object): # pylint: disable=too-many-instance-attributes CONFIG = "config" CONFIG_LOCAL = "config.local" + CREDENTIALPATH = "credentialpath" + LEVEL_LOCAL = 0 LEVEL_REPO = 1 LEVEL_GLOBAL = 2 @@ -162,7 +164,7 @@ class Config(object): # pylint: disable=too-many-instance-attributes } # aws specific options - SECTION_AWS_CREDENTIALPATH = "credentialpath" + SECTION_AWS_CREDENTIALPATH = CREDENTIALPATH SECTION_AWS_ENDPOINT_URL = "endpointurl" SECTION_AWS_LIST_OBJECTS = "listobjects" SECTION_AWS_REGION = "region" @@ -172,7 +174,7 @@ class Config(object): # pylint: disable=too-many-instance-attributes SECTION_AWS_ACL = "acl" # gcp specific options - SECTION_GCP_CREDENTIALPATH = SECTION_AWS_CREDENTIALPATH + SECTION_GCP_CREDENTIALPATH = CREDENTIALPATH SECTION_GCP_PROJECTNAME = "projectname" # azure specific option @@ -183,6 +185,11 @@ class Config(object): # pylint: disable=too-many-instance-attributes SECTION_OSS_ACCESS_KEY_SECRET = "oss_key_secret" SECTION_OSS_ENDPOINT = "oss_endpoint" + # GDrive options + SECTION_GDRIVE_CLIENT_ID = "gdrive_client_id" + SECTION_GDRIVE_CLIENT_SECRET = "gdrive_client_secret" + SECTION_GDRIVE_USER_CREDENTIALS_FILE = "gdrive_user_credentials_file" + SECTION_REMOTE_REGEX = r'^\s*remote\s*"(?P.*)"\s*$' SECTION_REMOTE_FMT = 'remote "{}"' SECTION_REMOTE_URL = "url" @@ -218,6 +225,9 @@ class Config(object): # pylint: disable=too-many-instance-attributes SECTION_OSS_ACCESS_KEY_ID: str, SECTION_OSS_ACCESS_KEY_SECRET: str, SECTION_OSS_ENDPOINT: str, + SECTION_GDRIVE_CLIENT_ID: str, + SECTION_GDRIVE_CLIENT_SECRET: str, + SECTION_GDRIVE_USER_CREDENTIALS_FILE: str, PRIVATE_CWD: str, Optional(SECTION_REMOTE_NO_TRAVERSE, default=True): Bool, } diff --git a/dvc/remote/__init__.py b/dvc/remote/__init__.py index 3ff90d365..e2c20a216 100644 --- a/dvc/remote/__init__.py +++ b/dvc/remote/__init__.py @@ -2,6 +2,7 @@ from __future__ import unicode_literals from .config import RemoteConfig from dvc.remote.azure import RemoteAZURE +from dvc.remote.gdrive import RemoteGDrive from dvc.remote.gs import RemoteGS from dvc.remote.hdfs import RemoteHDFS from dvc.remote.http import RemoteHTTP @@ -14,6 +15,7 @@ from dvc.remote.ssh import RemoteSSH REMOTES = [ RemoteAZURE, + RemoteGDrive, RemoteGS, RemoteHDFS, RemoteHTTP, diff --git a/dvc/remote/gdrive/__init__.py b/dvc/remote/gdrive/__init__.py new file mode 100644 index 000000000..fa0e25a28 --- /dev/null +++ b/dvc/remote/gdrive/__init__.py @@ -0,0 +1,301 @@ +from __future__ import unicode_literals + +import os +import posixpath +import logging +import threading + +from funcy import retry, compose, decorator, wrap_with +from funcy.py3 import cat + +from dvc.remote.gdrive.utils import TrackFileReadProgress, FOLDER_MIME_TYPE +from dvc.scheme import Schemes +from dvc.path_info import CloudURLInfo +from dvc.remote.base import RemoteBASE +from dvc.config import Config +from dvc.exceptions import DvcException +from dvc.utils import tmp_fname + +logger = logging.getLogger(__name__) + + +class GDriveRetriableError(DvcException): + def __init__(self, msg): + super(GDriveRetriableError, self).__init__(msg) + + +@decorator +def _wrap_pydrive_retriable(call): + from apiclient import errors + from pydrive.files import ApiRequestError + + try: + result = call() + except (ApiRequestError, errors.HttpError) as exception: + retry_codes = ["403", "500", "502", "503", "504"] + if any( + "HttpError {}".format(code) in str(exception) + for code in retry_codes + ): + raise GDriveRetriableError(msg="Google API request failed") + raise + return result + + +gdrive_retry = compose( + # 8 tries, start at 0.5s, multiply by golden ratio, cap at 10s + retry( + 8, GDriveRetriableError, timeout=lambda a: min(0.5 * 1.618 ** a, 10) + ), + _wrap_pydrive_retriable, +) + + +class RemoteGDrive(RemoteBASE): + scheme = Schemes.GDRIVE + path_cls = CloudURLInfo + REQUIRES = {"pydrive": "pydrive"} + GDRIVE_USER_CREDENTIALS_DATA = "GDRIVE_USER_CREDENTIALS_DATA" + DEFAULT_USER_CREDENTIALS_FILE = ".dvc/tmp/gdrive-user-credentials.json" + + def __init__(self, repo, config): + super(RemoteGDrive, self).__init__(repo, config) + self.no_traverse = False + self.path_info = self.path_cls(config[Config.SECTION_REMOTE_URL]) + self.config = config + self.init_drive() + + def init_drive(self): + self.client_id = self.config.get(Config.SECTION_GDRIVE_CLIENT_ID, None) + self.client_secret = self.config.get( + Config.SECTION_GDRIVE_CLIENT_SECRET, None + ) + if not self.client_id or not self.client_secret: + raise DvcException( + "Please specify Google Drive's client id and " + "secret in DVC's config. Learn more at " + "https://man.dvc.org/remote/add." + ) + self.gdrive_user_credentials_path = ( + tmp_fname(".dvc/tmp/") + if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA) + else self.config.get( + Config.SECTION_GDRIVE_USER_CREDENTIALS_FILE, + self.DEFAULT_USER_CREDENTIALS_FILE, + ) + ) + + @gdrive_retry + def gdrive_upload_file( + self, args, no_progress_bar=True, from_file="", progress_name="" + ): + item = self.drive.CreateFile( + {"title": args["title"], "parents": [{"id": args["parent_id"]}]} + ) + self.upload_file(item, no_progress_bar, from_file, progress_name) + return item + + def upload_file(self, item, no_progress_bar, from_file, progress_name): + with open(from_file, "rb") as opened_file: + if not no_progress_bar: + opened_file = TrackFileReadProgress(progress_name, opened_file) + # PyDrive doesn't like content property setting for empty files + # https://github.com/gsuitedevs/PyDrive/issues/121 + if os.stat(from_file).st_size: + item.content = opened_file + item.Upload() + + @gdrive_retry + def gdrive_download_file( + self, file_id, to_file, progress_name, no_progress_bar + ): + from dvc.progress import Tqdm + + gdrive_file = self.drive.CreateFile({"id": file_id}) + with Tqdm( + desc=progress_name, + total=int(gdrive_file["fileSize"]), + disable=no_progress_bar, + ): + gdrive_file.GetContentFile(to_file) + + def gdrive_list_item(self, query): + file_list = self.drive.ListFile({"q": query, "maxResults": 1000}) + + # Isolate and decorate fetching of remote drive items in pages + get_list = gdrive_retry(lambda: next(file_list, None)) + + # Fetch pages until None is received, lazily flatten the thing + return cat(iter(get_list, None)) + + def cache_root_dirs(self): + cached_dirs = {} + cached_ids = {} + for dir1 in self.gdrive_list_item( + "'{}' in parents and trashed=false".format(self.root_id) + ): + remote_path = posixpath.join(self.path_info.path, dir1["title"]) + cached_dirs.setdefault(remote_path, []).append(dir1["id"]) + cached_ids[dir1["id"]] = dir1["title"] + return cached_dirs, cached_ids + + @property + def cached_dirs(self): + if not hasattr(self, "_cached_dirs"): + self.drive + return self._cached_dirs + + @property + def cached_ids(self): + if not hasattr(self, "_cached_ids"): + self.drive + return self._cached_ids + + @property + @wrap_with(threading.RLock()) + def drive(self): + if not hasattr(self, "_gdrive"): + from pydrive.auth import GoogleAuth + from pydrive.drive import GoogleDrive + + if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA): + with open( + self.gdrive_user_credentials_path, "w" + ) as credentials_file: + credentials_file.write( + os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA) + ) + + GoogleAuth.DEFAULT_SETTINGS["client_config_backend"] = "settings" + GoogleAuth.DEFAULT_SETTINGS["client_config"] = { + "client_id": self.client_id, + "client_secret": self.client_secret, + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + "revoke_uri": "https://oauth2.googleapis.com/revoke", + "redirect_uri": "", + } + GoogleAuth.DEFAULT_SETTINGS["save_credentials"] = True + GoogleAuth.DEFAULT_SETTINGS["save_credentials_backend"] = "file" + GoogleAuth.DEFAULT_SETTINGS[ + "save_credentials_file" + ] = self.gdrive_user_credentials_path + GoogleAuth.DEFAULT_SETTINGS["get_refresh_token"] = True + GoogleAuth.DEFAULT_SETTINGS["oauth_scope"] = [ + "https://www.googleapis.com/auth/drive", + "https://www.googleapis.com/auth/drive.appdata", + ] + + # Pass non existent settings path to force DEFAULT_SETTINGS loading + gauth = GoogleAuth(settings_file="") + gauth.CommandLineAuth() + + if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA): + os.remove(self.gdrive_user_credentials_path) + + self._gdrive = GoogleDrive(gauth) + + self.root_id = self.get_remote_id(self.path_info, create=True) + self._cached_dirs, self._cached_ids = self.cache_root_dirs() + + return self._gdrive + + @gdrive_retry + def create_remote_dir(self, parent_id, title): + item = self.drive.CreateFile( + { + "title": title, + "parents": [{"id": parent_id}], + "mimeType": FOLDER_MIME_TYPE, + } + ) + item.Upload() + return item + + @gdrive_retry + def get_remote_item(self, name, parents_ids): + if not parents_ids: + return None + query = " or ".join( + "'{}' in parents".format(parent_id) for parent_id in parents_ids + ) + + query += " and trashed=false and title='{}'".format(name) + + # Limit found remote items count to 1 in response + item_list = self.drive.ListFile( + {"q": query, "maxResults": 1} + ).GetList() + return next(iter(item_list), None) + + def resolve_remote_item_from_path(self, path_parts, create): + parents_ids = ["root"] + current_path = "" + for path_part in path_parts: + current_path = posixpath.join(current_path, path_part) + remote_ids = self.get_remote_id_from_cache(current_path) + if remote_ids: + parents_ids = remote_ids + continue + item = self.get_remote_item(path_part, parents_ids) + if not item and create: + item = self.create_remote_dir(parents_ids[0], path_part) + elif not item: + return None + parents_ids = [item["id"]] + return item + + def get_remote_id_from_cache(self, remote_path): + if hasattr(self, "_cached_dirs"): + return self.cached_dirs.get(remote_path, []) + return [] + + def get_remote_id(self, path_info, create=False): + remote_ids = self.get_remote_id_from_cache(path_info.path) + + if remote_ids: + return remote_ids[0] + + file1 = self.resolve_remote_item_from_path( + path_info.path.split("/"), create + ) + return file1["id"] if file1 else "" + + def exists(self, path_info): + return self.get_remote_id(path_info) != "" + + def _upload(self, from_file, to_info, name, no_progress_bar): + dirname = to_info.parent + if dirname: + parent_id = self.get_remote_id(dirname, True) + else: + parent_id = to_info.bucket + + self.gdrive_upload_file( + {"title": to_info.name, "parent_id": parent_id}, + no_progress_bar, + from_file, + name, + ) + + def _download(self, from_info, to_file, name, no_progress_bar): + file_id = self.get_remote_id(from_info) + self.gdrive_download_file(file_id, to_file, name, no_progress_bar) + + def all(self): + if not self.cached_ids: + return + + query = " or ".join( + "'{}' in parents".format(dir_id) for dir_id in self.cached_ids + ) + + query += " and trashed=false" + for file1 in self.gdrive_list_item(query): + parent_id = file1["parents"][0]["id"] + path = posixpath.join(self.cached_ids[parent_id], file1["title"]) + try: + yield self.path_to_checksum(path) + except ValueError: + # We ignore all the non-cache looking files + logger.debug('Ignoring path as "non-cache looking"') diff --git a/dvc/remote/gdrive/utils.py b/dvc/remote/gdrive/utils.py new file mode 100644 index 000000000..781af811a --- /dev/null +++ b/dvc/remote/gdrive/utils.py @@ -0,0 +1,25 @@ +import os + +from dvc.progress import Tqdm + + +FOLDER_MIME_TYPE = "application/vnd.google-apps.folder" + + +class TrackFileReadProgress(object): + def __init__(self, progress_name, fobj): + self.progress_name = progress_name + self.fobj = fobj + file_size = os.fstat(fobj.fileno()).st_size + self.tqdm = Tqdm(desc=self.progress_name, total=file_size) + + def read(self, size): + self.tqdm.update(size) + return self.fobj.read(size) + + def close(self): + self.fobj.close() + self.tqdm.close() + + def __getattr__(self, attr): + return getattr(self.fobj, attr) diff --git a/dvc/scheme.py b/dvc/scheme.py index e12b768f5..5f7a8d1a2 100644 --- a/dvc/scheme.py +++ b/dvc/scheme.py @@ -9,5 +9,6 @@ class Schemes: HTTP = "http" HTTPS = "https" GS = "gs" + GDRIVE = "gdrive" LOCAL = "local" OSS = "oss" diff --git a/setup.py b/setup.py index cfd9e3447..b9b16f014 100644 --- a/setup.py +++ b/setup.py @@ -91,6 +91,10 @@ else: # Extra dependencies for remote integrations gs = ["google-cloud-storage==1.19.0"] +# google-api-python-client is internal dependency of pydrive. After merge of +# https://github.com/gsuitedevs/PyDrive/pull/180 into pydrive's master, +# usage of google-api-python-client can be removed from DVC. +gdrive = ["pydrive==1.3.1", "google-api-python-client>=1.2"] s3 = ["boto3==1.9.115"] azure = ["azure-storage-blob==2.1.0"] oss = ["oss2==2.6.1"] @@ -101,7 +105,7 @@ ssh = ["paramiko>=2.5.0"] # we can start shipping it by default. ssh_gssapi = ["paramiko[gssapi]>=2.5.0"] hdfs = ["pyarrow==0.14.0"] -all_remotes = gs + s3 + azure + ssh + oss +all_remotes = gs + s3 + azure + ssh + oss + gdrive if os.name != "nt" or sys.version_info[0] != 2: # NOTE: there are no pyarrow wheels for python2 on windows @@ -151,6 +155,7 @@ setup( extras_require={ "all": all_remotes, "gs": gs, + "gdrive": gdrive, "s3": s3, "azure": azure, "oss": oss, diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index 6eb94d7ea..ad0d33020 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -18,6 +18,7 @@ from dvc.config import Config from dvc.data_cloud import DataCloud from dvc.main import main from dvc.remote import RemoteAZURE +from dvc.remote import RemoteGDrive from dvc.remote import RemoteGS from dvc.remote import RemoteHDFS from dvc.remote import RemoteHTTP @@ -58,6 +59,11 @@ TEST_GCP_CREDS_FILE = os.path.abspath( # Ensure that absolute path is used os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = TEST_GCP_CREDS_FILE +TEST_GDRIVE_CLIENT_ID = ( + "719861249063-v4an78j9grdtuuuqg3lnm0sugna6v3lh.apps.googleusercontent.com" +) +TEST_GDRIVE_CLIENT_SECRET = "2fy_HyzSwkxkGzEken7hThXb" + def _should_test_aws(): do_test = env2bool("DVC_TEST_AWS", undefined=None) @@ -70,6 +76,13 @@ def _should_test_aws(): return False +def _should_test_gdrive(): + if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA): + return True + + return False + + def _should_test_gcp(): do_test = env2bool("DVC_TEST_GCP", undefined=None) if do_test is not None: @@ -202,6 +215,10 @@ def get_aws_url(): return "s3://" + get_aws_storagepath() +def get_gdrive_url(): + return "gdrive://root/" + str(uuid.uuid4()) + + def get_gcp_storagepath(): return TEST_GCP_REPO_BUCKET + "/" + str(uuid.uuid4()) @@ -375,6 +392,35 @@ class TestRemoteS3(TestDataCloudBase): return RemoteS3 +class TestRemoteGDrive(TestDataCloudBase): + def _should_test(self): + return _should_test_gdrive() + + def _setup_cloud(self): + self._ensure_should_run() + + repo = self._get_url() + + config = copy.deepcopy(TEST_CONFIG) + config[TEST_SECTION][Config.SECTION_REMOTE_URL] = repo + config[TEST_SECTION][ + Config.SECTION_GDRIVE_CLIENT_ID + ] = TEST_GDRIVE_CLIENT_ID + config[TEST_SECTION][ + Config.SECTION_GDRIVE_CLIENT_SECRET + ] = TEST_GDRIVE_CLIENT_SECRET + self.dvc.config.config = config + self.cloud = DataCloud(self.dvc) + + self.assertIsInstance(self.cloud.get_remote(), self._get_cloud_class()) + + def _get_url(self): + return get_gdrive_url() + + def _get_cloud_class(self): + return RemoteGDrive + + class TestRemoteGS(TestDataCloudBase): def _should_test(self): return _should_test_gcp() @@ -621,6 +667,36 @@ class TestRemoteS3CLI(TestDataCloudCLIBase): self._test_cloud(TEST_REMOTE) +class TestRemoteGDriveCLI(TestDataCloudCLIBase): + def _should_test(self): + return _should_test_gdrive() + + def _test(self): + url = get_gdrive_url() + + self.main(["remote", "add", TEST_REMOTE, url]) + self.main( + [ + "remote", + "modify", + TEST_REMOTE, + Config.SECTION_GDRIVE_CLIENT_ID, + TEST_GDRIVE_CLIENT_ID, + ] + ) + self.main( + [ + "remote", + "modify", + TEST_REMOTE, + Config.SECTION_GDRIVE_CLIENT_SECRET, + TEST_GDRIVE_CLIENT_SECRET, + ] + ) + + self._test_cloud(TEST_REMOTE) + + class TestRemoteGSCLI(TestDataCloudCLIBase): def _should_test(self): return _should_test_gcp() diff --git a/tests/unit/remote/gdrive/__init__.py b/tests/unit/remote/gdrive/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/remote/gdrive/conftest.py b/tests/unit/remote/gdrive/conftest.py new file mode 100644 index 000000000..035ca1509 --- /dev/null +++ b/tests/unit/remote/gdrive/conftest.py @@ -0,0 +1,9 @@ +import pytest + +from dvc.remote.gdrive import RemoteGDrive + + +@pytest.fixture +def gdrive(repo): + ret = RemoteGDrive(None, {"url": "gdrive://root/data"}) + return ret diff --git a/tests/unit/remote/gdrive/test_gdrive.py b/tests/unit/remote/gdrive/test_gdrive.py new file mode 100644 index 000000000..28e003748 --- /dev/null +++ b/tests/unit/remote/gdrive/test_gdrive.py @@ -0,0 +1,9 @@ +import mock +from dvc.remote.gdrive import RemoteGDrive + + +@mock.patch("dvc.remote.gdrive.RemoteGDrive.init_drive") +def test_init_drive(repo): + url = "gdrive://root/data" + gdrive = RemoteGDrive(repo, {"url": url}) + assert str(gdrive.path_info) == url From 72abfb0572cc2d45ff364003ca9e817a807f02e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= Date: Thu, 14 Nov 2019 15:33:40 +0100 Subject: [PATCH 09/16] remote: http: raise exception when download response with error status code --- dvc/exceptions.py | 8 ++++++++ dvc/remote/http.py | 9 +++++++-- tests/func/test_repro.py | 3 ++- tests/unit/remote/test_http.py | 24 ++++++++++++++++++++++++ tests/utils/httpd.py | 4 ++-- 5 files changed, 43 insertions(+), 5 deletions(-) diff --git a/dvc/exceptions.py b/dvc/exceptions.py index 1a167117a..3778be592 100644 --- a/dvc/exceptions.py +++ b/dvc/exceptions.py @@ -348,3 +348,11 @@ class NoOutputInExternalRepoError(DvcException): relpath(path, external_repo_path), external_repo_url ) ) + + +class HTTPErrorStatusCodeException(DvcException): + def __init__(self, code, reason): + super(HTTPErrorStatusCodeException, self).__init__( + "Server responded with error status code: '{}' and message: " + "'{}'".format(code, reason) + ) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index c22416625..d4ff1cd04 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -8,6 +8,7 @@ from funcy import cached_property, wrap_prop from dvc.config import Config from dvc.config import ConfigError from dvc.exceptions import DvcException +from dvc.exceptions import HTTPErrorStatusCodeException from dvc.progress import Tqdm from dvc.remote.base import RemoteBASE from dvc.scheme import Schemes @@ -37,7 +38,11 @@ class RemoteHTTP(RemoteBASE): ) def _download(self, from_info, to_file, name=None, no_progress_bar=False): - request = self._request("GET", from_info.url, stream=True) + response = self._request("GET", from_info.url, stream=True) + if response.status_code != 200: + raise HTTPErrorStatusCodeException( + response.status_code, response.reason + ) with Tqdm( total=None if no_progress_bar else self._content_length(from_info), leave=False, @@ -46,7 +51,7 @@ class RemoteHTTP(RemoteBASE): disable=no_progress_bar, ) as pbar: with open(to_file, "wb") as fd: - for chunk in request.iter_content(chunk_size=self.CHUNK_SIZE): + for chunk in response.iter_content(chunk_size=self.CHUNK_SIZE): fd.write(chunk) fd.flush() pbar.update(len(chunk)) diff --git a/tests/func/test_repro.py b/tests/func/test_repro.py index ae0727c14..f6e9a23a9 100644 --- a/tests/func/test_repro.py +++ b/tests/func/test_repro.py @@ -43,6 +43,7 @@ from tests.func.test_data_cloud import _should_test_ssh from tests.func.test_data_cloud import get_ssh_url from tests.func.test_data_cloud import TEST_AWS_REPO_BUCKET from tests.func.test_data_cloud import TEST_GCP_REPO_BUCKET +from tests.utils.httpd import ContentMD5Handler from tests.utils.httpd import StaticFileServer @@ -1176,7 +1177,7 @@ class TestReproExternalHTTP(TestReproExternalBase): self.dvc.remove("imported_file.dvc") - with StaticFileServer(handler="Content-MD5") as httpd: + with StaticFileServer(handler_class=ContentMD5Handler) as httpd: import_url = urljoin(self.get_remote(httpd.server_port), self.FOO) import_output = "imported_file" import_stage = self.dvc.imp_url(import_url, import_output) diff --git a/tests/unit/remote/test_http.py b/tests/unit/remote/test_http.py index 2b1ecdb07..d1c2a5efa 100644 --- a/tests/unit/remote/test_http.py +++ b/tests/unit/remote/test_http.py @@ -1,7 +1,11 @@ import pytest +from BaseHTTPServer import BaseHTTPRequestHandler from dvc.config import ConfigError +from dvc.exceptions import HTTPErrorStatusCodeException +from dvc.path_info import URLInfo from dvc.remote.http import RemoteHTTP +from tests.utils.httpd import StaticFileServer def test_no_traverse_compatibility(dvc_repo): @@ -13,3 +17,23 @@ def test_no_traverse_compatibility(dvc_repo): with pytest.raises(ConfigError): RemoteHTTP(dvc_repo, config) + + +@pytest.mark.parametrize("response_code", [404, 403, 500]) +def test_download_fails_on_error_code(response_code, dvc_repo): + class ErrorStatusRequestHandler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(response_code, message="Error message") + self.end_headers() + + with StaticFileServer(ErrorStatusRequestHandler) as httpd: + url = "http://localhost:{}/".format(httpd.server_port) + config = {"url": url} + + remote = RemoteHTTP(dvc_repo, config) + import os + + with pytest.raises(HTTPErrorStatusCodeException): + remote._download( + URLInfo(os.path.join(url, "file.txt")), "file.txt" + ) diff --git a/tests/utils/httpd.py b/tests/utils/httpd.py index ff902e71a..5883db7b7 100644 --- a/tests/utils/httpd.py +++ b/tests/utils/httpd.py @@ -45,9 +45,9 @@ class ContentMD5Handler(TestRequestHandler): class StaticFileServer: _lock = threading.Lock() - def __init__(self, handler="etag"): + def __init__(self, handler_class=ETagHandler): self._lock.acquire() - handler_class = ETagHandler if handler == "etag" else ContentMD5Handler + self.response_handler = handler_class self._httpd = HTTPServer(("localhost", 0), handler_class) self._thread = None From 7218ff202cb2d9043b1f0d6dca496ea793fbae83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= Date: Fri, 15 Nov 2019 11:08:22 +0100 Subject: [PATCH 10/16] remote: http: rename http error, refactor test --- dvc/exceptions.py | 7 +++---- dvc/remote/http.py | 7 ++----- tests/func/test_repro.py | 3 +-- tests/unit/remote/test_http.py | 20 ++++++++++---------- tests/utils/httpd.py | 1 - 5 files changed, 16 insertions(+), 22 deletions(-) diff --git a/dvc/exceptions.py b/dvc/exceptions.py index 3778be592..3ad8c72b8 100644 --- a/dvc/exceptions.py +++ b/dvc/exceptions.py @@ -350,9 +350,8 @@ class NoOutputInExternalRepoError(DvcException): ) -class HTTPErrorStatusCodeException(DvcException): +class HTTPError(DvcException): def __init__(self, code, reason): - super(HTTPErrorStatusCodeException, self).__init__( - "Server responded with error status code: '{}' and message: " - "'{}'".format(code, reason) + super(HTTPError, self).__init__( + "HTTP error: '{} {}'".format(code, reason) ) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index d4ff1cd04..4d5d15faa 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -7,8 +7,7 @@ from funcy import cached_property, wrap_prop from dvc.config import Config from dvc.config import ConfigError -from dvc.exceptions import DvcException -from dvc.exceptions import HTTPErrorStatusCodeException +from dvc.exceptions import DvcException, HTTPError from dvc.progress import Tqdm from dvc.remote.base import RemoteBASE from dvc.scheme import Schemes @@ -40,9 +39,7 @@ class RemoteHTTP(RemoteBASE): def _download(self, from_info, to_file, name=None, no_progress_bar=False): response = self._request("GET", from_info.url, stream=True) if response.status_code != 200: - raise HTTPErrorStatusCodeException( - response.status_code, response.reason - ) + raise HTTPError(response.status_code, response.reason) with Tqdm( total=None if no_progress_bar else self._content_length(from_info), leave=False, diff --git a/tests/func/test_repro.py b/tests/func/test_repro.py index f6e9a23a9..4e46fdde9 100644 --- a/tests/func/test_repro.py +++ b/tests/func/test_repro.py @@ -43,8 +43,7 @@ from tests.func.test_data_cloud import _should_test_ssh from tests.func.test_data_cloud import get_ssh_url from tests.func.test_data_cloud import TEST_AWS_REPO_BUCKET from tests.func.test_data_cloud import TEST_GCP_REPO_BUCKET -from tests.utils.httpd import ContentMD5Handler -from tests.utils.httpd import StaticFileServer +from tests.utils.httpd import StaticFileServer, ContentMD5Handler class TestRepro(TestDvc): diff --git a/tests/unit/remote/test_http.py b/tests/unit/remote/test_http.py index d1c2a5efa..924a05ae2 100644 --- a/tests/unit/remote/test_http.py +++ b/tests/unit/remote/test_http.py @@ -1,8 +1,12 @@ +try: + from http.server import BaseHTTPRequestHandler +except ImportError: + from BaseHTTPServer import BaseHTTPRequestHandler + import pytest -from BaseHTTPServer import BaseHTTPRequestHandler from dvc.config import ConfigError -from dvc.exceptions import HTTPErrorStatusCodeException +from dvc.exceptions import HTTPError from dvc.path_info import URLInfo from dvc.remote.http import RemoteHTTP from tests.utils.httpd import StaticFileServer @@ -19,11 +23,10 @@ def test_no_traverse_compatibility(dvc_repo): RemoteHTTP(dvc_repo, config) -@pytest.mark.parametrize("response_code", [404, 403, 500]) -def test_download_fails_on_error_code(response_code, dvc_repo): +def test_download_fails_on_error_code(dvc_repo): class ErrorStatusRequestHandler(BaseHTTPRequestHandler): def do_GET(self): - self.send_response(response_code, message="Error message") + self.send_response(404, message="Not found") self.end_headers() with StaticFileServer(ErrorStatusRequestHandler) as httpd: @@ -31,9 +34,6 @@ def test_download_fails_on_error_code(response_code, dvc_repo): config = {"url": url} remote = RemoteHTTP(dvc_repo, config) - import os - with pytest.raises(HTTPErrorStatusCodeException): - remote._download( - URLInfo(os.path.join(url, "file.txt")), "file.txt" - ) + with pytest.raises(HTTPError): + remote._download(URLInfo(url) / "file.txt", "file.txt") diff --git a/tests/utils/httpd.py b/tests/utils/httpd.py index 5883db7b7..92027ae75 100644 --- a/tests/utils/httpd.py +++ b/tests/utils/httpd.py @@ -47,7 +47,6 @@ class StaticFileServer: def __init__(self, handler_class=ETagHandler): self._lock.acquire() - self.response_handler = handler_class self._httpd = HTTPServer(("localhost", 0), handler_class) self._thread = None From e084b427e3c434ca1ecc89c751ccb387d492d4d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Redzy=C5=84ski?= Date: Thu, 21 Nov 2019 17:27:59 +0100 Subject: [PATCH 11/16] Update dvc/remote/http.py Co-Authored-By: Alexander Schepanovski --- dvc/remote/http.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index 4d5d15faa..e9ec4f888 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -41,7 +41,7 @@ class RemoteHTTP(RemoteBASE): if response.status_code != 200: raise HTTPError(response.status_code, response.reason) with Tqdm( - total=None if no_progress_bar else self._content_length(from_info), + total=None if no_progress_bar else self._content_length(response), leave=False, bytes=True, desc=from_info.url if name is None else name, From c102a781c248e6ee3b4f54bff1dbda36daa2ac32 Mon Sep 17 00:00:00 2001 From: pawel Date: Thu, 21 Nov 2019 17:45:29 +0100 Subject: [PATCH 12/16] HTTPError: reformat error message --- dvc/exceptions.py | 4 +--- dvc/remote/http.py | 1 - tests/unit/remote/test_http.py | 12 +----------- 3 files changed, 2 insertions(+), 15 deletions(-) diff --git a/dvc/exceptions.py b/dvc/exceptions.py index 3ad8c72b8..62c0a3466 100644 --- a/dvc/exceptions.py +++ b/dvc/exceptions.py @@ -352,6 +352,4 @@ class NoOutputInExternalRepoError(DvcException): class HTTPError(DvcException): def __init__(self, code, reason): - super(HTTPError, self).__init__( - "HTTP error: '{} {}'".format(code, reason) - ) + super(HTTPError, self).__init__("'{} {}'".format(code, reason)) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index e9ec4f888..bf7e93adb 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -50,7 +50,6 @@ class RemoteHTTP(RemoteBASE): with open(to_file, "wb") as fd: for chunk in response.iter_content(chunk_size=self.CHUNK_SIZE): fd.write(chunk) - fd.flush() pbar.update(len(chunk)) def exists(self, path_info): diff --git a/tests/unit/remote/test_http.py b/tests/unit/remote/test_http.py index 924a05ae2..c0d3bd42b 100644 --- a/tests/unit/remote/test_http.py +++ b/tests/unit/remote/test_http.py @@ -1,8 +1,3 @@ -try: - from http.server import BaseHTTPRequestHandler -except ImportError: - from BaseHTTPServer import BaseHTTPRequestHandler - import pytest from dvc.config import ConfigError @@ -24,12 +19,7 @@ def test_no_traverse_compatibility(dvc_repo): def test_download_fails_on_error_code(dvc_repo): - class ErrorStatusRequestHandler(BaseHTTPRequestHandler): - def do_GET(self): - self.send_response(404, message="Not found") - self.end_headers() - - with StaticFileServer(ErrorStatusRequestHandler) as httpd: + with StaticFileServer() as httpd: url = "http://localhost:{}/".format(httpd.server_port) config = {"url": url} From f7846a5bed646596e169250d5260e8ba1cab2cbb Mon Sep 17 00:00:00 2001 From: pawel Date: Thu, 21 Nov 2019 17:50:03 +0100 Subject: [PATCH 13/16] test: remote: http: clear test file name --- dvc/remote/http.py | 2 +- tests/unit/remote/test_http.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index bf7e93adb..b93adc395 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -41,7 +41,7 @@ class RemoteHTTP(RemoteBASE): if response.status_code != 200: raise HTTPError(response.status_code, response.reason) with Tqdm( - total=None if no_progress_bar else self._content_length(response), + total=None if no_progress_bar else self._content_length(from_info), leave=False, bytes=True, desc=from_info.url if name is None else name, diff --git a/tests/unit/remote/test_http.py b/tests/unit/remote/test_http.py index c0d3bd42b..dbd54bcd1 100644 --- a/tests/unit/remote/test_http.py +++ b/tests/unit/remote/test_http.py @@ -26,4 +26,4 @@ def test_download_fails_on_error_code(dvc_repo): remote = RemoteHTTP(dvc_repo, config) with pytest.raises(HTTPError): - remote._download(URLInfo(url) / "file.txt", "file.txt") + remote._download(URLInfo(url) / "missing.txt", "missing.txt") From 03c27192a6bb6172ad819ee39a3a787b57100714 Mon Sep 17 00:00:00 2001 From: pawel Date: Wed, 27 Nov 2019 11:57:06 +0100 Subject: [PATCH 14/16] remote: http: calculate length basing on response and not head call --- dvc/remote/http.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index b93adc395..98a1e50f5 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -41,7 +41,7 @@ class RemoteHTTP(RemoteBASE): if response.status_code != 200: raise HTTPError(response.status_code, response.reason) with Tqdm( - total=None if no_progress_bar else self._content_length(from_info), + total=None if no_progress_bar else self._content_length(response), leave=False, bytes=True, desc=from_info.url if name is None else name, @@ -55,13 +55,8 @@ class RemoteHTTP(RemoteBASE): def exists(self, path_info): return bool(self._request("HEAD", path_info.url)) - def _content_length(self, url_or_request): - headers = getattr( - url_or_request, - "headers", - self._request("HEAD", url_or_request).headers, - ) - res = headers.get("Content-Length") + def _content_length(self, response): + res = response.headers.get("Content-Length") return int(res) if res else None def get_file_checksum(self, path_info): From 845cd352f78b7b89b5a63f19529cb8ad445f81ea Mon Sep 17 00:00:00 2001 From: "Mr. Outis" Date: Fri, 29 Nov 2019 11:45:45 -0600 Subject: [PATCH 15/16] s3: add support for top level empty directories --- dvc/remote/s3.py | 12 ++++++++++++ tests/unit/remote/test_s3.py | 12 +++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 5594806f0..75084ad51 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -212,6 +212,15 @@ class RemoteS3(RemoteBASE): fname = next(self._list_paths(path_info, max_items=1), "") return path_info.path == fname or fname.startswith(dir_path.path) + def makedirs(self, path_info): + # We need to support creating empty directories, on S3, that means + # creating an object with an empty body and a trailing slash `/`. + # + # We are not creating directory objects for every parent prefix, + # it doesn't make sense. + dir_path = path_info / "" + self.s3.put_object(Bucket=path_info.bucket, Key=dir_path.path, Body="") + def isdir(self, path_info): # S3 doesn't have a concept for directories. # @@ -271,4 +280,7 @@ class RemoteS3(RemoteBASE): def walk_files(self, path_info, max_items=None): for fname in self._list_paths(path_info, max_items): + if fname.endswith("/"): + continue + yield path_info.replace(path=fname) diff --git a/tests/unit/remote/test_s3.py b/tests/unit/remote/test_s3.py index a82d005c4..96b729917 100644 --- a/tests/unit/remote/test_s3.py +++ b/tests/unit/remote/test_s3.py @@ -79,9 +79,11 @@ def test_walk_files(remote): remote.path_info / "data/subdir/1", remote.path_info / "data/subdir/2", remote.path_info / "data/subdir/3", + remote.path_info / "empty_file", + remote.path_info / "foo", ] - assert list(remote.walk_files(remote.path_info / "data")) == files + assert list(remote.walk_files(remote.path_info)) == files def test_copy_preserve_etag_across_buckets(remote): @@ -99,3 +101,11 @@ def test_copy_preserve_etag_across_buckets(remote): to_etag = RemoteS3.get_etag(s3, "another", "foo") assert from_etag == to_etag + + +def makedirs(remote): + empty_dir = remote.path_info / "empty_dir" + remote.remove(empty_dir) + assert not remote.exists(empty_dir) + remote.makedirs(empty_dir) + assert remote.exists(empty_dir) From 1f9a0f0e0fe25f5e17cffb04a404b2bd0deeca7e Mon Sep 17 00:00:00 2001 From: "Mr. Outis" Date: Fri, 29 Nov 2019 12:38:12 -0600 Subject: [PATCH 16/16] address @efiop's suggestions --- dvc/remote/s3.py | 4 ++-- tests/unit/remote/test_s3.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 75084ad51..8013f9bb9 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -213,11 +213,11 @@ class RemoteS3(RemoteBASE): return path_info.path == fname or fname.startswith(dir_path.path) def makedirs(self, path_info): - # We need to support creating empty directories, on S3, that means + # We need to support creating empty directories, which means # creating an object with an empty body and a trailing slash `/`. # # We are not creating directory objects for every parent prefix, - # it doesn't make sense. + # as it is not required. dir_path = path_info / "" self.s3.put_object(Bucket=path_info.bucket, Key=dir_path.path, Body="") diff --git a/tests/unit/remote/test_s3.py b/tests/unit/remote/test_s3.py index 96b729917..49fc3dbb8 100644 --- a/tests/unit/remote/test_s3.py +++ b/tests/unit/remote/test_s3.py @@ -103,8 +103,8 @@ def test_copy_preserve_etag_across_buckets(remote): assert from_etag == to_etag -def makedirs(remote): - empty_dir = remote.path_info / "empty_dir" +def test_makedirs(remote): + empty_dir = remote.path_info / "empty_dir" / "" remote.remove(empty_dir) assert not remote.exists(empty_dir) remote.makedirs(empty_dir)