diff --git a/dvc/repo/worktree.py b/dvc/repo/worktree.py index fde64001a..c8e99521f 100644 --- a/dvc/repo/worktree.py +++ b/dvc/repo/worktree.py @@ -1,5 +1,6 @@ import logging -from typing import TYPE_CHECKING, Any, Dict, Optional +from functools import partial +from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Union, cast from dvc.fs.callbacks import Callback @@ -8,13 +9,25 @@ if TYPE_CHECKING: from dvc.index import Index, IndexView from dvc.output import Output from dvc.repo import Repo + from dvc.repo.stage import StageInfo from dvc.stage import Stage from dvc.types import TargetType from dvc_data.hashfile.meta import Meta + from dvc_data.hashfile.tree import Tree + from dvc_data.index import DataIndex, DataIndexView + from dvc_objects.fs.base import FileSystem logger = logging.getLogger(__name__) +# for files, if our version's checksum (etag) matches the latest remote +# checksum, we do not need to push, even if the version IDs don't match +def _meta_checksum(fs: "FileSystem", meta: "Meta") -> Any: + if not meta or meta.isdir: + return meta + return getattr(meta, fs.PARAM_CHECKSUM) + + def worktree_view( index: "Index", targets: Optional["TargetType"] = None, @@ -82,7 +95,7 @@ def push_worktree(repo: "Repo", remote: "Remote") -> int: view = worktree_view(repo.index, push=True, latest_only=remote.worktree) new_index = view.data["repo"] if remote.worktree: - logger.debug("Indexing latest worktree for '%s'", remote.path) + logger.debug("indexing latest worktree for '%s'", remote.path) old_index = build(remote.path, remote.fs) logger.debug("Pushing worktree changes to '%s'", remote.path) else: @@ -90,17 +103,9 @@ def push_worktree(repo: "Repo", remote: "Remote") -> int: logger.debug("Pushing version-aware files to '%s'", remote.path) if remote.worktree: - - # for files, if our version's checksum (etag) matches the latest remote - # checksum, we do not need to push, even if the version IDs don't match - def _checksum(meta: "Meta") -> Any: - if not meta or meta.isdir: - return meta - return getattr(meta, remote.fs.PARAM_CHECKSUM) - diff_kwargs: Dict[str, Any] = { "meta_only": True, - "meta_cmp_key": _checksum, + "meta_cmp_key": partial(_meta_checksum, remote.fs), } else: diff_kwargs = {} @@ -121,41 +126,125 @@ def push_worktree(repo: "Repo", remote: "Remote") -> int: **diff_kwargs, ) if pushed: - _update_pushed_meta(repo, view) + for stage in view.stages: + for out in stage.outs: + workspace, _key = out.index_key + _update_out_meta(out, repo.index.data[workspace]) + stage.dvcfile.dump(stage, with_files=True, update_pipeline=False) return pushed -def _update_pushed_meta(repo: "Repo", view: "IndexView"): +def _update_out_meta( + out: "Output", index: Union["DataIndex", "DataIndexView"] +): from dvc_data.index.save import build_tree + _, key = out.index_key + entry = index[key] + repo = out.stage.repo + if out.isdir(): + old_tree = cast("Tree", out.get_obj()) + entry.hash_info = old_tree.hash_info + entry.meta = out.meta + for subkey, entry in index.iteritems(key): + if entry.meta.isdir: + continue + fs_path = repo.fs.path.join(repo.root_dir, *subkey) + meta, hash_info = old_tree.get( + repo.fs.path.relparts(fs_path, out.fs_path) + ) + entry.hash_info = hash_info + if meta.version_id is not None: + # preserve existing version IDs for unchanged files in + # this dir (entry will have the latest remote version + # ID after checkout) + entry.meta = meta + tree_meta, new_tree = build_tree(index, key) + out.obj = new_tree + out.hash_info = new_tree.hash_info + out.meta = tree_meta + else: + if entry.hash_info: + out.hash_info = entry.hash_info + if out.meta.version_id is None: + out.meta = entry.meta + + +def update_worktree_stages( + repo: "Repo", + stage_infos: Iterable["StageInfo"], + remote: "Remote", +): + from dvc.repo.index import IndexView + from dvc_data.index import build + + def outs_filter(out: "Output") -> bool: + return out.is_in_repo and out.use_cache and out.can_push + + view = IndexView( + repo.index, + stage_infos, + outs_filter=outs_filter, + ) + local_index = view.data["repo"] + logger.debug("indexing latest worktree for '%s'", remote.path) + remote_index = build(remote.path, remote.fs) for stage in view.stages: for out in stage.outs: - workspace, key = out.index_key - index = repo.index.data[workspace] - entry = index[key] - if out.isdir(): - old_tree = out.get_obj() - entry.hash_info = old_tree.hash_info - entry.meta = out.meta - for subkey, entry in index.iteritems(key): - if entry.meta.isdir: - continue - fs_path = repo.fs.path.join(repo.root_dir, *subkey) - meta, hash_info = old_tree.get( - repo.fs.path.relparts(fs_path, out.fs_path) + _workspace, key = out.index_key + if key not in remote_index: + logger.warning( + "Could not update '%s', it does not exist in the remote", + out, + ) + continue + entry = remote_index[key] + if entry.meta and entry.meta.isdir: + if not any( + subkey != key and subentry.meta and not subentry.meta.isdir + for subkey, subentry in remote_index.iteritems(key) + ): + logger.warning( + "Could not update '%s', directory is empty in the " + "remote", + out, ) - entry.hash_info = hash_info - if meta.version_id is not None: - # preserve existing version IDs for unchanged files in - # this dir (entry will have the latest remote version - # ID after checkout) - entry.meta = meta - tree_meta, new_tree = build_tree(index, key) - out.obj = new_tree - out.hash_info = new_tree.hash_info - out.meta = tree_meta - else: - out.hash_info = entry.hash_info - if out.meta.version_id is None: - out.meta = entry.meta + continue + _fetch_out_changes(out, local_index, remote_index, remote) + stage.save(merge_versioned=True) + for out in stage.outs: + _update_out_meta(out, remote_index) stage.dvcfile.dump(stage, with_files=True, update_pipeline=False) + + +def _fetch_out_changes( + out: "Output", + local_index: "DataIndex", + remote_index: "DataIndex", + remote: "Remote", +): + from dvc_data.index import DataIndex, checkout + + _, key = out.index_key + old = DataIndex() + new = DataIndex() + for _, entry in local_index.iteritems(key): + old.add(entry) + for _, entry in remote_index.iteritems(key): + new.add(entry) + total = len(new) + with Callback.as_tqdm_callback( + unit="file", desc=f"Updating '{out}'", disable=total == 0 + ) as cb: + cb.set_size(total) + checkout( + new, + out.repo.root_dir, + out.fs, + old=old, + delete=True, + update_meta=False, + meta_only=True, + meta_cmp_key=partial(_meta_checksum, remote.fs), + callback=cb, + )