diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index dcbcbaaef..7f7b25fb3 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -45,7 +45,7 @@ def fetch( run_cache=False, revs=None, odb: Optional["ObjectDB"] = None, -): +) -> int: """Download data items from a cloud and imported repositories Returns: @@ -58,6 +58,9 @@ def fetch( config.NoRemoteError: thrown when downloading only local files and no remote is configured """ + from dvc.repo.imports import save_imports + from dvc_data.hashfile.transfer import TransferResult + if isinstance(targets, str): targets = [targets] @@ -66,7 +69,6 @@ def fetch( if _remote.worktree: return _fetch_worktree(self, _remote) - downloaded_count = 0 failed_count = 0 try: @@ -75,39 +77,51 @@ def fetch( except DownloadError as exc: failed_count += exc.amount - result = _fetch( - self, - targets, + no_remote_msg: Optional[str] = None + result = TransferResult(set(), set()) + try: + d, f = _fetch( + self, + targets, + all_branches=all_branches, + all_tags=all_tags, + all_commits=all_commits, + with_deps=with_deps, + force=True, + remote=remote, + jobs=jobs, + recursive=recursive, + revs=revs, + odb=odb, + ) + result.transferred.update(d) + result.failed.update(f) + except NoRemoteError as exc: + no_remote_msg = str(exc) + + for rev in self.brancher( + revs=revs, all_branches=all_branches, all_tags=all_tags, all_commits=all_commits, - with_deps=with_deps, - force=True, - remote=remote, - jobs=jobs, - recursive=recursive, - revs=revs, - odb=odb, - ) - downloaded_count = len(result.transferred) + ): + imported = save_imports( + self, + targets, + unpartial=rev == "workspace", + recursive=recursive, + ) + result.transferred.update(imported) + result.failed.difference_update(imported) + failed_count += len(result.failed) - d, f = _fetch_partial_imports( - self, - targets, - all_branches=all_branches, - all_tags=all_tags, - all_commits=all_commits, - recursive=recursive, - revs=revs, - ) - downloaded_count += d - failed_count += f - if failed_count: + if no_remote_msg: + logger.error(no_remote_msg) raise DownloadError(failed_count) - return downloaded_count + return len(result.transferred) def _fetch( @@ -154,39 +168,3 @@ def _fetch( result.transferred.update(d) result.failed.update(f) return result - - -def _fetch_partial_imports(repo, targets, **kwargs): - from dvc.stage.exceptions import DataSourceChanged - - downloaded = 0 - failed = 0 - for stage in repo.partial_imports(targets, **kwargs): - dep = stage.deps[0] - out = stage.outs[0] - try: - if dep.changed_checksum(): - raise DataSourceChanged(f"{stage} ({dep})") - - logger.info("Importing '%s' -> '%s'", dep, out) - if stage.is_repo_import: - obj = dep.get_obj() - meta = dep.get_meta() - out.hash_info = obj.hash_info - out.meta = meta - else: - out.transfer( - dep.def_path, jobs=kwargs.get("jobs"), odb=repo.odb.local - ) - except DataSourceChanged as exc: - logger.warning(f"{exc}") - failed += 1 - continue - if not any( - kwargs.get(kw, None) - for kw in ("all_branches", "all_tags", "all_commits", "revs") - ): - stage.dump() - - downloaded += 1 - return downloaded, failed diff --git a/dvc/repo/imports.py b/dvc/repo/imports.py new file mode 100644 index 000000000..2bd4e6e06 --- /dev/null +++ b/dvc/repo/imports.py @@ -0,0 +1,108 @@ +import logging +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, List, Set, Tuple, Union + +if TYPE_CHECKING: + from dvc.dependency.base import Dependency + from dvc.repo import Repo + from dvc.repo.index import Index, IndexView + from dvc.stage import Stage + from dvc.types import TargetType + from dvc_data.hashfile import HashInfo + +logger = logging.getLogger(__name__) + + +def unfetched_view( + index: "Index", targets: "TargetType", **kwargs +) -> Tuple["IndexView", List["Dependency"]]: + """Return index view of imports which have not been fetched. + + Returns: + Tuple in the form (view, changed_deps) where changed_imports is a list + of import dependencies that cannot be fetched due to changed data + source. + """ + changed_deps: List["Dependency"] = [] + + def need_fetch(stage: "Stage") -> bool: + out = stage.outs[0] + if not out.changed_cache(): + return False + + dep = stage.deps[0] + if dep.changed_checksum(): + changed_deps.append(dep) + return False + + return True + + unfetched = index.targets_view(targets, filter_fn=need_fetch, **kwargs) + return unfetched, changed_deps + + +def unpartial_imports(index: Union["Index", "IndexView"]) -> int: + """Update any outs in the index which are no longer partial imports. + + Returns: + Total number of files which were unpartialed. + """ + updated = 0 + for out in index.outs: + # we need to use view[key] here and since the out fields have not been + # updated yet (out.get_entry() would return the partial-import state) + workspace, key = out.index_key + entry = index.data[workspace][key] + if out.stage.is_partial_import: + out.hash_info = entry.hash_info + out.meta = entry.meta + out.stage.dump() + updated += out.meta.nfiles if out.meta.nfiles is not None else 1 + return updated + + +def save_imports( + repo: "Repo", targets: "TargetType", unpartial: bool = False, **kwargs +) -> Set["HashInfo"]: + """Save (download) imports from their original source location. + + Imports which are already cached will not be downloaded. + + Returns: + Objects which were downloaded from source location. + """ + from dvc.stage.exceptions import DataSourceChanged + from dvc.utils.fs import makedirs + from dvc_data.index import checkout, md5, save + from dvc_objects.fs.callbacks import Callback + + downloaded: Set["HashInfo"] = set() + + unfetched, changed = unfetched_view(repo.index, targets, **kwargs) + for dep in changed: + logger.warning(str(DataSourceChanged(f"{dep.stage} ({dep})"))) + + data_view = unfetched.data["repo"] + if len(data_view): + cache = repo.odb.local + if not cache.fs.exists(cache.path): + makedirs(cache.path) + with TemporaryDirectory(dir=cache.path) as tmpdir: + with Callback.as_tqdm_callback( + desc="Downloading imports from source", + unit="files", + ) as cb: + checkout(data_view, tmpdir, cache.fs, callback=cb) + md5(data_view) + save(data_view, odb=cache, hardlink=True) + + if unpartial: + unpartial_imports(unfetched) + + downloaded.update( + entry.hash_info + for _, entry in data_view.iteritems() + if not entry.meta.isdir and entry.hash_info is not None + ) + + return downloaded