objects: use object IDs and references instead of naive objs in status/transfer (#6360)

* objects.status: use object IDs instead of naive objects

* objects.transfer: use object IDs instead of naive objects

* objects: add ReferenceHashFile

* odb: add ReferenceObjectDB

* repo: collect used object IDs instead of naive objs

* support basic/naive fs serialization

* objects.stage: stage file objects in addition to trees

* objects: remove save(), update transfer()

* output/dep: update for osave removal

* path_info: support both url and path attributes

* status/gc: use object IDs instead of naive objs

* diff: update staging usage

* commit: update staging usage

* tests: clean staging when needed

* tests: update for save/transfer/stage changes

* dvcfs: don't allow config serialization (repofs should be used instead)

* utils: return mtime as str, size as int

* tests: update unit tests for save/stage/transfer changes

* objects: use Repo.open when needed in ReferenceHashFile

* fix ls test staging

* objects.stage: don't skip hash check when loading from state

- individual hashes will still be loaded from state if they are still
  valid

* ReferenceHashFile: serialize repo fs paths relative to repo root

* preserve state when staging

* objects.tree: skip hash computation in digest() when we have already
loaded hash from state

* always check file existence in referencehashfile

* update repofs test since stage no longer allows nonexistent files

* odb: only protect on check() if hash was verified

* staging: use unique memfs namespace for staging per dest ODB

* tests: clean full staging namespace

* refodb: preserve repofs cache dir

* repodependency: stage imports to local ODB and strict check circular
imports

* fetch: fetch staging last

* staging: make dest odb mandatory and preserve state

* refodb: cache nonlocal fs's

* fix rebase errors

* repofs: support full instantiation from fs config

* repofs: preserve subrepo factory for erepos

* tests: clarify repofs dirty file behavior

dvc/objects/__init__.py

@@ -1,8 +1,5 @@
 import logging
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import TYPE_CHECKING, Optional
-
-from dvc.progress import Tqdm
+from typing import TYPE_CHECKING
 
 from .tree import Tree
 
@@ -15,50 +12,12 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-def save(
-    odb: "ObjectDB",
-    obj: "HashFile",
-    jobs: Optional[int] = None,
-    **kwargs,
-):
-    if isinstance(obj, Tree):
-        from .stage import get_staging
-
-        with ThreadPoolExecutor(max_workers=jobs) as executor:
-            for future in Tqdm(
-                as_completed(
-                    executor.submit(
-                        odb.add,
-                        entry.path_info,
-                        entry.fs,
-                        entry.hash_info,
-                        **kwargs,
-                    )
-                    for _, entry in obj
-                ),
-                desc="Saving files",
-                total=len(obj),
-                unit="file",
-            ):
-                future.result()
-
-            # if dir cache for this tree has already been staged, move the
-            # staged object when possible
-            staging = get_staging(odb)
-            if staging.exists(obj.hash_info):
-                obj = staging.get(obj.hash_info)
-                kwargs["move"] = True
-
-    assert obj.path_info and obj.fs and obj.hash_info
-    odb.add(obj.path_info, obj.fs, obj.hash_info, **kwargs)
-
-
-def check(odb: "ObjectDB", obj: "HashFile"):
+def check(odb: "ObjectDB", obj: "HashFile", **kwargs):
     if isinstance(obj, Tree):
         for _, entry in obj:
-            odb.check(entry.hash_info)
+            odb.check(entry.hash_info, **kwargs)
 
-    odb.check(obj.hash_info)
+    odb.check(obj.hash_info, **kwargs)
 
 
 def load(odb: "ObjectDB", hash_info: "HashInfo") -> "HashFile":

dvc/objects/db/base.py

@@ -13,7 +13,7 @@ from dvc.progress import Tqdm
 if TYPE_CHECKING:
     from dvc.fs.base import BaseFileSystem
     from dvc.hash_info import HashInfo
-    from dvc.types import AnyPath
+    from dvc.types import AnyPath, DvcPath
 
 logger = logging.getLogger(__name__)
 
@@ -76,7 +76,14 @@ class ObjectDB:
             hash_info,
         )
 
-    def _add_file(self, from_fs, from_info, to_info, move):
+    def _add_file(
+        self,
+        from_fs: "BaseFileSystem",
+        from_info: "AnyPath",
+        to_info: "DvcPath",
+        _hash_info: "HashInfo",
+        move: bool = False,
+    ):
         self.makedirs(to_info.parent)
         use_move = isinstance(from_fs, type(self.fs)) and move
         try:
@@ -131,8 +138,8 @@ class ObjectDB:
             pass
 
         cache_info = self.hash_to_path_info(hash_info.value)
-        # using our makedirs to create dirs with proper permissions
-        self._add_file(fs, path_info, cache_info, move)
+        self._add_file(fs, path_info, cache_info, hash_info, move)
+
         try:
             if verify:
                 self.check(hash_info, check_hash=True)
@@ -145,7 +152,7 @@ class ObjectDB:
         if callback:
             callback(1)
 
-    def hash_to_path_info(self, hash_):
+    def hash_to_path_info(self, hash_) -> "DvcPath":
         return self.path_info / hash_[0:2] / hash_[2:]
 
     # Override to return path as a string instead of PathInfo for clouds
@@ -199,9 +206,10 @@ class ObjectDB:
             self.fs.remove(obj.path_info)
             raise
 
-        # making cache file read-only so we don't need to check it
-        # next time
-        self.protect(obj.path_info)
+        if check_hash:
+            # making cache file read-only so we don't need to check it
+            # next time
+            self.protect(obj.path_info)
 
     def _list_paths(self, prefix=None, progress_callback=None):
         if prefix:
