index: checkout: keep cache objects protected after removing links

src/dvc_data/index/checkout.py

@@ -26,6 +26,7 @@
     from dvc_objects.fs.base import AnyFSPath, FileSystem
     from fsspec import Callback
 
+    from dvc_data.hashfile.db import HashFileDB
     from dvc_data.hashfile.hash_info import HashInfo
     from dvc_data.hashfile.state import StateBase
 
@@ -67,6 +68,34 @@ def _delete_files(
     fs.remove([fs.join(path, *(entry.key or ())) for entry in entries])
 
 
+def _cache_objects(
+    entries: list["DataIndexEntry"],
+    index: Optional["BaseDataIndex"],
+) -> list[tuple["HashFileDB", str]]:
+    """Find the cache objects that the given workspace entries may link to."""
+    if index is None:
+        return []
+
+    objects = []
+    for entry in entries:
+        if not entry.hash_info or not entry.hash_info.value:
+            continue
+        try:
+            storage = index.storage_map[entry.key].cache
+        except KeyError:
+            continue
+        if not isinstance(storage, ObjectStorage):
+            continue
+        objects.append((storage.odb, storage.odb.oid_to_path(entry.hash_info.value)))
+    return objects
+
+
+def _protect_objects(objects: list[tuple["HashFileDB", str]]) -> None:
+    for odb, obj_path in objects:
+        if odb.fs.exists(obj_path):
+            odb.protect(obj_path)
+
+
 def _create_files(  # noqa: C901, PLR0912, PLR0913
     entries,
     index: Optional["BaseDataIndex"],
@@ -353,11 +382,21 @@ def apply(
     for entry in diff.dirs_failed:
         onerror(None, fs.join(path, *entry.key), None)
 
+    # NOTE: On some filesystems (e.g. NTFS) a hardlink shares its read-only
+    # attribute with the cache object. Deleting a protected hardlink has to
+    # clear that attribute first, which leaves the cache object writable too,
+    # so we protect those objects again after the workspace files are gone.
+    deleted_objects = (
+        _cache_objects(diff.files_delete, diff.new)
+        if isinstance(fs, LocalFileSystem)
+        else []
+    )
     _delete_files(
         diff.files_delete,
         path,
         fs,
     )
+    _protect_objects(deleted_objects)
     _delete_dirs(diff.dirs_delete, path, fs)
     _create_dirs(diff.dirs_create, path, fs)
 

tests/index/test_checkout.py

@@ -1,5 +1,9 @@
+import os
+
+import pytest
 from dvc_objects.fs.local import LocalFileSystem
 
+from dvc_data.hashfile.db.local import LocalHashFileDB
 from dvc_data.hashfile.hash_info import HashInfo
 from dvc_data.hashfile.meta import Meta
 from dvc_data.index import DataIndex, DataIndexEntry, ObjectStorage
@@ -136,3 +140,59 @@ def test_checkout_delete_nested_dir(tmp_path, odb):
 
     assert not (tmp_path / "dir1" / "subdir1").exists()
     assert not (tmp_path / "dir1").exists()
+
+
+FOO_MD5 = "d3b07384d113edec49eaa6238ad5ff00"
+BAR_MD5 = "c157a79031e1c40f85931829bc5fc552"
+
+
+@pytest.fixture
+def hardlink_odb(tmp_path_factory):
+    path = tmp_path_factory.mktemp("odb")
+    odb = LocalHashFileDB(localfs, os.fspath(path), type="hardlink")
+    odb.add_bytes(FOO_MD5, b"foo\n")
+    odb.add_bytes(BAR_MD5, b"bar\n")
+    for oid in (FOO_MD5, BAR_MD5):
+        odb.protect(odb.oid_to_path(oid))
+    return odb
+
+
+def _file_index(odb, oid):
+    index = DataIndex(
+        {
+            ("file",): DataIndexEntry(
+                key=("file",),
+                meta=Meta(),
+                hash_info=HashInfo(name="md5", value=oid),
+            ),
+        }
+    )
+    index.storage_map.add_cache(ObjectStorage((), odb))
+    return index
+
+
+def test_checkout_relink_keeps_cache_protected(tmp_path, hardlink_odb):
+    # https://github.com/treeverse/dvc/issues/11119
+    old = _file_index(hardlink_odb, FOO_MD5)
+    apply(compare(None, old), str(tmp_path), localfs, links=["hardlink"])
+    assert localfs.is_hardlink(str(tmp_path / "file"))
+
+    new = _file_index(hardlink_odb, FOO_MD5)
+    apply(compare(old, new, relink=True), str(tmp_path), localfs, links=["hardlink"])
+
+    assert localfs.is_hardlink(str(tmp_path / "file"))
+    assert (tmp_path / "file").read_text() == "foo\n"
+    assert hardlink_odb.is_protected(hardlink_odb.oid_to_path(FOO_MD5))
+
+
+def test_checkout_modified_keeps_old_cache_protected(tmp_path, hardlink_odb):
+    # https://github.com/treeverse/dvc/issues/11119
+    old = _file_index(hardlink_odb, FOO_MD5)
+    apply(compare(None, old), str(tmp_path), localfs, links=["hardlink"])
+
+    new = _file_index(hardlink_odb, BAR_MD5)
+    apply(compare(old, new), str(tmp_path), localfs, links=["hardlink"])
+
+    assert (tmp_path / "file").read_text() == "bar\n"
+    assert hardlink_odb.is_protected(hardlink_odb.oid_to_path(FOO_MD5))
+    assert hardlink_odb.is_protected(hardlink_odb.oid_to_path(BAR_MD5))