From ce554b6ae78fa8ee939460b6c8bcafd92c6b6b28 Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Thu, 4 Feb 2021 13:10:42 +0300 Subject: [PATCH] tree: Introduce the new .ls() API (#5275) * s3: faster directory hashing * Introduce Tree.info() API * tests for s3 * Implement support for azure * New .ls() API * remove _iter_hashes from other trees * assert recursive * yield strings instead of file_infos * Add GSTree.ls() * Add DETAIL_FIELDS to the GSTree * WebDAVTree.ls() * proper webdav root directory initalization * disable arguments-differ --- dvc/tree/azure.py | 23 +++++++++++++++++ dvc/tree/base.py | 53 ++++++++++++++++++++++++++-------------- dvc/tree/gs.py | 19 ++++++++++++++ dvc/tree/s3.py | 18 ++++++++++++++ dvc/tree/webdav.py | 31 +++++++++++++++++++++++ tests/func/test_tree.py | 41 +++++++++++++++++++++++++++++++ tests/remotes/azure.py | 1 - tests/remotes/base.py | 4 +-- tests/remotes/gs.py | 1 - tests/remotes/hdfs.py | 1 - tests/remotes/http.py | 1 - tests/remotes/s3.py | 1 - tests/remotes/ssh.py | 1 - tests/remotes/webdav.py | 29 ++++++++++++++-------- tests/remotes/webhdfs.py | 1 - 15 files changed, 188 insertions(+), 37 deletions(-) diff --git a/dvc/tree/azure.py b/dvc/tree/azure.py index de4f6bca6..f4ee759ac 100644 --- a/dvc/tree/azure.py +++ b/dvc/tree/azure.py @@ -23,6 +23,8 @@ class AzureTree(BaseTree): "knack": "knack", } PARAM_CHECKSUM = "etag" + DETAIL_FIELDS = frozenset(("etag", "size")) + COPY_POLL_SECONDS = 5 LIST_OBJECT_PAGE_SIZE = 5000 @@ -149,6 +151,27 @@ class AzureTree(BaseTree): yield path_info.replace(path=fname) + def ls( + self, path_info, detail=False, recursive=False + ): # pylint: disable=arguments-differ + assert recursive + + container_client = self.blob_service.get_container_client( + path_info.bucket + ) + for blob in container_client.list_blobs( + name_starts_with=path_info.path + ): + if detail: + yield { + "type": "file", + "name": blob.name, + "size": blob.size, + "etag": blob.etag, + } + else: + yield blob.name + def remove(self, path_info): if path_info.scheme != self.scheme: raise NotImplementedError diff --git a/dvc/tree/base.py b/dvc/tree/base.py index d6cb93bd4..cebddfcee 100644 --- a/dvc/tree/base.py +++ b/dvc/tree/base.py @@ -2,13 +2,14 @@ import logging from concurrent.futures import ThreadPoolExecutor, as_completed from functools import partial from multiprocessing import cpu_count -from typing import Any, ClassVar, Dict, Optional +from typing import Any, ClassVar, Dict, FrozenSet, Optional from urllib.parse import urlparse from funcy import cached_property, decorator from dvc.dir_info import DirInfo from dvc.exceptions import DvcException, DvcIgnoreInCollectedDirError +from dvc.hash_info import HashInfo from dvc.ignore import DvcIgnore from dvc.path_info import URLInfo from dvc.progress import Tqdm @@ -63,6 +64,7 @@ class BaseTree: CHUNK_SIZE = 64 * 1024 * 1024 # 64 MiB PARAM_CHECKSUM: ClassVar[Optional[str]] = None + DETAIL_FIELDS: FrozenSet[str] = frozenset() state = StateNoop() @@ -202,6 +204,9 @@ class BaseTree: """ raise NotImplementedError + def ls(self, path_info, detail=False, **kwargs): + raise RemoteActionNotImplemented("ls", self.scheme) + def is_empty(self, path_info): return False @@ -287,7 +292,6 @@ class BaseTree: return self.path_info / hash_[0:2] / hash_[2:] def _calculate_hashes(self, file_infos): - file_infos = list(file_infos) with Tqdm( total=len(file_infos), unit="md5", @@ -298,24 +302,37 @@ class BaseTree: hash_infos = executor.map(worker, file_infos) return dict(zip(file_infos, hash_infos)) + def _iter_hashes(self, path_info, **kwargs): + if self.PARAM_CHECKSUM in self.DETAIL_FIELDS: + for details in self.ls(path_info, recursive=True, detail=True): + file_info = path_info.replace(path=details["name"]) + hash_info = HashInfo( + self.PARAM_CHECKSUM, + details[self.PARAM_CHECKSUM], + size=details.get("size"), + ) + yield file_info, hash_info + + return None + + file_infos = [] + for file_info in self.walk_files(path_info, **kwargs): + hash_info = self.state.get( # pylint: disable=assignment-from-none + file_info + ) + if not hash_info: + file_infos.append(file_info) + continue + yield file_info, hash_info + + yield from self._calculate_hashes(file_infos).items() + def _collect_dir(self, path_info, **kwargs): - - file_infos = set() - - for fname in self.walk_files(path_info, **kwargs): - if DvcIgnore.DVCIGNORE_FILE == fname.name: - raise DvcIgnoreInCollectedDirError(fname.parent) - - file_infos.add(fname) - - hash_infos = {fi: self.state.get(fi) for fi in file_infos} - not_in_state = {fi for fi, hi in hash_infos.items() if hi is None} - - new_hash_infos = self._calculate_hashes(not_in_state) - hash_infos.update(new_hash_infos) - dir_info = DirInfo() - for fi, hi in hash_infos.items(): + for fi, hi in self._iter_hashes(path_info, **kwargs): + if DvcIgnore.DVCIGNORE_FILE == fi.name: + raise DvcIgnoreInCollectedDirError(fi.parent) + # NOTE: this is lossy transformation: # "hey\there" -> "hey/there" # "hey/there" -> "hey/there" diff --git a/dvc/tree/gs.py b/dvc/tree/gs.py index 70f478c9f..f1fca17b0 100644 --- a/dvc/tree/gs.py +++ b/dvc/tree/gs.py @@ -56,6 +56,7 @@ class GSTree(BaseTree): PATH_CLS = CloudURLInfo REQUIRES = {"google-cloud-storage": "google.cloud.storage"} PARAM_CHECKSUM = "md5" + DETAIL_FIELDS = frozenset(("md5", "size")) def __init__(self, repo, config): super().__init__(repo, config) @@ -143,6 +144,24 @@ class GSTree(BaseTree): continue yield path_info.replace(fname) + def ls( + self, path_info, detail=False, recursive=False + ): # pylint: disable=arguments-differ + assert recursive + + for blob in self.gs.bucket(path_info.bucket).list_blobs( + prefix=path_info.path + ): + if detail: + yield { + "type": "file", + "name": blob.name, + "md5": blob.md5_hash.encode().hex(), + "size": blob.size, + } + else: + yield blob.name + def remove(self, path_info): if path_info.scheme != "gs": raise NotImplementedError diff --git a/dvc/tree/s3.py b/dvc/tree/s3.py index eb62786fd..568b7e83c 100644 --- a/dvc/tree/s3.py +++ b/dvc/tree/s3.py @@ -23,6 +23,7 @@ class S3Tree(BaseTree): PATH_CLS = CloudURLInfo REQUIRES = {"boto3": "boto3"} PARAM_CHECKSUM = "etag" + DETAIL_FIELDS = frozenset(("etag", "size")) def __init__(self, repo, config): super().__init__(repo, config) @@ -227,6 +228,23 @@ class S3Tree(BaseTree): yield path_info.replace(path=fname) + def ls( + self, path_info, detail=False, recursive=False + ): # pylint: disable=arguments-differ + assert recursive + + with self._get_bucket(path_info.bucket) as bucket: + for obj_summary in bucket.objects.filter(Prefix=path_info.path): + if detail: + yield { + "type": "file", + "name": obj_summary.key, + "size": obj_summary.size, + "etag": obj_summary.e_tag.strip('"'), + } + else: + yield obj_summary.key + def remove(self, path_info): if path_info.scheme != "s3": raise NotImplementedError diff --git a/dvc/tree/webdav.py b/dvc/tree/webdav.py index 4c149d7e8..3481ae95d 100644 --- a/dvc/tree/webdav.py +++ b/dvc/tree/webdav.py @@ -44,6 +44,7 @@ class WebDAVTree(BaseTree): # pylint:disable=abstract-method CHUNK_SIZE = 2 ** 16 PARAM_CHECKSUM = "etag" + DETAIL_FIELDS = frozenset(("etag", "size")) # Constructor def __init__(self, repo, config): @@ -189,6 +190,36 @@ class WebDAVTree(BaseTree): # pylint:disable=abstract-method # Yield path info to non directory yield info + def ls( + self, path_info, detail=False, recursive=False + ): # pylint: disable=arguments-differ + dirs = deque([path_info.path]) + + while dirs: + for entry in self._client.list(dirs.pop(), get_info=True): + path = entry["path"] + if entry["isdir"]: + dirs.append(path) + continue + + if detail: + yield { + "type": "file", + "name": path, + "size": entry["size"], + "etag": entry["etag"], + } + else: + yield path + + if not recursive: + for entry in dirs: + if detail: + yield {"type": "directory", "name": entry} + else: + yield entry + return None + # Removes file/directory def remove(self, path_info): # Use webdav client clean (DELETE) method to remove file/directory diff --git a/tests/func/test_tree.py b/tests/func/test_tree.py index 63942bcf7..dc1c402d8 100644 --- a/tests/func/test_tree.py +++ b/tests/func/test_tree.py @@ -279,3 +279,44 @@ def test_tree_upload_fobj(dvc, tmp_dir, cloud): assert tree.exists(to_info) with tree.open(to_info, "rb") as stream: assert stream.read() == b"foo" + + +@pytest.mark.parametrize( + "cloud", + [ + pytest.lazy_fixture("s3"), + pytest.lazy_fixture("azure"), + pytest.lazy_fixture("gs"), + pytest.lazy_fixture("webdav"), + ], +) +def test_tree_ls(dvc, cloud): + cloud.gen({"data": {"foo": "foo", "bar": {"baz": "baz"}, "quux": "quux"}}) + tree = get_cloud_tree(dvc, **cloud.config) + path_info = tree.path_info + + assert { + os.path.basename(file_key) + for file_key in tree.ls(path_info / "data", recursive=True) + } == {"foo", "baz", "quux"} + + +@pytest.mark.parametrize( + "cloud", + [ + pytest.lazy_fixture("s3"), + pytest.lazy_fixture("azure"), + pytest.lazy_fixture("gs"), + pytest.lazy_fixture("webdav"), + ], +) +def test_tree_ls_with_etag(dvc, cloud): + cloud.gen({"data": {"foo": "foo", "bar": {"baz": "baz"}, "quux": "quux"}}) + tree = get_cloud_tree(dvc, **cloud.config) + path_info = tree.path_info + + for details in tree.ls(path_info / "data", recursive=True, detail=True): + assert ( + tree.get_file_hash(path_info.replace(path=details["name"])).value + == details[tree.PARAM_CHECKSUM] + ) diff --git a/tests/remotes/azure.py b/tests/remotes/azure.py index 2d7d845af..de26cb829 100644 --- a/tests/remotes/azure.py +++ b/tests/remotes/azure.py @@ -49,7 +49,6 @@ class Azure(Base, CloudURLInfo): def mkdir(self, mode=0o777, parents=False, exist_ok=False): assert mode == 0o777 assert parents - assert not exist_ok def write_bytes(self, contents): self.blob_client.upload_blob(contents, overwrite=True) diff --git a/tests/remotes/base.py b/tests/remotes/base.py index 71b0fad67..6e27c1d18 100644 --- a/tests/remotes/base.py +++ b/tests/remotes/base.py @@ -36,11 +36,11 @@ class Base(URLInfo): if isinstance(contents, dict): if not contents: - path.mkdir(parents=True) + path.mkdir(parents=True, exist_ok=True) else: self._gen(contents, prefix=path) else: - path.parent.mkdir(parents=True) + path.parent.mkdir(parents=True, exist_ok=True) if isinstance(contents, bytes): path.write_bytes(contents) else: diff --git a/tests/remotes/gs.py b/tests/remotes/gs.py index 6909e5540..6870c5121 100644 --- a/tests/remotes/gs.py +++ b/tests/remotes/gs.py @@ -79,7 +79,6 @@ class GCP(Base, CloudURLInfo): def mkdir(self, mode=0o777, parents=False, exist_ok=False): assert mode == 0o777 assert parents - assert not exist_ok def write_bytes(self, contents): assert isinstance(contents, bytes) diff --git a/tests/remotes/hdfs.py b/tests/remotes/hdfs.py index cb1c5af22..c49da7703 100644 --- a/tests/remotes/hdfs.py +++ b/tests/remotes/hdfs.py @@ -48,7 +48,6 @@ class HDFS(Base, URLInfo): # pylint: disable=abstract-method def mkdir(self, mode=0o777, parents=False, exist_ok=False): assert mode == 0o777 assert parents - assert not exist_ok with self._hdfs() as _hdfs: # NOTE: fs.create_dir creates parents by default diff --git a/tests/remotes/http.py b/tests/remotes/http.py index 83f80ab1d..49775b7b8 100644 --- a/tests/remotes/http.py +++ b/tests/remotes/http.py @@ -20,7 +20,6 @@ class HTTP(Base, HTTPURLInfo): def mkdir(self, mode=0o777, parents=False, exist_ok=False): assert mode == 0o777 assert parents - assert not exist_ok def write_bytes(self, contents): assert isinstance(contents, bytes) diff --git a/tests/remotes/s3.py b/tests/remotes/s3.py index 23c85cceb..e90da9611 100644 --- a/tests/remotes/s3.py +++ b/tests/remotes/s3.py @@ -74,7 +74,6 @@ class S3(Base, CloudURLInfo): def mkdir(self, mode=0o777, parents=False, exist_ok=False): assert mode == 0o777 assert parents - assert not exist_ok def write_bytes(self, contents): self._s3.put_object( diff --git a/tests/remotes/ssh.py b/tests/remotes/ssh.py index b078775e2..ebdeeea21 100644 --- a/tests/remotes/ssh.py +++ b/tests/remotes/ssh.py @@ -81,7 +81,6 @@ class SSHMocked(Base, URLInfo): def mkdir(self, mode=0o777, parents=False, exist_ok=False): assert mode == 0o777 assert parents - assert not exist_ok with self._ssh() as _ssh: _ssh.makedirs(self.path) diff --git a/tests/remotes/webdav.py b/tests/remotes/webdav.py index 26f0b1fe6..9c4938a69 100644 --- a/tests/remotes/webdav.py +++ b/tests/remotes/webdav.py @@ -1,6 +1,8 @@ # pylint:disable=abstract-method import os +import uuid from pathlib import Path +from tempfile import TemporaryDirectory from wsgiref.simple_server import make_server import pytest @@ -8,28 +10,32 @@ from funcy import first from wsgidav.wsgidav_app import WsgiDAVApp from dvc.path_info import WebDAVURLInfo -from dvc.utils.fs import makedirs from tests.utils.httpd import run_server_on_thread from .base import Base AUTH = {"user1": {"password": "password1"}} +_WEBDAV_ROOT = TemporaryDirectory() +_WEBDAV_DIR = Path(_WEBDAV_ROOT.name) class Webdav(Base, WebDAVURLInfo): @staticmethod - def get_url(port): # pylint: disable=arguments-differ - return f"webdav://localhost:{port}" + def get_url(port, root_id): # pylint: disable=arguments-differ + return f"webdav://localhost:{port}/{root_id}/" def mkdir(self, mode=0o777, parents=False, exist_ok=False): - assert parents - makedirs(self.path, exist_ok=exist_ok) + self.dir_path.mkdir(parents=parents, exist_ok=True) def write_bytes(self, contents): - Path(self.path).write_bytes(contents) + self.dir_path.write_bytes(contents) def write_text(self, contents, encoding=None, errors=None): - Path(self.path).write_text(contents, encoding=encoding, errors=errors) + self.dir_path.write_text(contents, encoding=encoding, errors=errors) + + @property + def dir_path(self): + return _WEBDAV_DIR / self.path[1:] @pytest.fixture @@ -37,8 +43,10 @@ def webdav_server(test_config, tmp_path_factory): test_config.requires("webdav") host, port = "localhost", 0 - directory = os.fspath(tmp_path_factory.mktemp("http")) - dirmap = {"/": directory} + root_id = str(uuid.uuid4()) + root_dir = _WEBDAV_DIR / root_id + root_dir.mkdir() + dirmap = {f"/{root_id}": os.fspath(_WEBDAV_DIR / root_id)} app = WsgiDAVApp( { @@ -49,13 +57,14 @@ def webdav_server(test_config, tmp_path_factory): } ) server = make_server(host, port, app) + server.root_id = root_id with run_server_on_thread(server) as httpd: yield httpd @pytest.fixture def webdav(webdav_server): - url = Webdav.get_url(webdav_server.server_port) + url = Webdav.get_url(webdav_server.server_port, webdav_server.root_id) ret = Webdav(url) user, secrets = first(AUTH.items()) ret.config = {"url": url, "user": user, **secrets} diff --git a/tests/remotes/webhdfs.py b/tests/remotes/webhdfs.py index 0e5d97e9f..08f673fa5 100644 --- a/tests/remotes/webhdfs.py +++ b/tests/remotes/webhdfs.py @@ -35,7 +35,6 @@ class WebHDFS(Base, URLInfo): # pylint: disable=abstract-method def mkdir(self, mode=0o777, parents=False, exist_ok=False): assert mode == 0o777 assert parents - assert not exist_ok with self._webhdfs() as _hdfs: # NOTE: hdfs.makekdirs always creates parents