cache: move unrelated helpers out of CloudCache class (#5439)

* cache: move transfer to objects

Moving it as is for now.

* cache: move link() to checkout

dvc/cache/base.py

@@ -1,21 +1,14 @@
 import itertools
 import logging
-from concurrent import futures
 from concurrent.futures import ThreadPoolExecutor
 from copy import copy
 from typing import Optional
 
 from funcy import decorator
-from shortuuid import uuid
 
-from dvc.dir_info import DirInfo
-from dvc.exceptions import CacheLinkError, DvcException
 from dvc.progress import Tqdm
-from dvc.remote.slow_link_detection import (  # type: ignore[attr-defined]
-    slow_link_guard,
-)
 
-from ..objects import HashFile, ObjectFormatError, Tree
+from ..objects import HashFile, ObjectFormatError
 
 logger = logging.getLogger(__name__)
 
@@ -43,60 +36,12 @@ class CloudCache:
         )
         self.cache_type_confirmed = False
 
-    def link(self, from_info, to_info):
-        self._link(from_info, to_info, self.cache_types)
-
     def move(self, from_info, to_info):
         self.tree.move(from_info, to_info)
 
     def makedirs(self, path_info):
         self.tree.makedirs(path_info)
 
-    def _link(self, from_info, to_info, link_types):
-        assert self.tree.isfile(from_info)
-
-        self.makedirs(to_info.parent)
-
-        self._try_links(from_info, to_info, link_types)
-
-    def _verify_link(self, path_info, link_type):
-        if self.cache_type_confirmed:
-            return
-
-        is_link = getattr(self.tree, f"is_{link_type}", None)
-        if is_link and not is_link(path_info):
-            self.tree.remove(path_info)
-            raise DvcException(f"failed to verify {link_type}")
-
-        self.cache_type_confirmed = True
-
-    @slow_link_guard
-    def _try_links(self, from_info, to_info, link_types):
-        while link_types:
-            link_method = getattr(self.tree, link_types[0])
-            try:
-                self._do_link(from_info, to_info, link_method)
-                self._verify_link(to_info, link_types[0])
-                return
-
-            except DvcException as exc:
-                logger.debug(
-                    "Cache type '%s' is not supported: %s", link_types[0], exc
-                )
-                del link_types[0]
-
-        raise CacheLinkError([to_info])
-
-    def _do_link(self, from_info, to_info, link_method):
-        if self.tree.exists(to_info):
-            raise DvcException(f"Link '{to_info}' already exists!")
-
-        link_method(from_info, to_info)
-
-        logger.debug(
-            "Created '%s': %s -> %s", self.cache_types[0], from_info, to_info,
-        )
-
     def get(self, hash_info):
         """ get raw object """
         return HashFile(
@@ -129,107 +74,6 @@ class CloudCache:
         if callback:
             callback(1)
 
-    def _transfer_file(self, from_tree, from_info):
-        from dvc.utils import tmp_fname
-        from dvc.utils.stream import HashedStreamReader
-
-        tmp_info = self.tree.path_info / tmp_fname()
-        with from_tree.open(
-            from_info, mode="rb", chunk_size=from_tree.CHUNK_SIZE
-        ) as stream:
-            stream_reader = HashedStreamReader(stream)
-            # Since we don't know the hash beforehand, we'll
-            # upload it to a temporary location and then move
-            # it.
-            self.tree.upload_fobj(
-                stream_reader,
-                tmp_info,
-                total=from_tree.getsize(from_info),
-                desc=from_info.name,
-            )
-
-        hash_info = stream_reader.hash_info
-        return tmp_info, hash_info
-
-    def _transfer_directory_contents(self, from_tree, from_info, jobs, pbar):
-        rel_path_infos = {}
-        from_infos = from_tree.walk_files(from_info)
-
-        def create_tasks(executor, amount):
-            for entry_info in itertools.islice(from_infos, amount):
-                pbar.total += 1
-                task = executor.submit(
-                    pbar.wrap_fn(self._transfer_file), from_tree, entry_info
-                )
-                rel_path_infos[task] = entry_info.relative_to(from_info)
-                yield task
-
-        pbar.total = 0
-        with ThreadPoolExecutor(max_workers=jobs) as executor:
-            tasks = set(create_tasks(executor, jobs * 5))
-
-            while tasks:
-                done, tasks = futures.wait(
-                    tasks, return_when=futures.FIRST_COMPLETED
-                )
-                tasks.update(create_tasks(executor, len(done)))
-                for task in done:
-                    yield rel_path_infos.pop(task), task.result()
-
-    def _transfer_directory(
-        self, from_tree, from_info, jobs, no_progress_bar=False
-    ):
-        dir_info = DirInfo()
-
-        with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar:
-            for (
-                entry_info,
-                (entry_tmp_info, entry_hash),
-            ) in self._transfer_directory_contents(
-                from_tree, from_info, jobs, pbar
-            ):
-                self.add(entry_tmp_info, self.tree, entry_hash)
-                dir_info.trie[entry_info.parts] = entry_hash
-
-        return Tree.save_dir_info(self, dir_info)
-
-    def transfer(self, from_tree, from_info, jobs=None, no_progress_bar=False):
-        jobs = jobs or min((from_tree.jobs, self.tree.jobs))
-
-        if from_tree.isdir(from_info):
-            return self._transfer_directory(
-                from_tree,
-                from_info,
-                jobs=jobs,
-                no_progress_bar=no_progress_bar,
-            )
-        tmp_info, hash_info = self._transfer_file(from_tree, from_info)
-        self.add(tmp_info, self.tree, hash_info)
-        return hash_info
-
-    def cache_is_copy(self, path_info):
-        """Checks whether cache uses copies."""
-        if self.cache_type_confirmed:
-            return self.cache_types[0] == "copy"
-
-        if set(self.cache_types) <= {"copy"}:
-            return True
-
-        workspace_file = path_info.with_name("." + uuid())
-        test_cache_file = self.tree.path_info / ".cache_type_test_file"
-        if not self.tree.exists(test_cache_file):
-            self.makedirs(test_cache_file.parent)
-            with self.tree.open(test_cache_file, "wb") as fobj:
-                fobj.write(bytes(1))
-        try:
-            self.link(test_cache_file, workspace_file)
-        finally:
-            self.tree.remove(workspace_file)
-            self.tree.remove(test_cache_file)
-
-        self.cache_type_confirmed = True
-        return self.cache_types[0] == "copy"
-
     # Override to return path as a string instead of PathInfo for clouds
     # which support string paths (see local)
     def hash_to_path(self, hash_):

dvc/cache/local.py

@@ -55,11 +55,6 @@ class LocalCache(CloudCache):
 
         makedirs(path_info, exist_ok=True, mode=self._dir_mode)
 
-    def link(self, from_info, to_info):
-        super().link(from_info, to_info)
-        if self.cache_types[0] not in ("symlink", "hardlink"):
-            os.chmod(to_info, self._file_mode)
-
     def hash_to_path(self, hash_):
         # NOTE: `self.cache_path` is already normalized so we can simply use
         # `os.sep` instead of `os.path.join`. This results in this helper
@@ -85,12 +80,6 @@ class LocalCache(CloudCache):
 
         return ret
 
-    def _verify_link(self, path_info, link_type):
-        if link_type == "hardlink" and self.tree.getsize(path_info) == 0:
-            return
-
-        super()._verify_link(path_info, link_type)
-
     def _list_paths(self, prefix=None, progress_callback=None):
         assert self.tree.path_info is not None
         if prefix:

dvc/checkout.py

@@ -1,8 +1,18 @@
 import logging
 
+from shortuuid import uuid
+
 import dvc.prompt as prompt
-from dvc.exceptions import CheckoutError, ConfirmRemoveError
+from dvc.exceptions import (
+    CacheLinkError,
+    CheckoutError,
+    ConfirmRemoveError,
+    DvcException,
+)
 from dvc.objects import ObjectFormatError, check, load
+from dvc.remote.slow_link_detection import (  # type: ignore[attr-defined]
+    slow_link_guard,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -61,6 +71,80 @@ def _remove(path_info, tree, cache, force=False):
     tree.remove(path_info)
 
 
+def _verify_link(cache, path_info, link_type):
+    if link_type == "hardlink" and cache.tree.getsize(path_info) == 0:
+        return
+
+    if cache.cache_type_confirmed:
+        return
+
+    is_link = getattr(cache.tree, f"is_{link_type}", None)
+    if is_link and not is_link(path_info):
+        cache.tree.remove(path_info)
+        raise DvcException(f"failed to verify {link_type}")
+
+    cache.cache_type_confirmed = True
+
+
+def _do_link(cache, from_info, to_info, link_method):
+    if cache.tree.exists(to_info):
+        raise DvcException(f"Link '{to_info}' already exists!")
+
+    link_method(from_info, to_info)
+
+    logger.debug(
+        "Created '%s': %s -> %s", cache.cache_types[0], from_info, to_info,
+    )
+
+
+@slow_link_guard
+def _try_links(cache, from_info, to_info, link_types):
+    while link_types:
+        link_method = getattr(cache.tree, link_types[0])
+        try:
+            _do_link(cache, from_info, to_info, link_method)
+            _verify_link(cache, to_info, link_types[0])
+            return
+
+        except DvcException as exc:
+            logger.debug(
+                "Cache type '%s' is not supported: %s", link_types[0], exc
+            )
+            del link_types[0]
+
+    raise CacheLinkError([to_info])
+
+
+def _link(cache, from_info, to_info):
+    assert cache.tree.isfile(from_info)
+    cache.makedirs(to_info.parent)
+    _try_links(cache, from_info, to_info, cache.cache_types)
+
+
+def _cache_is_copy(cache, path_info):
+    """Checks whether cache uses copies."""
+    if cache.cache_type_confirmed:
+        return cache.cache_types[0] == "copy"
+
+    if set(cache.cache_types) <= {"copy"}:
+        return True
+
+    workspace_file = path_info.with_name("." + uuid())
+    test_cache_file = cache.tree.path_info / ".cache_type_test_file"
+    if not cache.tree.exists(test_cache_file):
+        cache.makedirs(test_cache_file.parent)
+        with cache.tree.open(test_cache_file, "wb") as fobj:
+            fobj.write(bytes(1))
+    try:
+        _link(cache, test_cache_file, workspace_file)
+    finally:
+        cache.tree.remove(workspace_file)
+        cache.tree.remove(test_cache_file)
+
+    cache.cache_type_confirmed = True
+    return cache.cache_types[0] == "copy"
+
+
 def _checkout_file(
     path_info, tree, obj, cache, force, progress_callback=None, relink=False,
 ):
@@ -71,15 +155,15 @@ def _checkout_file(
         if not relink and _changed(path_info, tree, obj, cache):
             modified = True
             _remove(path_info, tree, cache, force=force)
-            cache.link(cache_info, path_info)
+            _link(cache, cache_info, path_info)
         else:
-            if tree.iscopy(path_info) and cache.cache_is_copy(path_info):
+            if tree.iscopy(path_info) and _cache_is_copy(cache, path_info):
                 cache.unprotect(path_info)
             else:
                 _remove(path_info, tree, cache, force=force)
-                cache.link(cache_info, path_info)
+                _link(cache, cache_info, path_info)
     else:
-        cache.link(cache_info, path_info)
+        _link(cache, cache_info, path_info)
         modified = True
 
     tree.state.save(path_info, obj.hash_info)

dvc/objects.py

@@ -1,5 +1,9 @@
+import itertools
 import json
 import logging
+from concurrent import futures
+
+from dvc.progress import Tqdm
 
 from .dir_info import DirInfo
 from .exceptions import DvcException
@@ -172,8 +176,6 @@ class Tree(HashFile):
         return cls(obj.path_info, obj.tree, hash_info)
 
     def save(self, odb, **kwargs):
-        from dvc.progress import Tqdm
-
         assert self.src_hash_info.dir_info
         hi = self.save_dir_info(
             odb, self.src_hash_info.dir_info, self.hash_info
@@ -253,3 +255,86 @@ def merge(odb, ancestor_info, our_info, their_info):
     hash_info = Tree.save_dir_info(odb, merged)
     hash_info.size = _get_dir_size(odb, merged)
     return hash_info
+
+
+def _transfer_file(odb, from_tree, from_info):
+    from dvc.utils import tmp_fname
+    from dvc.utils.stream import HashedStreamReader
+
+    tmp_info = odb.tree.path_info / tmp_fname()
+    with from_tree.open(
+        from_info, mode="rb", chunk_size=from_tree.CHUNK_SIZE
+    ) as stream:
+        stream_reader = HashedStreamReader(stream)
+        # Since we don't know the hash beforehand, we'll
+        # upload it to a temporary location and then move
+        # it.
+        odb.tree.upload_fobj(
+            stream_reader,
+            tmp_info,
+            total=from_tree.getsize(from_info),
+            desc=from_info.name,
+        )
+
+    hash_info = stream_reader.hash_info
+    return tmp_info, hash_info
+
+
+def _transfer_directory_contents(odb, from_tree, from_info, jobs, pbar):
+    rel_path_infos = {}
+    from_infos = from_tree.walk_files(from_info)
+
+    def create_tasks(executor, amount):
+        for entry_info in itertools.islice(from_infos, amount):
+            pbar.total += 1
+            task = executor.submit(
+                pbar.wrap_fn(_transfer_file), odb, from_tree, entry_info
+            )
+            rel_path_infos[task] = entry_info.relative_to(from_info)
+            yield task
+
+    pbar.total = 0
+    with futures.ThreadPoolExecutor(max_workers=jobs) as executor:
+        tasks = set(create_tasks(executor, jobs * 5))
+
+        while tasks:
+            done, tasks = futures.wait(
+                tasks, return_when=futures.FIRST_COMPLETED
+            )
+            tasks.update(create_tasks(executor, len(done)))
+            for task in done:
+                yield rel_path_infos.pop(task), task.result()
+
+
+def _transfer_directory(
+    odb, from_tree, from_info, jobs, no_progress_bar=False
+):
+    dir_info = DirInfo()
+
+    with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar:
+        for (
+            entry_info,
+            (entry_tmp_info, entry_hash),
+        ) in _transfer_directory_contents(
+            odb, from_tree, from_info, jobs, pbar
+        ):
+            odb.add(entry_tmp_info, odb.tree, entry_hash)
+            dir_info.trie[entry_info.parts] = entry_hash
+
+    return Tree.save_dir_info(odb, dir_info)
+
+
+def transfer(odb, from_tree, from_info, jobs=None, no_progress_bar=False):
+    jobs = jobs or min((from_tree.jobs, odb.tree.jobs))
+
+    if from_tree.isdir(from_info):
+        return _transfer_directory(
+            odb,
+            from_tree,
+            from_info,
+            jobs=jobs,
+            no_progress_bar=no_progress_bar,
+        )
+    tmp_info, hash_info = _transfer_file(odb, from_tree, from_info)
+    odb.add(tmp_info, odb.tree, hash_info)
+    return hash_info

dvc/remote/base.py

@@ -512,8 +512,14 @@ class Remote:
         return ret
 
     def transfer(self, from_tree, from_info, jobs=None, no_progress_bar=False):
-        return self.cache.transfer(
-            from_tree, from_info, jobs=jobs, no_progress_bar=no_progress_bar
+        from dvc.objects import transfer
+
+        return transfer(
+            self.cache,
+            from_tree,
+            from_info,
+            jobs=jobs,
+            no_progress_bar=no_progress_bar,
         )
 
     @staticmethod

dvc/repo/add.py

@@ -165,11 +165,15 @@ def _process_stages(
                 command="add",
             )
         else:
+            from dvc.objects import transfer
             from dvc.tree import get_cloud_tree
 
             from_tree = get_cloud_tree(repo, url=target)
-            out.hash_info = out.cache.transfer(
-                from_tree, from_tree.path_info, jobs=kwargs.get("jobs"),
+            out.hash_info = transfer(
+                out.cache,
+                from_tree,
+                from_tree.path_info,
+                jobs=kwargs.get("jobs"),
             )
             out.checkout()
 

dvc/system.py

@@ -2,7 +2,6 @@ import errno
 import logging
 import os
 import platform
-import shutil
 import sys
 
 from dvc.exceptions import DvcException
@@ -27,10 +26,6 @@ class System:
     def is_unix():
         return os.name != "nt"
 
-    @staticmethod
-    def copy(src, dest):
-        return shutil.copyfile(src, dest)
-
     @staticmethod
     def hardlink(source, link_name):
         try:
@@ -98,6 +93,8 @@ class System:
 
     @staticmethod
     def reflink(source, link_name):
+        from dvc.utils.fs import umask
+
         source, link_name = os.fspath(source), os.fspath(link_name)
 
         system = platform.system()
@@ -116,6 +113,10 @@ class System:
         if ret != 0:
             raise DvcException("reflink is not supported")
 
+        # NOTE: reflink has a new inode, but has the same mode as the src,
+        # so we need to chmod it to look like a normal copy.
+        os.chmod(link_name, 0o666 & ~umask)
+
     @staticmethod
     def _getdirinfo(path):
         from collections import namedtuple

dvc/tree/local.py

@@ -154,7 +154,7 @@ class LocalTree(BaseTree):
     def copy(self, from_info, to_info):
         tmp_info = to_info.parent / tmp_fname("")
         try:
-            System.copy(from_info, tmp_info)
+            copyfile(from_info, tmp_info)
             os.rename(tmp_info, to_info)
         except Exception:
             self.remove(tmp_info)

dvc/utils/fs.py

@@ -195,9 +195,6 @@ def copyfile(src, dest, no_progress_bar=False, name=None):
 
     try:
         System.reflink(src, dest)
-        # NOTE: reflink has a new inode, but has the same mode as the src,
-        # so we need to chmod it to look like a normal copy.
-        os.chmod(dest, 0o666 & ~umask)
     except DvcException:
         with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
             with Tqdm.wrapattr(

tests/func/test_add.py

@@ -937,8 +937,8 @@ def test_add_file_in_symlink_dir(make_tmp_dir, tmp_dir, dvc, external):
 def test_add_with_cache_link_error(tmp_dir, dvc, mocker, caplog):
     tmp_dir.gen("foo", "foo")
 
-    mocker.patch.object(
-        dvc.cache.local, "_do_link", side_effect=DvcException("link failed")
+    mocker.patch(
+        "dvc.checkout._do_link", side_effect=DvcException("link failed")
     )
     with caplog.at_level(logging.WARNING, logger="dvc"):
         dvc.add("foo")

tests/func/test_s3.py

@@ -55,6 +55,8 @@ def test_copy_singlepart_preserve_etag():
     [S3Tree.PATH_CLS("s3://bucket/"), S3Tree.PATH_CLS("s3://bucket/ns/")],
 )
 def test_link_created_on_non_nested_path(base_info, tmp_dir, dvc, scm):
+    from dvc.checkout import _link
+
     tree = S3Tree(dvc, {"url": str(base_info.parent)})
     cache = CloudCache(tree)
     s3 = cache.tree.s3.meta.client
@@ -62,7 +64,7 @@ def test_link_created_on_non_nested_path(base_info, tmp_dir, dvc, scm):
     s3.put_object(
         Bucket=base_info.bucket, Key=(base_info / "from").path, Body="data"
     )
-    cache.link(base_info / "from", base_info / "to")
+    _link(cache, base_info / "from", base_info / "to")
 
     assert cache.tree.exists(base_info / "from")
     assert cache.tree.exists(base_info / "to")