From f4b14c2d4a462c593ac7b777322e171d2407f1fd Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Sat, 13 Feb 2021 04:10:38 +0300 Subject: [PATCH] gdrive: Implement .ls() protocol / proper tree walking with .walk_files() (#5433) * gdrive: .ls()/.walk()/.walk_dir() * better walking * proper tests * ._ls_recursive() * use posxipath for tests * store ids instead of raw names * add gdrive fixture * add dependency * tests: temporarily disable gdrive tests Co-authored-by: Ruslan Kuprieiev --- dvc/tree/gdrive.py | 117 +++++++++++++++++++++++-------- setup.py | 1 + tests/func/test_tree.py | 31 +++++++- tests/remotes/gdrive.py | 43 +++++++++++- tests/unit/remote/test_gdrive.py | 18 +++-- 5 files changed, 167 insertions(+), 43 deletions(-) diff --git a/dvc/tree/gdrive.py b/dvc/tree/gdrive.py index cb2b4abaa..86e8cc33c 100644 --- a/dvc/tree/gdrive.py +++ b/dvc/tree/gdrive.py @@ -88,6 +88,7 @@ class GDriveURLInfo(CloudURLInfo): class GDriveTree(BaseTree): scheme = Schemes.GDRIVE PATH_CLS = GDriveURLInfo + PARAM_CHECKSUM = "checksum" REQUIRES = {"pydrive2": "pydrive2"} # Always prefer traverse for GDrive since API usage quotas are a concern. TRAVERSE_WEIGHT_MULTIPLIER = 1 @@ -285,20 +286,21 @@ class GDriveTree(BaseTree): ), } - self._cache_path_id(self._path, cache["root_id"], cache) + self._cache_path_id(self._path, cache["root_id"], cache=cache) for item in self._gdrive_list( "'{}' in parents and trashed=false".format(cache["root_id"]) ): item_path = (self.path_info / item["title"]).path - self._cache_path_id(item_path, item["id"], cache) + self._cache_path_id(item_path, item["id"], cache=cache) return cache - def _cache_path_id(self, path, item_id, cache=None): + def _cache_path_id(self, path, *item_ids, cache=None): cache = cache or self._ids_cache - cache["dirs"][path].append(item_id) - cache["ids"][item_id] = path + for item_id in item_ids: + cache["dirs"][path].append(item_id) + cache["ids"][item_id] = path @cached_property def _list_params(self): @@ -477,7 +479,7 @@ class GDriveTree(BaseTree): return self._ids_cache["dirs"].get(path, []) return [] - def _path_to_item_ids(self, path, create, use_cache): + def _path_to_item_ids(self, path, create=False, use_cache=True): item_ids = self._get_cached_item_ids(path, use_cache) if item_ids: return item_ids @@ -510,35 +512,88 @@ class GDriveTree(BaseTree): else: return True - def _list_paths(self, prefix=None): - if not self._ids_cache["ids"]: - return - - if prefix: - dir_ids = self._ids_cache["dirs"].get(prefix[:2]) - if not dir_ids: - return - else: - dir_ids = self._ids_cache["ids"] - parents_query = " or ".join( - f"'{dir_id}' in parents" for dir_id in dir_ids + def _gdrive_list_ids(self, query_ids): + query = " or ".join( + f"'{query_id}' in parents" for query_id in query_ids ) - query = f"({parents_query}) and trashed=false" + query = f"({query}) and trashed=false" + return self._gdrive_list(query) - for item in self._gdrive_list(query): - parent_id = item["parents"][0]["id"] - yield posixpath.join( - self._ids_cache["ids"][parent_id], item["title"] - ) + def _ls_recursive(self, path_info, detail=False): + root_path = path_info.path + seen_paths = set() + + dir_ids = [self._ids_cache["ids"].copy()] + while dir_ids: + query_ids = { + dir_id: dir_name + for dir_id, dir_name in dir_ids.pop().items() + if posixpath.commonpath([root_path, dir_name]) == root_path + if dir_id not in seen_paths + } + if not query_ids: + continue + + seen_paths |= query_ids.keys() + + new_query_ids = {} + dir_ids.append(new_query_ids) + for item in self._gdrive_list_ids(query_ids): + parent_id = item["parents"][0]["id"] + item_path = posixpath.join(query_ids[parent_id], item["title"]) + if item["mimeType"] == FOLDER_MIME_TYPE: + new_query_ids[item["id"]] = item_path + self._cache_path_id(item_path, item["id"]) + continue + + if detail: + yield { + "type": "file", + "name": item_path, + "size": item["fileSize"], + "checksum": item["md5Checksum"], + } + else: + yield item_path + + def ls( + self, path_info, detail=False, recursive=False + ): # pylint: disable=arguments-differ + if recursive: + yield from self._ls_recursive(path_info, detail=detail) + return None + + cached = path_info.path in self._ids_cache["dirs"] + if cached: + dir_ids = self._ids_cache["dirs"][path_info.path] + else: + dir_ids = self._path_to_item_ids(path_info.path) + + if not dir_ids: + return None + + root_path = path_info.path + for item in self._gdrive_list_ids(dir_ids): + item_path = posixpath.join(root_path, item["title"]) + if detail: + if item["mimeType"] == FOLDER_MIME_TYPE: + yield {"type": "directory", "name": item_path} + else: + yield { + "type": "file", + "name": item_path, + "size": item["fileSize"], + "checksum": item["md5Checksum"], + } + else: + yield item_path + + if not cached: + self._cache_path_id(root_path, *dir_ids) def walk_files(self, path_info, **kwargs): - use_prefix = kwargs.pop("prefix", False) - if path_info == self.path_info or not use_prefix: - prefix = None - else: - prefix = path_info.path - for fname in self._list_paths(prefix=prefix, **kwargs): - yield path_info.replace(fname) + for filename in self.ls(path_info, recursive=True): + yield path_info.replace(path=filename) def remove(self, path_info): item_id = self._get_item_id(path_info) diff --git a/setup.py b/setup.py index 308c9f798..57a3867af 100644 --- a/setup.py +++ b/setup.py @@ -152,6 +152,7 @@ tests_requirements = [ "mypy", "wsgidav", "crc32c", + "gdrivefs @ git+https://github.com/intake/gdrivefs.git", ] setup( diff --git a/tests/func/test_tree.py b/tests/func/test_tree.py index d66223f43..88471761d 100644 --- a/tests/func/test_tree.py +++ b/tests/func/test_tree.py @@ -281,6 +281,34 @@ def test_tree_upload_fobj(dvc, tmp_dir, cloud): assert stream.read() == b"foo" +@pytest.mark.parametrize("cloud", [pytest.lazy_fixture("gdrive")]) +def test_tree_ls(dvc, cloud): + cloud.gen( + { + "directory": { + "foo": "foo", + "bar": "bar", + "baz": {"quux": "quux", "egg": {"foo": "foo"}}, + "empty": {}, + } + } + ) + tree = get_cloud_tree(dvc, **cloud.config) + path_info = cloud / "directory" + + assert {os.path.basename(file_key) for file_key in tree.ls(path_info)} == { + "foo", + "bar", + "baz", + "empty", + } + assert set(tree.ls(path_info / "empty")) == set() + assert { + (detail["type"], os.path.basename(detail["name"])) + for detail in tree.ls(path_info / "baz", detail=True) + } == {("file", "quux"), ("directory", "egg")} + + @pytest.mark.parametrize( "cloud", [ @@ -288,9 +316,10 @@ def test_tree_upload_fobj(dvc, tmp_dir, cloud): pytest.lazy_fixture("azure"), pytest.lazy_fixture("gs"), pytest.lazy_fixture("webdav"), + pytest.lazy_fixture("gdrive"), ], ) -def test_tree_ls(dvc, cloud): +def test_tree_ls_recursive(dvc, cloud): cloud.gen({"data": {"foo": "foo", "bar": {"baz": "baz"}, "quux": "quux"}}) tree = get_cloud_tree(dvc, **cloud.config) path_info = tree.path_info diff --git a/tests/remotes/gdrive.py b/tests/remotes/gdrive.py index 21e755991..ef0eba8d3 100644 --- a/tests/remotes/gdrive.py +++ b/tests/remotes/gdrive.py @@ -1,17 +1,24 @@ # pylint:disable=abstract-method import os import uuid +from functools import partial, partialmethod import pytest from funcy import cached_property from dvc.path_info import CloudURLInfo from dvc.tree.gdrive import GDriveTree +from dvc.utils import tmp_fname from .base import Base TEST_GDRIVE_REPO_BUCKET = "root" +CREDENTIALS_DIR = os.path.join(os.path.expanduser("~"), ".config", "pydata") +CREDENTIALS_FILE = os.path.join( + CREDENTIALS_DIR, "pydata_google_credentials.json" +) + class GDrive(Base, CloudURLInfo): @staticmethod @@ -36,9 +43,43 @@ class GDrive(Base, CloudURLInfo): # NOTE: `get_url` should always return new random url return "gdrive://" + GDrive._get_storagepath() + @cached_property + def client(self): + import pydata_google_auth + from gdrivefs import GoogleDriveFileSystem + + tmp_path = tmp_fname() + with open(tmp_path, "w") as stream: + stream.write(os.getenv(GDriveTree.GDRIVE_CREDENTIALS_DATA)) + + GoogleDriveFileSystem._connect_cache = partial( + pydata_google_auth.load_user_credentials, tmp_path + ) + return GoogleDriveFileSystem(token="cache") + + def mkdir(self, mode=0o777, parents=False, exist_ok=False): + if not self.client.exists(self.path): + self.client.mkdir(self.path) + + def write_bytes(self, contents): + with self.client.open(self.path, mode="wb") as stream: + stream.write(contents) + + def _read(self, mode): + with self.client.open(self.path, mode=mode) as stream: + return stream.read() + + def cleanup(self): + self.client.delete(self.path, recursive=True) + + read_text = partialmethod(_read, mode="r") + read_bytes = partialmethod(_read, mode="rb") + @pytest.fixture def gdrive(test_config, make_tmp_dir): + pytest.skip("temporarily disabled") + test_config.requires("gdrive") if not GDrive.should_test(): pytest.skip("no gdrive") @@ -49,4 +90,4 @@ def gdrive(test_config, make_tmp_dir): ret = GDrive(GDrive.get_url()) tree = GDriveTree(tmp_dir.dvc, ret.config) tree._gdrive_create_dir("root", tree.path_info.path) - return ret + yield ret diff --git a/tests/unit/remote/test_gdrive.py b/tests/unit/remote/test_gdrive.py index 9c159aa36..e03c5eac5 100644 --- a/tests/unit/remote/test_gdrive.py +++ b/tests/unit/remote/test_gdrive.py @@ -1,5 +1,3 @@ -import os - import pytest from dvc.tree.gdrive import GDriveAuthError, GDriveTree @@ -20,18 +18,18 @@ class TestRemoteGDrive: tree = GDriveTree(dvc, self.CONFIG) assert str(tree.path_info) == self.CONFIG["url"] - def test_drive(self, dvc): + def test_drive(self, dvc, monkeypatch): tree = GDriveTree(dvc, self.CONFIG) - os.environ[ - GDriveTree.GDRIVE_CREDENTIALS_DATA - ] = USER_CREDS_TOKEN_REFRESH_ERROR + monkeypatch.setenv( + GDriveTree.GDRIVE_CREDENTIALS_DATA, USER_CREDS_TOKEN_REFRESH_ERROR + ) with pytest.raises(GDriveAuthError): assert tree._drive - os.environ[GDriveTree.GDRIVE_CREDENTIALS_DATA] = "" + monkeypatch.setenv(GDriveTree.GDRIVE_CREDENTIALS_DATA, "") tree = GDriveTree(dvc, self.CONFIG) - os.environ[ - GDriveTree.GDRIVE_CREDENTIALS_DATA - ] = USER_CREDS_MISSED_KEY_ERROR + monkeypatch.setenv( + GDriveTree.GDRIVE_CREDENTIALS_DATA, USER_CREDS_MISSED_KEY_ERROR + ) with pytest.raises(GDriveAuthError): assert tree._drive