diff --git a/dvc/repo/update.py b/dvc/repo/update.py index 74a7522df..798da8b14 100644 --- a/dvc/repo/update.py +++ b/dvc/repo/update.py @@ -1,7 +1,6 @@ from typing import TYPE_CHECKING, List from dvc.exceptions import InvalidArgumentError -from dvc.stage.exceptions import StageUpdateError from . import locked @@ -33,11 +32,7 @@ def update( # noqa: C901 "--to-remote can't be used with --no-download" ) - if ( - not to_remote - and remote - and not self.cloud.get_remote(name=remote).worktree - ): + if not to_remote and remote: raise InvalidArgumentError( "--remote can't be used without --to-remote" ) @@ -62,9 +57,6 @@ def update( # noqa: C901 stage.dump() if other_stage_infos: - remote_obj = self.cloud.get_remote(name=remote) - if not remote_obj.worktree: - raise StageUpdateError(other_stage_infos[0].stage.relpath) if rev: raise InvalidArgumentError( "--rev can't be used with worktree update" @@ -80,7 +72,6 @@ def update( # noqa: C901 update_worktree_stages( self, other_stage_infos, - remote_obj, ) stages = import_stages | { diff --git a/dvc/repo/worktree.py b/dvc/repo/worktree.py index da5ba5f3f..e935ab93c 100644 --- a/dvc/repo/worktree.py +++ b/dvc/repo/worktree.py @@ -15,6 +15,7 @@ from typing import ( from funcy import first from dvc.fs.callbacks import Callback +from dvc.stage.exceptions import StageUpdateError if TYPE_CHECKING: from dvc.data_cloud import Remote @@ -207,7 +208,7 @@ def push_worktree( for out in view.outs: workspace, _key = out.index_key - _update_out_meta(out, repo.index.data[workspace], remote_obj.name) + _merge_push_meta(out, repo.index.data[workspace], remote_obj.name) stages.add(out.stage) for stage in stages: @@ -215,11 +216,17 @@ def push_worktree( return pushed -def _update_out_meta( +def _merge_push_meta( out: "Output", index: Union["DataIndex", "DataIndexView"], remote: Optional[str] = None, ): + """Merge existing output meta with newly pushed meta. + + Existing version IDs for unchanged files will be preserved to reduce merge + conflicts (i.e. the DVC output's version ID may not match the pushed/latest + version ID as long when the file content of both versions is the same). + """ from dvc_data.index.save import build_tree _, key = out.index_key @@ -260,10 +267,8 @@ def _update_out_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 @@ -274,54 +279,80 @@ def update_worktree_stages( 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) + remote_indexes: Dict[str, Tuple["Remote", "DataIndex"]] = {} for stage in view.stages: for out in stage.outs: - _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 - and 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, - ) - continue - _fetch_out_changes(out, local_index, remote_index, remote) - stage.save() - for out in stage.outs: - _update_out_meta(out, remote_index) + _update_worktree_out(repo, out, local_index, remote_indexes) stage.dump(with_files=True, update_pipeline=False) +def _update_worktree_out( + repo: "Repo", + out: "Output", + local_index: Union["DataIndex", "DataIndexView"], + remote_indexes: Dict[str, Tuple["Remote", "DataIndex"]], +): + from dvc_data.index import build + + remote_name = out.remote or out.meta.remote + if not remote_name: + logger.warning( + "Could not update '%s', it was never pushed to a remote", + out, + ) + return + + if remote_name in remote_indexes: + remote, remote_index = remote_indexes[remote_name] + else: + remote = repo.cloud.get_remote(remote_name, "update") + if not remote.worktree: + raise StageUpdateError(out.stage.relpath) + logger.debug("indexing latest worktree for '%s'", remote.path) + remote_index = build(remote.path, remote.fs) + remote_indexes[remote_name] = remote, remote_index + _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, + ) + return + + entry = remote_index[key] + if ( + entry.meta + and entry.meta.isdir + and 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, + ) + return + + _fetch_out_changes(out, local_index, remote_index, remote) + _update_out_meta( + repo, + out, + local_index, + remote_index, + remote, + ) + + def _fetch_out_changes( out: "Output", local_index: Union["DataIndex", "DataIndexView"], remote_index: Union["DataIndex", "DataIndexView"], remote: "Remote", ): - from dvc_data.index import DataIndex, checkout + from dvc_data.index import 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) + old, new = _get_diff_indexes(out, local_index, remote_index) total = len(new) with Callback.as_tqdm_callback( unit="file", desc=f"Updating '{out}'", disable=total == 0 @@ -338,3 +369,87 @@ def _fetch_out_changes( meta_cmp_key=partial(_meta_checksum, remote.fs), callback=cb, ) + out.save() + + +def _get_diff_indexes( + out: "Output", + local_index: Union["DataIndex", "DataIndexView"], + remote_index: Union["DataIndex", "DataIndexView"], +) -> Tuple["DataIndex", "DataIndex"]: + from dvc_data.index import DataIndex + + _, 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) + return old, new + + +def _update_out_meta( + repo: "Repo", + out: "Output", + local_index: Union["DataIndex", "DataIndexView"], + remote_index: Union["DataIndex", "DataIndexView"], + remote: "Remote", +): + from dvc_data.index.save import build_tree + + index = _get_update_diff_index( + repo, out, local_index, remote_index, remote + ) + + _, key = out.index_key + entry = index[key] + if out.isdir(): + 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 + out.meta = entry.meta + if out.meta: + out.meta.remote = remote.name + + +def _get_update_diff_index( + repo: "Repo", + out: "Output", + local_index: Union["DataIndex", "DataIndexView"], + remote_index: Union["DataIndex", "DataIndexView"], + remote: "Remote", +) -> "DataIndex": + from dvc_data.index import DataIndex + from dvc_data.index.diff import ADD, MODIFY, UNCHANGED, diff + + old, new = _get_diff_indexes(out, local_index, remote_index) + index = DataIndex() + for change in diff( + old, + new, + meta_only=True, + meta_cmp_key=partial(_meta_checksum, remote.fs), + with_unchanged=True, + ): + if change.typ == ADD or change.typ == MODIFY: + entry = change.new + # preserve md5's which were calculated in out.save() after + # downloading + if out.isdir(): + if not entry.meta.isdir: + fs_path = repo.fs.path.join(repo.root_dir, *entry.key) + tree = cast("Tree", out.obj) + _, entry.hash_info = tree.get( # type: ignore[misc] + repo.fs.path.relparts(fs_path, out.fs_path) + ) + else: + entry.hash_info = out.hash_info + index[change.new.key] = change.new + elif change.typ == UNCHANGED: + index[change.old.key] = change.old + return index diff --git a/dvc/testing/remote_tests.py b/dvc/testing/remote_tests.py index 30b189446..1ac6dde60 100644 --- a/dvc/testing/remote_tests.py +++ b/dvc/testing/remote_tests.py @@ -5,6 +5,7 @@ import pytest from dvc.stage.cache import RunCacheNotSupported from dvc.utils.fs import remove +from dvc_data.hashfile.tree import Tree # pylint: disable=unused-argument @@ -247,8 +248,8 @@ class TestRemoteWorktree: def test_update( self, tmp_dir, dvc, remote_worktree # pylint: disable=W0613 ): - tmp_dir.dvc_gen("foo", "foo") - tmp_dir.dvc_gen( + (foo_stage,) = tmp_dir.dvc_gen("foo", "foo") + (data_dir_stage,) = tmp_dir.dvc_gen( { "data_dir": { "data_sub_dir": {"data_sub": "data_sub"}, @@ -258,11 +259,41 @@ class TestRemoteWorktree: } ) dvc.push() + orig_foo = foo_stage.reload().outs[0] + orig_data_dir = data_dir_stage.reload().outs[0] (remote_worktree / "foo").write_text("bar") (remote_worktree / "data_dir" / "data").write_text("modified") (remote_worktree / "data_dir" / "new_data").write_text("new data") dvc.update([str(tmp_dir / "foo.dvc"), str(tmp_dir / "data_dir.dvc")]) + updated_foo = foo_stage.reload().outs[0] + updated_data_dir = data_dir_stage.reload().outs[0] + + assert updated_foo.meta.version_id + assert updated_foo.meta.version_id != orig_foo.meta.version_id + updated_data_dir = data_dir_stage.reload().outs[0] + orig_tree = orig_data_dir.get_obj() + updated_tree = Tree.from_list(updated_data_dir.files, hash_name="md5") + assert orig_tree.get(("data_sub_dir", "data_sub")) == updated_tree.get( + ("data_sub_dir", "data_sub") + ) + orig_meta, _ = orig_tree.get(("data",)) + updated_meta, _ = updated_tree.get(("data",)) + assert orig_meta.version_id + assert updated_meta.version_id + assert orig_meta.version_id != updated_meta.version_id + meta, hash_info = updated_tree.get(("new_data",)) + assert meta + assert hash_info + + assert (tmp_dir / "foo").read_text() == "bar" + assert (tmp_dir / "data_dir" / "data").read_text() == "modified" + assert (tmp_dir / "data_dir" / "new_data").read_text() == "new data" + + remove(dvc.odb.local.path) + remove(tmp_dir / "foo") + remove(tmp_dir / "data_dir") + dvc.pull() assert (tmp_dir / "foo").read_text() == "bar" assert (tmp_dir / "data_dir" / "data").read_text() == "modified" assert (tmp_dir / "data_dir" / "new_data").read_text() == "new data"