@@ -378,17 +386,20 @@ class ObjectDB:
     def _remove_unpacked_dir(self, hash_):
         pass
 
-    def gc(self, used, jobs=None):
+    def gc(self, used, jobs=None, cache_odb=None, shallow=True):
         from ..tree import Tree
 
         if self.read_only:
             raise ObjectDBPermissionError("Cannot gc read-only ODB")
+        if not cache_odb:
+            cache_odb = self
         used_hashes = set()
-        for obj in used:
-            used_hashes.add(obj.hash_info.value)
-            if isinstance(obj, Tree):
+        for hash_info in used:
+            used_hashes.add(hash_info.value)
+            if hash_info.isdir and not shallow:
+                tree = Tree.load(cache_odb, hash_info)
                 used_hashes.update(
-                    entry_obj.hash_info.value for _, entry_obj in obj
+                    entry_obj.hash_info.value for _, entry_obj in tree
                 )
 
         removed = False

dvc/objects/db/reference.py

@@ -0,0 +1,77 @@
+import io
+import logging
+import os
+from typing import TYPE_CHECKING, Dict
+
+from dvc.scheme import Schemes
+
+from ..errors import ObjectFormatError
+from ..reference import ReferenceHashFile
+from .base import ObjectDB
+
+if TYPE_CHECKING:
+    from dvc.fs.base import BaseFileSystem
+    from dvc.hash_info import HashInfo
+    from dvc.types import AnyPath, DvcPath
+
+logger = logging.getLogger(__name__)
+
+
+class ReferenceObjectDB(ObjectDB):
+    """Reference ODB.
+
+    File objects are stored as ReferenceHashFiles which reference paths outside
+    of the staging ODB fs. Tree objects are stored natively.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._fs_cache: Dict[tuple, "BaseFileSystem"] = {}
+
+    def get(self, hash_info: "HashInfo"):
+        if hash_info.isdir:
+            return super().get(hash_info)
+        path_info = self.hash_to_path(hash_info.value)
+        try:
+            with self.fs.open(path_info, "rb") as fobj:
+                ref_file = ReferenceHashFile.from_bytes(
+                    fobj.read(), fs_cache=self._fs_cache
+                )
+        except OSError:
+            raise FileNotFoundError
+        try:
+            ref_file.check(self, check_hash=False)
+        except ObjectFormatError:
+            self.fs.remove(path_info)
+            raise
+        return ref_file
+
+    def _add_file(
+        self,
+        from_fs: "BaseFileSystem",
+        from_info: "AnyPath",
+        to_info: "DvcPath",
+        hash_info: "HashInfo",
+        move: bool = False,
+    ):
+        self.makedirs(to_info.parent)
+        if hash_info.isdir:
+            return super()._add_file(
+                from_fs, from_info, to_info, hash_info, move
+            )
+        ref_file = ReferenceHashFile(from_info, from_fs, hash_info)
+        ref_fobj = io.BytesIO(ref_file.to_bytes())
+        ref_fobj.seek(0)
+        try:
+            self.fs.upload_fobj(ref_fobj, to_info)
+        except OSError as exc:
+            if isinstance(exc, FileExistsError) or (
+                os.name == "nt"
+                and exc.__context__
+                and isinstance(exc.__context__, FileExistsError)
+            ):
+                logger.debug("'%s' file already exists, skipping", to_info)
+            else:
+                raise
+        if from_fs.scheme != Schemes.LOCAL:
+            self._fs_cache[ReferenceHashFile.config_tuple(from_fs)] = from_fs

dvc/objects/file.py

@@ -8,7 +8,9 @@ from .errors import ObjectFormatError
 if TYPE_CHECKING:
     from dvc.fs.base import BaseFileSystem
     from dvc.hash_info import HashInfo
-    from dvc.types import DvcPath
+    from dvc.types import AnyPath
+
+    from .db.base import ObjectDB
 
 logger = logging.getLogger(__name__)
 
@@ -16,7 +18,7 @@ logger = logging.getLogger(__name__)
 class HashFile:
     def __init__(
         self,
-        path_info: Optional["DvcPath"],
+        path_info: Optional["AnyPath"],
         fs: Optional["BaseFileSystem"],
         hash_info: "HashInfo",
         name: Optional[str] = None,
@@ -57,10 +59,9 @@ class HashFile:
             )
         )
 
-    def check(self, odb, check_hash=True):
-        from .stage import get_file_hash
-
+    def check(self, odb: "ObjectDB", check_hash: bool = True):
         if not check_hash:
+            assert self.fs
             if not self.fs.exists(self.path_info):
                 raise FileNotFoundError(
                     errno.ENOENT, os.strerror(errno.ENOENT), self.path_info
@@ -68,6 +69,11 @@ class HashFile:
             else:
                 return None
 
+        self._check_hash(odb)
+
+    def _check_hash(self, odb):
+        from .stage import get_file_hash
+
         actual = get_file_hash(
             self.path_info, self.fs, self.hash_info.name, odb.state
         )

dvc/objects/reference.py

@@ -0,0 +1,135 @@
+import errno
+import logging
+import os
+import pickle
+from typing import TYPE_CHECKING, Optional, Tuple
+
+from .errors import ObjectFormatError
+from .file import HashFile
+
+if TYPE_CHECKING:
+    from dvc.fs.base import BaseFileSystem
+    from dvc.hash_info import HashInfo
+    from dvc.types import AnyPath
+
+    from .db.base import ObjectDB
+
+logger = logging.getLogger(__name__)
+
+
+class ReferenceHashFile(HashFile):
+    PARAM_PATH = "path"
+    PARAM_HASH = "hash"
+    PARAM_MTIME = "mtime"
+    PARAM_SIZE = "size"
+    PARAM_FS_CONFIG = "fs_config"
+
+    def __init__(
+        self,
+        path_info: "AnyPath",
+        fs: "BaseFileSystem",
+        hash_info: "HashInfo",
+        mtime: Optional[str] = None,
+        size: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(path_info, fs, hash_info, **kwargs)
+        cur_mtime, cur_size = self._get_mtime_and_size()
+        self.mtime = cur_mtime if mtime is None else mtime
+        self.hash_info.size = cur_size if size is None else size
+
+    def __str__(self):
+        return f"ref object {self.hash_info} -> {self.path_info}"
+
+    def check(self, odb: "ObjectDB", check_hash: bool = True):
+        assert self.fs
+        if not self.fs.exists(self.path_info):
+            raise FileNotFoundError(
+                errno.ENOENT, os.strerror(errno.ENOENT), self.path_info
+            )
+        if not (self.mtime, self.size) == self._get_mtime_and_size():
+            raise ObjectFormatError(f"{self} is changed")
+        if check_hash:
+            self._check_hash(odb)
+
+    def _get_mtime_and_size(self) -> Tuple[Optional[str], int]:
+        from dvc.utils.fs import get_mtime_and_size
+
+        assert self.fs
+        if hasattr(self.fs, "stat"):
+            try:
+                return get_mtime_and_size(self.path_info, self.fs)
+            except FileNotFoundError:
+                pass
+        return None, self.fs.getsize(self.path_info)
+
+    def to_bytes(self):
+        from dvc.fs.repo import RepoFileSystem
+        from dvc.path_info import PathInfo
+
+        # NOTE: dumping reference FS's this way is insecure, as the
+        # fully parsed remote FS config will include credentials
+        #
+        # ReferenceHashFiles should currently only be serialized in
+        # memory and not to disk
+        path_info = self.path_info
+        if isinstance(self.fs, RepoFileSystem):
+            path_info = PathInfo(path_info).relative_to(self.fs.root_dir)
+
+        dict_ = {
+            self.PARAM_PATH: path_info,
+            self.PARAM_HASH: self.hash_info,
+            self.PARAM_MTIME: self.mtime,
+            self.PARAM_SIZE: self.size,
+            self.PARAM_FS_CONFIG: self.config_tuple(self.fs),
+        }
+        try:
+            return pickle.dumps(dict_)
+        except pickle.PickleError as exc:
+            raise ObjectFormatError(f"Could not pickle {self}") from exc
+
+    @classmethod
+    def from_bytes(cls, data: bytes, fs_cache: Optional[dict] = None):
+        from dvc.fs import get_fs_cls
+        from dvc.fs.repo import RepoFileSystem
+
+        try:
+            dict_ = pickle.loads(data)
+        except pickle.PickleError as exc:
+            raise ObjectFormatError("ReferenceHashFile is corrupted") from exc
+
+        try:
+            path_info = dict_[cls.PARAM_PATH]
+            hash_info = dict_[cls.PARAM_HASH]
+        except KeyError as exc:
+            raise ObjectFormatError("ReferenceHashFile is corrupted") from exc
+
+        scheme, config_pairs = dict_.get(cls.PARAM_FS_CONFIG)
+        fs = fs_cache.get((scheme, config_pairs)) if fs_cache else None
+        if not fs:
+            config = dict(config_pairs)
+            if RepoFileSystem.PARAM_REPO_URL in config:
+                fs = RepoFileSystem(**config)
+                path_info = fs.root_dir / path_info
+            else:
+                fs_cls = get_fs_cls(config, scheme=path_info.scheme)
+                fs = fs_cls(**config)
+        return ReferenceHashFile(
+            path_info,
+            fs,
+            hash_info,
+            mtime=dict_.get(cls.PARAM_MTIME),
+            size=dict_.get(cls.PARAM_SIZE),
+        )
+
+    @staticmethod
+    def config_tuple(fs: "BaseFileSystem"):
+        return (
+            fs.scheme,
+            tuple(
+                (key, value)
+                for key, value in sorted(
+                    fs.config.items(), key=lambda item: item[0]
+                )
+            ),
+        )

dvc/objects/stage.py

@@ -1,6 +1,10 @@
+import hashlib
+import logging
+import os
+import pathlib
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Dict, Tuple
 
 from dvc.exceptions import DvcIgnoreInCollectedDirError
 from dvc.hash_info import HashInfo
@@ -8,32 +12,34 @@ from dvc.ignore import DvcIgnore
 from dvc.progress import Tqdm
 from dvc.utils import file_md5
 
+from .db.reference import ReferenceObjectDB
 from .file import HashFile
 
 if TYPE_CHECKING:
     from dvc.fs.base import BaseFileSystem
-    from dvc.types import DvcPath
+    from dvc.types import AnyPath, DvcPath
 
     from .db.base import ObjectDB
 
-
-_STAGING_DIR = "staging"
-_STAGING_MEMFS_PATH = f"dvc-{_STAGING_DIR}"
+logger = logging.getLogger(__name__)
 
 
-def _upload_file(path_info, fs, odb):
+_STAGING_MEMFS_PATH = "dvc-staging"
+
+
+def _upload_file(path_info, fs, odb, upload_odb):
     from dvc.utils import tmp_fname
     from dvc.utils.stream import HashedStreamReader
 
-    tmp_info = odb.path_info / tmp_fname()
+    tmp_info = upload_odb.path_info / tmp_fname()
     with fs.open(path_info, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream:
         stream = HashedStreamReader(stream)
-        odb.fs.upload_fobj(
+        upload_odb.fs.upload_fobj(
             stream, tmp_info, desc=path_info.name, size=fs.getsize(path_info)
         )
 
-    obj = HashFile(tmp_info, odb.fs, stream.hash_info)
-    return path_info, obj
+    odb.add(tmp_info, upload_odb.fs, stream.hash_info)
+    return path_info, odb.get(stream.hash_info)
 
 
 def _get_file_hash(path_info, fs, name):
@@ -71,15 +77,20 @@ def get_file_hash(path_info, fs, name, state=None):
     return hash_info
 
 
-def _get_file_obj(path_info, fs, name, odb=None, upload=False):
-    if upload:
-        assert odb and name == "md5"
-        return _upload_file(path_info, fs, odb)
-
+def _get_file_obj(
+    path_info, fs, name, odb=None, upload_odb=None, dry_run=False
+):
     state = odb.state if odb else None
-    obj = HashFile(
-        path_info, fs, get_file_hash(path_info, fs, name, state=state)
-    )
+    hash_info = get_file_hash(path_info, fs, name, state=state)
+    if upload_odb and not dry_run:
+        assert odb and name == "md5"
+        return _upload_file(path_info, fs, odb, upload_odb)
+
+    if dry_run:
+        obj = HashFile(path_info, fs, hash_info)
+    else:
+        odb.add(path_info, fs, hash_info, move=False)
+        obj = odb.get(hash_info)
     return path_info, obj
 
 
@@ -87,9 +98,9 @@ def _build_objects(
     path_info,
     fs,
     name,
-    odb=None,
-    upload=False,
     dvcignore=None,
+    jobs=None,
+    no_progress_bar=False,
     **kwargs,
 ):
     if dvcignore:
@@ -99,33 +110,24 @@ def _build_objects(
     with Tqdm(
         unit="md5",
         desc="Computing file/dir hashes (only done once)",
-        disable=kwargs.pop("no_progress_bar", False),
+        disable=no_progress_bar,
     ) as pbar:
         worker = pbar.wrap_fn(
             partial(
                 _get_file_obj,
                 fs=fs,
                 name=name,
-                odb=odb,
-                upload=upload,
+                **kwargs,
             )
         )
         with ThreadPoolExecutor(
-            max_workers=kwargs.pop("jobs", fs.hash_jobs)
+            max_workers=jobs if jobs is not None else fs.hash_jobs
         ) as executor:
             yield from executor.map(worker, walk_iterator)
 
 
-def _iter_objects(path_info, fs, name, upload=False, **kwargs):
-    if not upload and name in fs.DETAIL_FIELDS:
-        for details in fs.find(path_info, detail=True):
-            file_info = path_info.replace(path=details["name"])
-            hash_info = HashInfo(name, details[name], size=details.get("size"))
-            yield file_info, HashFile(file_info, fs, hash_info)
-
-        return None
-
-    yield from _build_objects(path_info, fs, name, upload=upload, **kwargs)
+def _iter_objects(path_info, fs, name, **kwargs):
+    yield from _build_objects(path_info, fs, name, **kwargs)
 
 
 def _build_tree(path_info, fs, name, **kwargs):
@@ -145,14 +147,13 @@ def _build_tree(path_info, fs, name, **kwargs):
         # Yes, this is a BUG, as long as we permit "/" in
         # filenames on Windows and "\" on Unix
         tree.add(file_info.relative_to(path_info).parts, obj)
-    tree.digest()
     return tree
 
 
-def _get_tree_obj(path_info, fs, name, odb=None, **kwargs):
+def _get_tree_obj(path_info, fs, fs_info, name, odb=None, **kwargs):
     from .tree import Tree
 
-    value = fs.info(path_info).get(name)
+    value = fs_info.get(name)
     if odb and value:
         hash_info = HashInfo(name, value)
         try:
@@ -168,40 +169,59 @@ def _get_tree_obj(path_info, fs, name, odb=None, **kwargs):
             pass
 
     tree = _build_tree(path_info, fs, name, odb=odb, **kwargs)
+    state = odb.state if odb and odb.state else None
+    hash_info = None
+    if state:
+        hash_info = state.get(  # pylint: disable=assignment-from-none
+            path_info, fs
+        )
+    tree.digest(hash_info=hash_info)
+    odb.add(tree.path_info, tree.fs, tree.hash_info, move=True)
+    raw = odb.get(tree.hash_info)
+    # cleanup unneeded memfs tmpfile and return tree based on the
+    # ODB fs/path
+    if odb.fs != tree.fs:
+        tree.fs.remove(tree.path_info)
+    tree.fs = raw.fs
+    tree.path_info = raw.path_info
     return tree
 
 
-def get_staging(odb: Optional["ObjectDB"] = None) -> "ObjectDB":
+_url_cache: Dict[str, str] = {}
+
+
+def _make_staging_url(path_info: "AnyPath"):
+    from dvc.path_info import CloudURLInfo
+    from dvc.scheme import Schemes
+
+    url = CloudURLInfo(f"{Schemes.MEMORY}://{_STAGING_MEMFS_PATH}")
+    if path_info:
+        if isinstance(path_info, (str, pathlib.PurePath)):
+            path = os.path.abspath(path_info)
+        else:
+            path = str(path_info)
+        if path not in _url_cache:
+            _url_cache[path] = hashlib.sha256(path.encode("utf-8")).hexdigest()
+        url /= _url_cache[path]
+    return url
+
+
+def _get_staging(odb: "ObjectDB") -> "ObjectDB":
     """Return an ODB that can be used for staging objects.
 
