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/exceptions.py b/dvc/exceptions.py index 92f46a882..62c0a3466 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, + ) ) @@ -349,3 +348,8 @@ class NoOutputInExternalRepoError(DvcException): relpath(path, external_repo_path), external_repo_url ) ) + + +class HTTPError(DvcException): + def __init__(self, code, reason): + super(HTTPError, self).__init__("'{} {}'".format(code, reason)) 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) 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/__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/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/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( 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/remote/gs.py b/dvc/remote/gs.py index f89654461..584977636 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/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/http.py b/dvc/remote/http.py index 7f96e0b8e..98a1e50f5 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -1,12 +1,13 @@ 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 -from dvc.exceptions import DvcException +from dvc.exceptions import DvcException, HTTPError from dvc.progress import Tqdm from dvc.remote.base import RemoteBASE from dvc.scheme import Schemes @@ -36,30 +37,26 @@ 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 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, 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)) 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): @@ -81,6 +78,7 @@ class RemoteHTTP(RemoteBASE): return etag + @wrap_prop(threading.Lock()) @cached_property def _session(self): import requests diff --git a/dvc/remote/local.py b/dvc/remote/local.py index da922904c..ff88dfdcc 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -83,17 +83,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/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..8013f9bb9 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 @@ -210,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, which means + # creating an object with an empty body and a trailing slash `/`. + # + # We are not creating directory objects for every parent prefix, + # as it is not required. + 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. # @@ -269,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/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index e8469f968..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) @@ -264,7 +262,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: 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 ae460aa3c..b9b16f014 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", @@ -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, @@ -161,6 +166,7 @@ setup( ":python_version=='2.7'": [ "futures", "pathlib2", + "contextlib2", "zc.lockfile>=1.2.1", ], ":python_version>='3.0'": ["flufl.lock>=3.2"], diff --git a/tests/func/test_api.py b/tests/func/test_api.py index a8f293853..959d9a198 100644 --- a/tests/func/test_api.py +++ b/tests/func/test_api.py @@ -32,13 +32,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) @@ -55,7 +55,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_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/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 ee6c57d1f..20e40fde7 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): diff --git a/tests/func/test_repro.py b/tests/func/test_repro.py index ae0727c14..4e46fdde9 100644 --- a/tests/func/test_repro.py +++ b/tests/func/test_repro.py @@ -43,7 +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 StaticFileServer +from tests.utils.httpd import StaticFileServer, ContentMD5Handler class TestRepro(TestDvc): @@ -1176,7 +1176,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/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 diff --git a/tests/unit/remote/test_http.py b/tests/unit/remote/test_http.py index 2b1ecdb07..dbd54bcd1 100644 --- a/tests/unit/remote/test_http.py +++ b/tests/unit/remote/test_http.py @@ -1,7 +1,10 @@ import pytest from dvc.config import ConfigError +from dvc.exceptions import HTTPError +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 +16,14 @@ def test_no_traverse_compatibility(dvc_repo): with pytest.raises(ConfigError): RemoteHTTP(dvc_repo, config) + + +def test_download_fails_on_error_code(dvc_repo): + with StaticFileServer() as httpd: + url = "http://localhost:{}/".format(httpd.server_port) + config = {"url": url} + + remote = RemoteHTTP(dvc_repo, config) + + with pytest.raises(HTTPError): + remote._download(URLInfo(url) / "missing.txt", "missing.txt") diff --git a/tests/unit/remote/test_remote_dir.py b/tests/unit/remote/test_remote_dir.py index 827fcab0e..e919d4040 100644 --- a/tests/unit/remote/test_remote_dir.py +++ b/tests/unit/remote/test_remote_dir.py @@ -92,7 +92,17 @@ def test_copy_preserve_etag_across_buckets(remote): remote.copy(from_info, to_info) - from_etag = RemoteS3.get_etag(s3, "bucket", "foo") + from_etag = RemoteS3.get_etag(s3, from_info.bucket, from_info.path) to_etag = RemoteS3.get_etag(s3, "another", "foo") assert from_etag == to_etag + + +@pytest.mark.parametrize("remote", remotes, indirect=True) +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) + assert remote.exists(empty_dir) + assert remote.isdir(empty_dir) diff --git a/tests/utils/httpd.py b/tests/utils/httpd.py index ff902e71a..92027ae75 100644 --- a/tests/utils/httpd.py +++ b/tests/utils/httpd.py @@ -45,9 +45,8 @@ 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._httpd = HTTPServer(("localhost", 0), handler_class) self._thread = None