provide a callback to handle error during the stage collection (#5020)

also provide a callable to properly handle error messages
add tests

dvc/repo/__init__.py

@@ -8,7 +8,6 @@ from funcy import cached_property, cat
 from git import InvalidGitRepositoryError
 
 from dvc.config import Config
-from dvc.dvcfile import is_valid_filename
 from dvc.exceptions import FileMissingError
 from dvc.exceptions import IsADirectoryError as DvcIsADirectoryError
 from dvc.exceptions import NotDvcRepoError, OutputNotFoundError
@@ -183,6 +182,7 @@ class Repo:
         self.metrics = Metrics(self)
         self.plots = Plots(self)
         self.params = Params(self)
+        self.stage_collection_error_handler = None
         self._lock_depth = 0
 
     @cached_property
@@ -371,25 +371,8 @@ class Repo:
         NOTE: For large repos, this could be an expensive
               operation. Consider using some memoization.
         """
-        return self._collect_stages()
-
-    def _collect_stages(self):
-        stages = []
-        outs = set()
-
-        for root, dirs, files in self.tree.walk(self.root_dir):
-            for file_name in filter(is_valid_filename, files):
-                file_path = os.path.join(root, file_name)
-                new_stages = self.stage.load_file(file_path)
-                stages.extend(new_stages)
-                outs.update(
-                    out.fspath
-                    for stage in new_stages
-                    for out in stage.outs
-                    if out.scheme == "local"
-                )
-            dirs[:] = [d for d in dirs if os.path.join(root, d) not in outs]
-        return stages
+        error_handler = self.stage_collection_error_handler
+        return self.stage.collect_repo(onerror=error_handler)
 
     def find_outs_by_path(self, path, outs=None, recursive=False, strict=True):
         if not outs:

dvc/repo/stage.py

@@ -6,11 +6,15 @@ from contextlib import suppress
 from typing import Iterable, List, NamedTuple, Optional, Set, Tuple
 
 from dvc.dvcfile import PIPELINE_FILE, Dvcfile, is_valid_filename
-from dvc.exceptions import NoOutputOrStageError, OutputNotFoundError
+from dvc.exceptions import (
+    DvcException,
+    NoOutputOrStageError,
+    OutputNotFoundError,
+)
 from dvc.path_info import PathInfo
 from dvc.repo.graph import collect_inside_path, collect_pipeline
 from dvc.stage.exceptions import StageFileDoesNotExistError, StageNotFound
-from dvc.utils import parse_target
+from dvc.utils import parse_target, relpath
 
 logger = logging.getLogger(__name__)
 
@@ -316,3 +320,37 @@ class StageLoad:
                 raise NoOutputOrStageError(target, exc.file) from exc
 
         return [StageInfo(stage) for stage in stages]
+
+    def collect_repo(self, onerror=None):
+        """Collects all of the stages present in the DVC repo.
+
+        Args:
+            onerror (optional): callable that will be called with two args:
+                the filepath whose collection failed and the exc instance.
+                The decision to raise error is with the callback.
+                Otherwise, the file will be skipped, and dvc will partially
+                load the stages.
+        """
+        stages = []
+        outs = set()
+        for root, dirs, files in self.tree.walk(self.repo.root_dir):
+            for file_name in filter(is_valid_filename, files):
+                file_path = os.path.join(root, file_name)
+
+                try:
+                    new_stages = self.load_file(file_path)
+                except DvcException as exc:
+                    if onerror:
+                        onerror(relpath(file_path), exc)
+                        continue
+                    raise
+
+                stages.extend(new_stages)
+                outs.update(
+                    out.fspath
+                    for stage in new_stages
+                    for out in stage.outs
+                    if out.scheme == "local"
+                )
+            dirs[:] = [d for d in dirs if os.path.join(root, d) not in outs]
+        return stages

tests/func/test_remove.py

@@ -23,7 +23,7 @@ def test_remove(tmp_dir, scm, dvc, run_copy, remove_outs):
     for stage in [stage1, stage2, stage3]:
         dvc.remove(stage.addressing, outs=remove_outs)
         out_exists = (out.exists for out in stage.outs)
-        assert stage not in dvc._collect_stages()
+        assert stage not in dvc.stage.collect_repo()
         if remove_outs:
             assert not any(out_exists)
         else:

tests/func/test_stage_load.py

@@ -10,6 +10,7 @@ from dvc.path_info import PathInfo
 from dvc.repo import Repo
 from dvc.stage.exceptions import (
     StageFileDoesNotExistError,
+    StageFileFormatError,
     StageNameUnspecified,
     StageNotFound,
 )
@@ -434,3 +435,19 @@ def test_collect_optimization_on_stage_name(tmp_dir, dvc, mocker, run_copy):
     # Should read stage directly instead of collecting the whole graph
     assert dvc.stage.collect("copy-foo-bar") == [stage]
     assert dvc.stage.collect_granular("copy-foo-bar") == [(stage, None)]
+
+
+def test_collect_repo_callback(tmp_dir, dvc, mocker):
+    mock = mocker.Mock()
+    dvc.stage_collection_error_handler = mock
+
+    (stage,) = tmp_dir.dvc_gen("foo", "foo")
+    dump_yaml(tmp_dir / PIPELINE_FILE, {"stages": {"cmd": "echo hello world"}})
+
+    dvc._reset()
+    assert dvc.stages == [stage]
+    mock.assert_called_once()
+
+    file_path, exc = mock.call_args[0]
+    assert file_path == PIPELINE_FILE
+    assert isinstance(exc, StageFileFormatError)