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 <[email protected]>

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)

setup.py

@@ -152,6 +152,7 @@ tests_requirements = [
     "mypy",
     "wsgidav",
     "crc32c",
+    "gdrivefs @ git+https://github.com/intake/gdrivefs.git",
 ]
 
 setup(

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

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

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