-    If odb.fs is local, .dvc/tmp/staging will be returned. Otherwise
-    the the global (temporary) memfs ODB will be returned.
+    Staging will be a reference ODB stored in the the global memfs.
     """
 
     from dvc.fs.memory import MemoryFileSystem
-    from dvc.path_info import CloudURLInfo, PathInfo
-    from dvc.scheme import Schemes
 
-    from .db import get_odb
-
-    if odb and odb.fs.scheme == Schemes.LOCAL and odb.tmp_dir:
-        fs = odb.fs
-        path_info: "DvcPath" = PathInfo(odb.tmp_dir) / _STAGING_DIR
-        config = odb.config
-    else:
-        fs = MemoryFileSystem()
-        path_info = CloudURLInfo(f"{Schemes.MEMORY}://{_STAGING_MEMFS_PATH}")
-        config = {}
-    return get_odb(fs, path_info, **config)
-
-
-def is_memfs_staging(odb: "ObjectDB"):
-    staging = get_staging()
-    return odb.fs == staging.fs and odb.path_info == staging.path_info
+    fs = MemoryFileSystem()
+    path_info = _make_staging_url(odb.path_info)
+    state = odb.state
+    return ReferenceObjectDB(fs, path_info, state=state)
 
 
 def _load_from_state(odb, staging, path_info, fs, name):
-    from . import load
+    from . import check, load
     from .errors import ObjectFormatError
     from .tree import Tree
 
@@ -212,6 +232,7 @@ def _load_from_state(odb, staging, path_info, fs, name):
             if odb_.exists(hash_info):
                 try:
                     obj = load(odb_, hash_info)
+                    check(odb_, obj, check_hash=False)
                     if isinstance(obj, Tree):
                         obj.hash_info.nfiles = len(obj)
                         for key, entry in obj:
@@ -222,8 +243,8 @@ def _load_from_state(odb, staging, path_info, fs, name):
                         obj.path_info = path_info
                     assert obj.hash_info.name == name
                     obj.hash_info.size = hash_info.size
-                    return obj
-                except ObjectFormatError:
+                    return odb_, obj
+                except (ObjectFormatError, FileNotFoundError):
                     pass
     raise FileNotFoundError
 
@@ -248,17 +269,31 @@ def _stage_external_tree_info(odb, tree, name):
 
 
 def stage(
-    odb: Optional["ObjectDB"],
+    odb: "ObjectDB",
     path_info: "DvcPath",
     fs: "BaseFileSystem",
     name: str,
     upload: bool = False,
+    dry_run: bool = False,
     **kwargs,
-) -> "HashFile":
+) -> Tuple["ObjectDB", "HashFile"]:
+    """Stage (prepare) objects from the given path for addition to an ODB.
+
+    Returns at tuple of (staging_odb, object) where addition to the ODB can
+    be completed by transferring the object from staging to the dest ODB.
+
+    If dry_run is True, object hashes will be computed and returned, but file
+    objects themselves will not be added to the staging ODB (i.e. the resulting
+    file objects cannot transferred from staging to another ODB).
+
+    If upload is True, files will be uploaded to a temporary path on the dest
+    ODB filesystem, and staged objects will reference the uploaded path rather
+    than the original source path.
+    """
     assert path_info and path_info.scheme == fs.scheme
 
     details = fs.info(path_info)
-    staging = get_staging(odb)
+    staging = _get_staging(odb)
     if odb:
         try:
             return _load_from_state(odb, staging, path_info, fs, name)
@@ -267,26 +302,29 @@ def stage(
 
     if details["type"] == "directory":
         obj = _get_tree_obj(
-            path_info, fs, name, odb=odb, upload=upload, **kwargs
+            path_info,
+            fs,
+            details,
+            name,
+            odb=staging,
+            upload_odb=odb if upload else None,
+            dry_run=dry_run,
+            **kwargs,
         )
-        if name == "md5":
-            if odb and odb.exists(obj.hash_info):
-                raw = odb.get(obj.hash_info)
-            else:
-                staging.add(obj.path_info, obj.fs, obj.hash_info)
-                raw = staging.get(obj.hash_info)
-            # cleanup unneeded memfs tmpfile and return obj based on staging
-            # ODB fs/path
-            if obj.fs != raw.fs:
-                obj.fs.remove(obj.path_info)
-            obj.fs = raw.fs
-            obj.path_info = raw.path_info
-        else:
+        logger.debug("staged tree '%s'", obj)
+        if name != "md5":
             obj = _stage_external_tree_info(odb, obj, name)
     else:
-        _, obj = _get_file_obj(path_info, fs, name, odb=odb, upload=upload)
+        _, obj = _get_file_obj(
+            path_info,
+            fs,
+            name,
+            odb=staging,
+            upload_odb=odb if upload else None,
+            dry_run=dry_run,
+        )
 
     if odb and odb.state and obj.hash_info:
         odb.state.save(path_info, fs, obj.hash_info)
 
-    return obj
+    return staging, obj

dvc/objects/status.py

@@ -2,6 +2,7 @@ import logging
 from typing import TYPE_CHECKING, Dict, Iterable, NamedTuple, Optional, Set
 
 from dvc.hash_info import HashInfo
+from dvc.scheme import Schemes
 
 from .tree import Tree
 
@@ -14,23 +15,21 @@ logger = logging.getLogger(__name__)
 
 
 class StatusResult(NamedTuple):
-    exists: Set["HashFile"]
-    missing: Set["HashFile"]
+    exists: Set["HashInfo"]
+    missing: Set["HashInfo"]
 
 
 class CompareStatusResult(NamedTuple):
-    ok: Set["HashFile"]
-    missing: Set["HashFile"]
-    new: Set["HashFile"]
-    deleted: Set["HashFile"]
+    ok: Set["HashInfo"]
+    missing: Set["HashInfo"]
+    new: Set["HashInfo"]
+    deleted: Set["HashInfo"]
 
 
-def _indexed_dir_hashes(odb, index, hash_infos, name, cache_odb):
-    from . import load
-
+def _indexed_dir_hashes(odb, index, dir_objs, name, cache_odb):
     # Validate our index by verifying all indexed .dir hashes
     # still exist on the remote
-    dir_hashes = {hash_info.value for hash_info in hash_infos}
+    dir_hashes = set(dir_objs.keys())
     indexed_dirs = set(index.dir_hashes())
     indexed_dir_exists = set()
     if indexed_dirs:
@@ -51,10 +50,12 @@ def _indexed_dir_hashes(odb, index, hash_infos, name, cache_odb):
     # If .dir hash exists in the ODB, assume directory contents
     # also exists
     for dir_hash in dir_exists:
-        try:
-            tree = load(cache_odb, HashInfo(name, dir_hash))
-        except FileNotFoundError:
-            continue
+        tree = dir_objs.get(dir_hash)
+        if not tree:
+            try:
+                tree = Tree.load(cache_odb, HashInfo(name, dir_hash))
+            except FileNotFoundError:
+                continue
         file_hashes = [entry.hash_info.value for _, entry in tree]
         if dir_hash not in index:
             logger.debug(
@@ -67,21 +68,13 @@ def _indexed_dir_hashes(odb, index, hash_infos, name, cache_odb):
         yield tree.hash_info.value
 
 
-def _status_staging(objs: Iterable["HashFile"]) -> "StatusResult":
-    exists: Set["HashFile"] = set()
-    for obj in objs:
-        if isinstance(obj, Tree):
-            exists.update(entry for _, entry in obj)
-        exists.add(obj)
-    return StatusResult(exists, set())
-
-
 def status(
     odb: "ObjectDB",
-    objs: Iterable["HashFile"],
+    obj_ids: Iterable["HashInfo"],
     name: Optional[str] = None,
     index: Optional["ObjectDBIndexBase"] = None,
     cache_odb: Optional["ObjectDB"] = None,
+    shallow: bool = True,
     **kwargs,
 ) -> "StatusResult":
     """Return status of whether or not the specified objects exist odb.
