diff --git a/dvc/remote/gdrive.py b/dvc/remote/gdrive.py index 4bffb60b4..3ad4a7800 100644 --- a/dvc/remote/gdrive.py +++ b/dvc/remote/gdrive.py @@ -1,8 +1,9 @@ import os import posixpath import logging -import threading import re +import threading +from urllib.parse import urlparse from funcy import retry, compose, decorator, wrap_with from funcy.py3 import cat @@ -23,12 +24,22 @@ class GDriveRetriableError(DvcException): pass +class GDrivePathNotFound(DvcException): + def __init__(self, path_info): + super().__init__("Google Drive path '{}' not found.".format(path_info)) + + class GDriveAccessTokenRefreshError(DvcException): - pass + def __init__(self): + super().__init__("Google Drive access token refreshment is failed.") class GDriveMissedCredentialKeyError(DvcException): - pass + def __init__(self, path): + super().__init__( + "Google Drive user credentials file '{}' " + "misses value for key.".format(path) + ) @decorator @@ -49,17 +60,32 @@ def _wrap_pydrive_retriable(call): gdrive_retry = compose( - # 8 tries, start at 0.5s, multiply by golden ratio, cap at 10s + # 15 tries, start at 0.5s, multiply by golden ratio, cap at 20s retry( - 8, GDriveRetriableError, timeout=lambda a: min(0.5 * 1.618 ** a, 10) + 15, GDriveRetriableError, timeout=lambda a: min(0.5 * 1.618 ** a, 20) ), _wrap_pydrive_retriable, ) +class GDriveURLInfo(CloudURLInfo): + def __init__(self, url): + super().__init__(url) + + # GDrive URL host part is case sensitive, + # we are restoring it here. + p = urlparse(url) + self.host = p.netloc + assert self.netloc == self.host + + # Normalize path. Important since we have a cache (path to ID) + # and don't want to deal with different variations of path in it. + self._spath = re.sub("/{2,}", "/", self._spath.rstrip("/")) + + class RemoteGDrive(RemoteBASE): scheme = Schemes.GDRIVE - path_cls = CloudURLInfo + path_cls = GDriveURLInfo REQUIRES = {"pydrive2": "pydrive2"} DEFAULT_NO_TRAVERSE = False DEFAULT_VERIFY = True @@ -69,32 +95,32 @@ class RemoteGDrive(RemoteBASE): def __init__(self, repo, config): super().__init__(repo, config) - self.path_info = self.path_cls(config[Config.SECTION_REMOTE_URL]) - - bucket = re.search( - "{}://(.*)".format(self.scheme), - config[Config.SECTION_REMOTE_URL], - re.IGNORECASE, - ) - self.bucket = ( - bucket.group(1).split("/")[0] if bucket else self.path_info.bucket - ) - + url = config[Config.SECTION_REMOTE_URL] + self.path_info = self.path_cls(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( + if not self.path_info.bucket: + raise DvcException( + "Empty Google Drive URL '{}'. Learn more at " + "{}.".format( + url, format_link("https://man.dvc.org/remote/add") + ) + ) + + self._bucket = self.path_info.bucket + 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: + 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 " + "secret in DVC config. Learn more at " "{}.".format(format_link("https://man.dvc.org/remote/add")) ) - self.gdrive_user_credentials_path = ( + self._gdrive_user_credentials_path = ( tmp_fname(os.path.join(self.repo.tmp_dir, "")) if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA) else self.config.get( @@ -104,15 +130,134 @@ class RemoteGDrive(RemoteBASE): ), ) ) - self.remote_drive_id = None + + self._list_params = None + self._gdrive = None + + self._cache_initialized = False + self._remote_root_id = None + self._cached_dirs = None + self._cached_ids = None + + @property + @wrap_with(threading.RLock()) + def drive(self): + from pydrive2.auth import RefreshError + + if not self._gdrive: + from pydrive2.auth import GoogleAuth + from pydrive2.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="") + + try: + gauth.CommandLineAuth() + except RefreshError as exc: + raise GDriveAccessTokenRefreshError from exc + except KeyError as exc: + raise GDriveMissedCredentialKeyError( + self._gdrive_user_credentials_path + ) from exc + # Handle pydrive2.auth.AuthenticationError and other auth failures + except Exception as exc: + raise DvcException( + "Google Drive authentication failed" + ) from exc + finally: + if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA): + os.remove(self._gdrive_user_credentials_path) + + self._gdrive = GoogleDrive(gauth) + + return self._gdrive + + @wrap_with(threading.RLock()) + def _initialize_cache(self): + if self._cache_initialized: + return + + cached_dirs = {} + cached_ids = {} + self._remote_root_id = self._get_remote_id(self.path_info) + for dir1 in self.gdrive_list_item( + "'{}' in parents and trashed=false".format(self._remote_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"] + + self._cached_dirs = cached_dirs + self._cached_ids = cached_ids + self._cache_initialized = True + + @property + def cached_dirs(self): + if not self._cache_initialized: + self._initialize_cache() + return self._cached_dirs + + @property + def cached_ids(self): + if not self._cache_initialized: + self._initialize_cache() + return self._cached_ids + + @property + def remote_root_id(self): + if not self._cache_initialized: + self._initialize_cache() + return self._remote_root_id + + @property + def list_params(self): + if not self._list_params: + params = {"corpora": "default"} + if self._bucket != "root" and self._bucket != "appDataFolder": + params["driveId"] = self._get_remote_drive_id(self._bucket) + params["corpora"] = "drive" + self._list_params = params + return self._list_params @gdrive_retry def gdrive_upload_file( - self, args, no_progress_bar=True, from_file="", progress_name="" + self, + parent_id, + title, + no_progress_bar=True, + from_file="", + progress_name="", ): - parent = {"id": args["parent_id"]} item = self.drive.CreateFile( - {"title": args["title"], "parents": [parent]} + {"title": title, "parents": [{"id": parent_id}]} ) with open(from_file, "rb") as fobj: @@ -139,7 +284,7 @@ class RemoteGDrive(RemoteBASE): # it does not create a file on the remote gdrive_file = self.drive.CreateFile(param) bar_format = ( - "Donwloading {desc:{ncols_desc}.{ncols_desc}}... " + "Downloading {desc:{ncols_desc}.{ncols_desc}}... " + Tqdm.format_sizeof(int(gdrive_file["fileSize"]), "B", 1024) ) with Tqdm( @@ -148,10 +293,8 @@ class RemoteGDrive(RemoteBASE): gdrive_file.GetContentFile(to_file) def gdrive_list_item(self, query): - param = {"q": query, "maxResults": 1000, "corpora": self.corpora} - - if self.remote_drive_id: - param["driveId"] = self.remote_drive_id + param = {"q": query, "maxResults": 1000} + param.update(self.list_params) file_list = self.drive.ListFile(param) @@ -161,113 +304,23 @@ class RemoteGDrive(RemoteBASE): # 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.remote_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 - def corpora(self): - if not hasattr(self, "_corpora"): - self.drive - return self._corpora - - @property @wrap_with(threading.RLock()) - def drive(self): - from pydrive2.auth import RefreshError + def gdrive_create_dir(self, parent_id, title, remote_path): + if parent_id == self.remote_root_id: + cached = self.cached_dirs.get(remote_path, []) + if cached: + return cached[0] - if not hasattr(self, "_gdrive"): - from pydrive2.auth import GoogleAuth - from pydrive2.drive import GoogleDrive + item = self._create_remote_dir(parent_id, title) - 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) - ) + if parent_id == self.remote_root_id: + self.cached_dirs.setdefault(remote_path, []).append(item["id"]) + self.cached_ids[item["id"]] = item["title"] - 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", - # drive.appdata grants access to appDataFolder GDrive directory - "https://www.googleapis.com/auth/drive.appdata", - ] - - # Pass non existent settings path to force DEFAULT_SETTINGS loading - gauth = GoogleAuth(settings_file="") - - try: - gauth.CommandLineAuth() - except RefreshError as exc: - raise GDriveAccessTokenRefreshError( - "Google Drive's access token refreshment is failed" - ) from exc - except KeyError as exc: - raise GDriveMissedCredentialKeyError( - "Google Drive's user credentials file '{}' " - "misses value for key '{}'".format( - self.gdrive_user_credentials_path, str(exc) - ) - ) - # Handle pydrive2.auth.AuthenticationError and others auth failures - except Exception as exc: - raise DvcException( - "Google Drive authentication failed" - ) from exc - finally: - if os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA): - os.remove(self.gdrive_user_credentials_path) - - self._gdrive = GoogleDrive(gauth) - - if self.bucket != "root" and self.bucket != "appDataFolder": - self.remote_drive_id = self.get_remote_drive_id(self.bucket) - self._corpora = "drive" if self.remote_drive_id else "default" - self.remote_root_id = self.get_remote_id( - self.path_info, create=True - ) - - self._cached_dirs, self._cached_ids = self.cache_root_dirs() - - return self._gdrive + return item["id"] @gdrive_retry - def create_remote_dir(self, parent_id, title): + def _create_remote_dir(self, parent_id, title): parent = {"id": parent_id} item = self.drive.CreateFile( {"title": title, "parents": [parent], "mimeType": FOLDER_MIME_TYPE} @@ -276,14 +329,14 @@ class RemoteGDrive(RemoteBASE): return item @gdrive_retry - def delete_remote_file(self, remote_id): + def _delete_remote_file(self, remote_id): param = {"id": remote_id} # it does not create a file on the remote item = self.drive.CreateFile(param) item.Delete() @gdrive_retry - def get_remote_item(self, name, parents_ids): + def _get_remote_item(self, name, parents_ids): if not parents_ids: return None query = "({})".format( @@ -295,84 +348,78 @@ class RemoteGDrive(RemoteBASE): query += " and trashed=false and title='{}'".format(name) - param = { - "q": query, - # Remote might contain items with duplicated titles - "maxResults": 1, - "corpora": self.corpora, - } - - if self.remote_drive_id: - param["driveId"] = self.remote_drive_id + # Remote might contain items with duplicated path (titles). + # We thus limit number of items. + param = {"q": query, "maxResults": 1} + param.update(self.list_params) # Limit found remote items count to 1 in response item_list = self.drive.ListFile(param).GetList() return next(iter(item_list), None) @gdrive_retry - def get_remote_drive_id(self, remote_id): + def _get_remote_drive_id(self, remote_id): param = {"id": remote_id} # it does not create a file on the remote item = self.drive.CreateFile(param) item.FetchMetadata("driveId") return item.get("driveId", None) - def resolve_remote_item_from_path(self, path_parts, create): - parents_ids = [self.bucket] - 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, []) + def _get_cached_remote_ids(self, path): + if not path: + return [self._bucket] + if self._cache_initialized: + if path == self.path_info.path: + return [self.remote_root_id] + return self.cached_dirs.get(path, []) return [] - def get_remote_id(self, path_info, create=False): - if not path_info.path and path_info.bucket: - # Case sensitive base path - return self.bucket - - remote_ids = self.get_remote_id_from_cache(path_info.path) - + def _path_to_remote_ids(self, path, create): + remote_ids = self._get_cached_remote_ids(path) if remote_ids: - return remote_ids[0] + return remote_ids - file1 = self.resolve_remote_item_from_path( - path_info.path.split("/"), create - ) - return file1["id"] if file1 else "" + parent_path, part = posixpath.split(path) + parent_ids = self._path_to_remote_ids(parent_path, create) + item = self._get_remote_item(part, parent_ids) + + if not item: + return ( + [self.gdrive_create_dir(parent_ids[0], part, path)] + if create + else [] + ) + + return [item["id"]] + + def _get_remote_id(self, path_info, create=False): + assert path_info.bucket == self._bucket + + remote_ids = self._path_to_remote_ids(path_info.path, create) + if not remote_ids: + raise GDrivePathNotFound(path_info) + + return remote_ids[0] def exists(self, path_info): - return self.get_remote_id(path_info) != "" + try: + self._get_remote_id(path_info) + except GDrivePathNotFound: + return False + else: + return True 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 + assert dirname + parent_id = self._get_remote_id(dirname, True) self.gdrive_upload_file( - {"title": to_info.name, "parent_id": parent_id}, - no_progress_bar, - from_file, - name, + parent_id, to_info.name, 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) + file_id = self._get_remote_id(from_info) self.gdrive_download_file(file_id, to_file, name, no_progress_bar) def all(self): @@ -396,5 +443,5 @@ class RemoteGDrive(RemoteBASE): logger.debug('Ignoring path as "non-cache looking"') def remove(self, path_info): - remote_id = self.get_remote_id(path_info) - self.delete_remote_file(remote_id) + remote_id = self._get_remote_id(path_info) + self._delete_remote_file(remote_id) diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index f4de70631..2572bb9a4 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -194,24 +194,30 @@ class TestRemoteS3(S3, TestDataCloudBase): return RemoteS3 +def setup_gdrive_cloud(remote_url, dvc): + config = copy.deepcopy(TEST_CONFIG) + config[TEST_SECTION][Config.SECTION_REMOTE_URL] = remote_url + config[TEST_SECTION][ + Config.SECTION_GDRIVE_CLIENT_ID + ] = TEST_GDRIVE_CLIENT_ID + config[TEST_SECTION][ + Config.SECTION_GDRIVE_CLIENT_SECRET + ] = TEST_GDRIVE_CLIENT_SECRET + + dvc.config.config = config + remote = DataCloud(dvc).get_remote() + remote._create_remote_dir("root", remote.path_info.path) + + class TestRemoteGDrive(GDrive, TestDataCloudBase): def _setup_cloud(self): self._ensure_should_run() - repo = self.get_url() + setup_gdrive_cloud(self.get_url(), self.dvc) - 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()) + remote = self.cloud.get_remote() + self.assertIsInstance(remote, self._get_cloud_class()) def _get_cloud_class(self): return RemoteGDrive @@ -303,7 +309,12 @@ class TestDataCloudCLIBase(TestDvc): def get_url(): raise NotImplementedError + def _setup_cloud(self): + pass + def _test_cloud(self, remote=None): + self._setup_cloud() + args = ["-v", "-j", "2"] if remote: args += ["-r", remote] @@ -402,6 +413,9 @@ class TestRemoteS3CLI(S3, TestDataCloudCLIBase): class TestRemoteGDriveCLI(GDrive, TestDataCloudCLIBase): + def _setup_cloud(self): + setup_gdrive_cloud(self.get_url(), self.dvc) + def _test(self): url = self.get_url() diff --git a/tests/remotes.py b/tests/remotes.py index 35876b676..b44637451 100644 --- a/tests/remotes.py +++ b/tests/remotes.py @@ -150,9 +150,10 @@ class GDrive: def should_test(): return os.getenv(RemoteGDrive.GDRIVE_USER_CREDENTIALS_DATA) is not None - @staticmethod - def get_url(): - return "gdrive://root/" + str(uuid.uuid4()) + def get_url(self): + if not getattr(self, "_remote_url", None): + self._remote_url = "gdrive://root/" + str(uuid.uuid4()) + return self._remote_url class Azure: