fetch: use data index save() for fetching imports

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

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