@@ -93,39 +86,41 @@ def status(
         exists: objs that exist in odb
         missing: objs that do not exist in ODB
     """
-    from .stage import is_memfs_staging
-
     logger.debug("Preparing to collect status from '%s'", odb.path_info)
     if not name:
         name = odb.fs.PARAM_CHECKSUM
 
-    if is_memfs_staging(odb):
-        # assume memfs staged objects already exist
-        return _status_staging(objs)
+    if cache_odb is None:
+        cache_odb = odb
 
-    hash_objs: Dict[str, "HashFile"] = {}
-    hashes: Set[str] = set()
-    dir_infos: Set["HashInfo"] = set()
-    exists: Set[str] = set()
-    for obj in objs:
-        assert obj.hash_info and obj.hash_info.value
-        if isinstance(obj, Tree):
-            for _, entry in obj:
-                assert entry.hash_info and entry.hash_info.value
-                hash_objs[entry.hash_info.value] = entry
-                hashes.add(entry.hash_info.value)
+    hash_infos: Dict[str, "HashInfo"] = {}
+    dir_objs: Dict[str, Optional["HashFile"]] = {}
+    for hash_info in obj_ids:
+        assert hash_info.value
+        if hash_info.isdir:
+            if shallow:
+                tree = None
+            else:
+                tree = Tree.load(cache_odb, hash_info)
+                for _, entry in tree:
+                    assert entry.hash_info and entry.hash_info.value
+                    hash_infos[entry.hash_info.value] = entry.hash_info
             if index:
-                dir_infos.add(obj.hash_info)
-        hashes.add(obj.hash_info.value)
-        hash_objs[obj.hash_info.value] = obj
+                dir_objs[hash_info.value] = tree
+        hash_infos[hash_info.value] = hash_info
+
+    if odb.fs.scheme == Schemes.MEMORY:
+        # assume memfs staged objects already exist
+        return StatusResult(set(hash_infos.values()), set())
+
+    hashes: Set[str] = set(hash_infos.keys())
+    exists: Set[str] = set()
 
     logger.debug("Collecting status from '%s'", odb.path_info)
     if index and hashes:
-        if dir_infos:
-            if cache_odb is None:
-                cache_odb = odb
+        if dir_objs:
             exists = hashes.intersection(
-                _indexed_dir_hashes(odb, index, dir_infos, name, cache_odb)
+                _indexed_dir_hashes(odb, index, dir_objs, name, cache_odb)
             )
             hashes.difference_update(exists)
         if hashes:
@@ -137,15 +132,15 @@ def status(
             odb.hashes_exist(hashes, name=str(odb.path_info), **kwargs)
         )
     return StatusResult(
-        {hash_objs[hash_] for hash_ in exists},
-        {hash_objs[hash_] for hash_ in (hashes - exists)},
+        {hash_infos[hash_] for hash_ in exists},
+        {hash_infos[hash_] for hash_ in (hashes - exists)},
     )
 
 
 def compare_status(
     src: "ObjectDB",
     dest: "ObjectDB",
-    objs: Iterable["HashFile"],
+    obj_ids: Iterable["HashInfo"],
     log_missing: bool = True,
     check_deleted: bool = True,
     src_index: Optional["ObjectDBIndexBase"] = None,
@@ -160,11 +155,17 @@ def compare_status(
         new: hashes that only exist in src
         deleted: hashes that only exist in dest
     """
-    dest_exists, dest_missing = status(dest, objs, index=dest_index, **kwargs)
+    if "cache_odb" not in kwargs:
+        kwargs["cache_odb"] = src
+    dest_exists, dest_missing = status(
+        dest, obj_ids, index=dest_index, **kwargs
+    )
     # for transfer operations we can skip src status check when all objects
     # already exist in dest
     if dest_missing or check_deleted:
-        src_exists, src_missing = status(src, objs, index=src_index, **kwargs)
+        src_exists, src_missing = status(
+            src, obj_ids, index=src_index, **kwargs
+        )
     else:
         src_exists = dest_exists
         src_missing = set()
@@ -176,7 +177,8 @@ def compare_status(
     )
     if log_missing and result.missing:
         missing_desc = "\n".join(
-            f"name: {obj.name}, {obj.hash_info}" for obj in result.missing
+            f"name: {hash_info.obj_name}, {hash_info}"
+            for hash_info in result.missing
         )
         logger.warning(
             "Some of the cache files do not exist neither locally "

dvc/objects/transfer.py

@@ -4,14 +4,15 @@ import logging
 from concurrent import futures
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial, wraps
-from typing import TYPE_CHECKING, Iterable, Optional, Set
+from typing import TYPE_CHECKING, Callable, Iterable, Optional
+
+from funcy import split
 
 from dvc.progress import Tqdm
 
-from .file import HashFile
-from .tree import Tree
-
 if TYPE_CHECKING:
+    from dvc.hash_info import HashInfo
+
     from .db.base import ObjectDB
     from .db.index import ObjectDBIndexBase
 
@@ -20,9 +21,9 @@ logger = logging.getLogger(__name__)
 
 def _log_exceptions(func):
     @wraps(func)
-    def wrapper(odb, obj, *args, **kwargs):
+    def wrapper(path_info, *args, **kwargs):
         try:
-            func(odb, obj, *args, **kwargs)
+            func(path_info, *args, **kwargs)
             return 0
         except Exception as exc:  # pylint: disable=broad-except
             # NOTE: this means we ran out of file descriptors and there is no
@@ -31,23 +32,17 @@ def _log_exceptions(func):
             if isinstance(exc, OSError) and exc.errno == errno.EMFILE:
                 raise
 
-            logger.exception(
-                "failed to transfer '%s' to '%s'",
-                obj.path_info,
-                odb.get(obj.hash_info).path_info,
-            )
+            logger.exception("failed to transfer '%s'", path_info)
             return 1
 
     return wrapper
 
 
-def _transfer(src, dest, dir_objs, file_objs, missing, jobs, verify, **kwargs):
-    from . import save
-    from .stage import is_memfs_staging
-
-    is_staged = is_memfs_staging(src)
-    func = _log_exceptions(save)
-    total = len(dir_objs) + len(file_objs)
+def _transfer(
+    src, dest, dir_ids, file_ids, missing_ids, jobs, verify, move, **kwargs
+):
+    func = _log_exceptions(dest.add)
+    total = len(dir_ids) + len(file_ids)
     if total == 0:
         return 0
     with Tqdm(total=total, unit="file", desc="Transferring") as pbar:
@@ -59,37 +54,43 @@ def _transfer(src, dest, dir_objs, file_objs, missing, jobs, verify, **kwargs):
                 jobs,
                 func,
                 src,
-                dest,
-                is_staged,
                 verify,
+                move,
             )
-            processor.save_func = func
+            processor.add_func = func
             _do_transfer(
                 src,
                 dest,
-                dir_objs,
-                file_objs,
-                {obj.hash_info for obj in missing},
+                dir_ids,
+                file_ids,
+                missing_ids,
                 processor,
+                verify=verify,
+                move=move,
                 **kwargs,
             )
     return total
 
 
-def _create_tasks(executor, jobs, func, src, dest, is_staged, verify, objs):
+def _create_tasks(executor, jobs, func, src, verify, move, obj_ids):
     fails = 0
-    obj_iter = iter(objs)
+    hash_iter = iter(obj_ids)
+
+    def submit(hash_info):
+        obj = src.get(hash_info)
+        return executor.submit(
+            func,
+            obj.path_info,
+            obj.fs,
+            obj.hash_info,
+            verify=verify,
+            move=move,
+        )
 
     def create_taskset(amount):
         return {
-            executor.submit(
-                func,
-                dest,
-                _raw_obj(src, obj, is_staged),
-                move=False,
-                verify=verify,
-            )
-            for obj in itertools.islice(obj_iter, amount)
+            submit(hash_info)
+            for hash_info in itertools.islice(hash_iter, amount)
         }
 
     tasks = create_taskset(jobs * 5)
@@ -100,39 +101,46 @@ def _create_tasks(executor, jobs, func, src, dest, is_staged, verify, objs):
     return fails
 
 
-def _raw_obj(odb, obj, is_staged=False):
-    if is_staged:
-        return HashFile(obj.path_info, obj.fs, obj.hash_info)
-    return odb.get(obj.hash_info)
-
-
 def _do_transfer(
-    src,
-    dest,
-    dir_objs,
-    file_objs,
-    missing_hashes,
-    processor,
+    src: "ObjectDB",
+    dest: "ObjectDB",
+    dir_ids: Iterable["HashInfo"],
+    file_ids: Iterable["HashInfo"],
+    missing_ids: Iterable["HashInfo"],
+    processor: Callable,
     src_index: Optional["ObjectDBIndexBase"] = None,
     dest_index: Optional["ObjectDBIndexBase"] = None,
+    cache_odb: Optional["ObjectDB"] = None,
     **kwargs,
 ):
     from dvc.exceptions import FileTransferError
+    from dvc.objects.errors import ObjectFormatError
 
     total_fails = 0
     succeeded_dir_objs = []
-    all_file_objs = set(file_objs)
+    all_file_ids = set(file_ids)
 
-    for dir_obj in dir_objs:
-        bound_file_objs = set()
-        directory_hashes = {entry.hash_info for _, entry in dir_obj}
+    for dir_hash in dir_ids:
+        from .tree import Tree
 
-        for file_obj in all_file_objs.copy():
-            if file_obj.hash_info in directory_hashes:
-                bound_file_objs.add(file_obj)
-                all_file_objs.remove(file_obj)
+        bound_file_ids = set()
+        dir_obj: Optional["Tree"] = None
+        for odb in (cache_odb, src):
+            if odb is not None:
+                try:
+                    dir_obj = Tree.load(odb, dir_hash)
+                    break
+                except (FileNotFoundError, ObjectFormatError):
+                    pass
+        assert dir_obj
+        entry_ids = {entry.hash_info for _, entry in dir_obj}
 
-        dir_fails = processor(bound_file_objs)
+        for file_hash in all_file_ids.copy():
+            if file_hash in entry_ids:
+                bound_file_ids.add(file_hash)
+                all_file_ids.remove(file_hash)
+
+        dir_fails = processor(bound_file_ids)
         if dir_fails:
             logger.debug(
                 "failed to upload full contents of '%s', "
@@ -145,7 +153,7 @@ def _do_transfer(
                 dest.get(dir_obj.hash_info).path_info,
             )
             total_fails += dir_fails + 1
-        elif directory_hashes.intersection(missing_hashes):
+        elif entry_ids.intersection(missing_ids):
             # if for some reason a file contained in this dir is
             # missing both locally and in the remote, we want to
             # push whatever file content we have, but should not
@@ -156,15 +164,19 @@ def _do_transfer(
                 dir_obj.name,
             )
         else:
-            is_dir_failed = processor.save_func(
-                dest, src.get(dir_obj.hash_info), move=False
+            raw_obj = src.get(dir_obj.hash_info)
+            is_dir_failed = processor.add_func(  # type: ignore[attr-defined]
+                raw_obj.path_info,
+                raw_obj.fs,
+                raw_obj.hash_info,
+                **kwargs,
             )
             total_fails += is_dir_failed
             if not is_dir_failed:
                 succeeded_dir_objs.append(dir_obj)
 
     # insert the rest
-    total_fails += processor(all_file_objs)
+    total_fails += processor(all_file_ids)
     if total_fails:
         if src_index:
             src_index.clear()
@@ -179,15 +191,17 @@ def _do_transfer(
                 dir_obj.hash_info,
                 len(file_hashes),
             )
+            assert dir_obj.hash_info and dir_obj.hash_info.value
             dest_index.update([dir_obj.hash_info.value], file_hashes)
 
 
 def transfer(
     src: "ObjectDB",
     dest: "ObjectDB",
-    objs: Iterable["HashFile"],
+    obj_ids: Iterable["HashInfo"],
     jobs: Optional[int] = None,
     verify: bool = False,
+    move: bool = False,
     **kwargs,
 ) -> int:
     """Transfer (copy) the specified objects from one ODB to another.
@@ -201,27 +215,25 @@ def transfer(
         src.path_info,
         dest.path_info,
     )
-    status = compare_status(src, dest, objs, check_deleted=False, **kwargs)
+    if src == dest:
+        return 0
+
+    status = compare_status(src, dest, obj_ids, check_deleted=False, **kwargs)
     if not status.new:
         return 0
 
-    files: Set["HashFile"] = set()
-    dirs: Set["HashFile"] = set()
-    for obj in status.new:
-        if isinstance(obj, Tree):
-            dirs.add(obj)
-        else:
-            files.add(obj)
+    dir_ids, file_ids = split(lambda hash_info: hash_info.isdir, status.new)
     if jobs is None:
         jobs = dest.fs.jobs
 
     return _transfer(
         src,
         dest,
-        dirs,
-        files,
+        set(dir_ids),
+        set(file_ids),
         status.missing,
         jobs,
         verify,
+        move,
         **kwargs,
     )

dvc/objects/tree.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import posixpath
-from typing import Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
 from funcy import cached_property
 
@@ -9,6 +9,9 @@ from .errors import ObjectFormatError
 from .file import HashFile
 from .stage import get_file_hash
 
+if TYPE_CHECKING:
+    from dvc.hash_info import HashInfo
+
 logger = logging.getLogger(__name__)
 
 
@@ -29,7 +32,7 @@ class Tree(HashFile):
         self.__dict__.pop("trie", None)
         self._dict[key] = obj
 
-    def digest(self):
+    def digest(self, hash_info: Optional["HashInfo"] = None):
         from dvc.fs.memory import MemoryFileSystem
         from dvc.path_info import CloudURLInfo
         from dvc.utils import tmp_fname
@@ -40,8 +43,12 @@ class Tree(HashFile):
             fobj.write(self.as_bytes())
         self.fs = memfs
         self.path_info = path_info
-        self.hash_info = get_file_hash(path_info, memfs, "md5")
-        self.hash_info.value += ".dir"
+        if hash_info:
+            self.hash_info = hash_info
+        else:
+            self.hash_info = get_file_hash(path_info, memfs, "md5")
+            assert self.hash_info.value
+            self.hash_info.value += ".dir"
         try:
             self.hash_info.size = sum(obj.size for _, obj in self)
         except TypeError: