diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index ecfc39221..49ce64a68 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -23,6 +23,7 @@ if TYPE_CHECKING: from dvc.stage import Stage from dvc.types import DictStrAny from dvc_data.hashfile.state import StateBase + from dvc_data.index import DataIndex from .experiments import Experiments from .index import Index @@ -179,6 +180,7 @@ class Repo: self._fs_conf = {"repo_factory": repo_factory} self._fs = fs or localfs self._scm = scm + self._data_index = None if rev and not fs: self._scm = scm = SCM(root_dir or os.curdir) @@ -341,6 +343,31 @@ class Repo: # fs. self._reset() + @property + def data_index(self) -> Optional["DataIndex"]: + from appdirs import user_cache_dir + from fsspec.utils import tokenize + + from dvc_data.index import DataIndex + + if not self.config["feature"].get("data_index_cache"): + return None + + if self._data_index is None: + cache_dir = user_cache_dir(self.config.APPNAME, self.config.APPAUTHOR) + index_dir = os.path.join( + cache_dir, + "index", + "data", + # scm.root_dir and repo.root_dir don't match for subrepos + tokenize((self.scm.root_dir, self.root_dir)), + ) + os.makedirs(index_dir, exist_ok=True) + + self._data_index = DataIndex.open(os.path.join(index_dir, "db.db")) + + return self._data_index + def __repr__(self): return f"{self.__class__.__name__}: '{self.root_dir}'" @@ -553,6 +580,8 @@ class Repo: def close(self): self.scm.close() self.state.close() + if self._data_index is not None: + self._data_index.close() def _reset_cached_indecies(self): self.__dict__.pop("index", None) diff --git a/dvc/repo/index.py b/dvc/repo/index.py index dc241efa0..1b9fef083 100644 --- a/dvc/repo/index.py +++ b/dvc/repo/index.py @@ -358,31 +358,13 @@ class Index: prefix: "DataIndexKey" loaded = False - if self.repo.config["feature"].get("data_index_cache"): - import os - - from appdirs import user_cache_dir - from fsspec.utils import tokenize - - cache_dir = user_cache_dir( - self.repo.config.APPNAME, self.repo.config.APPAUTHOR - ) - index_dir = os.path.join( - cache_dir, - "index", - "data", - # scm.root_dir and repo.root_dir don't match for subrepos - tokenize((self.repo.scm.root_dir, self.repo.root_dir)), - ) - os.makedirs(index_dir, exist_ok=True) - - index = DataIndex.open(os.path.join(index_dir, "db.db")) - prefix = (self.data_tree.hash_info.value,) - if prefix in index.ls((), detail=False): - loaded = True - else: + index = self.repo.data_index + if index is None: index = DataIndex() - prefix = () + + prefix = (self.data_tree.hash_info.value,) + if prefix in index.ls((), detail=False): + loaded = True try: if not loaded: diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index d1509b649..42b23ab42 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -199,7 +199,8 @@ def test_verify_hashes(tmp_dir, scm, dvc, mocker, tmp_path_factory, local_remote hash_spy = mocker.spy(dvc_data.hashfile.hash, "file_md5") dvc.pull() - assert hash_spy.call_count == 0 + # NOTE: 1 is for index.data_tree building + assert hash_spy.call_count == 1 # Removing cache will invalidate existing state entries dvc.cache.local.clear() @@ -207,7 +208,7 @@ def test_verify_hashes(tmp_dir, scm, dvc, mocker, tmp_path_factory, local_remote dvc.config["remote"]["upstream"]["verify"] = True dvc.pull() - assert hash_spy.call_count == 4 + assert hash_spy.call_count == 6 @flaky(max_runs=3, min_passes=1)