Implement dependency graph walker (#2502)

## Changes
There is redundant code for walking a dependency graph 'properly' i.e.
from roots to leaves, populating python globals etc...
This PR makes it trivial to reuse the correct code, and reuses it for
local code linting and jobs linting.
This will also be used by #2371 

### Linked issues
None

### Functionality
None

### Tests
- [x] ran existing unit tests

---------

Co-authored-by: Eric Vergnaud <[email protected]>
Co-authored-by: Cor <[email protected]>

Author: ericvergnaud

src/databricks/labs/ucx/source_code/graph.py

@@ -4,7 +4,8 @@
 import logging
 from dataclasses import dataclass
 from pathlib import Path
-from collections.abc import Callable
+from collections.abc import Callable, Iterable, Iterator
+from typing import TypeVar, Generic
 
 from astroid import (  # type: ignore
     NodeNG,
@@ -581,3 +582,45 @@ def finalize(self) -> InheritedContext:
             return self
         tree = self._tree.renumber(-1)
         return InheritedContext(tree, self.found)
+
+
+T = TypeVar("T")
+
+
+class DependencyGraphWalker(abc.ABC, Generic[T]):
+
+    def __init__(self, graph: DependencyGraph, walked_paths: set[Path], path_lookup: PathLookup):
+        self._graph = graph
+        self._walked_paths = walked_paths
+        self._path_lookup = path_lookup
+
+    def __iter__(self) -> Iterator[T]:
+        for dependency in self._graph.root_dependencies:
+            # the dependency is a root, so its path is the one to use
+            # for computing lineage and building python global context
+            root_path = dependency.path
+            yield from self._iter_one(dependency, self._graph, root_path)
+
+    def _iter_one(self, dependency: Dependency, graph: DependencyGraph, root_path: Path) -> Iterable[T]:
+        if dependency.path in self._walked_paths:
+            return
+        self._walked_paths.add(dependency.path)
+        self._log_walk_one(dependency)
+        if dependency.path.is_file() or is_a_notebook(dependency.path):
+            inherited_tree = graph.root.build_inherited_tree(root_path, dependency.path)
+            path_lookup = self._path_lookup.change_directory(dependency.path.parent)
+            yield from self._process_dependency(dependency, path_lookup, inherited_tree)
+            maybe_graph = graph.locate_dependency(dependency.path)
+            # missing graph problems have already been reported while building the graph
+            if maybe_graph.graph:
+                child_graph = maybe_graph.graph
+                for child_dependency in child_graph.local_dependencies:
+                    yield from self._iter_one(child_dependency, child_graph, root_path)
+
+    def _log_walk_one(self, dependency: Dependency):
+        logger.debug(f'Analyzing dependency: {dependency}')
+
+    @abc.abstractmethod
+    def _process_dependency(
+        self, dependency: Dependency, path_lookup: PathLookup, inherited_tree: Tree | None
+    ) -> Iterable[T]: ...

src/databricks/labs/ucx/source_code/jobs.py

@@ -20,16 +20,18 @@
 from databricks.labs.ucx.assessment.crawlers import runtime_version_tuple
 from databricks.labs.ucx.hive_metastore.migration_status import MigrationIndex
 from databricks.labs.ucx.mixins.cached_workspace_path import WorkspaceCache
-from databricks.labs.ucx.source_code.base import CurrentSessionState, is_a_notebook, LocatedAdvice
+from databricks.labs.ucx.source_code.base import CurrentSessionState, LocatedAdvice
 from databricks.labs.ucx.source_code.graph import (
     Dependency,
     DependencyGraph,
     DependencyProblem,
     DependencyResolver,
     SourceContainer,
     WrappingLoader,
+    DependencyGraphWalker,
 )
 from databricks.labs.ucx.source_code.linters.context import LinterContext
+from databricks.labs.ucx.source_code.linters.python_ast import Tree
 from databricks.labs.ucx.source_code.notebooks.sources import FileLinter
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
 
@@ -389,51 +391,53 @@ def _lint_job(self, job: jobs.Job) -> list[JobProblem]:
         return problems
 
     def _lint_task(self, task: jobs.Task, job: jobs.Job, linted_paths: set[Path]) -> Iterable[LocatedAdvice]:
-        dependency: Dependency = WorkflowTask(self._ws, task, job)
+        root_dependency: Dependency = WorkflowTask(self._ws, task, job)
         # we can load it without further preparation since the WorkflowTask is merely a wrapper
-        container = dependency.load(self._path_lookup)
+        container = root_dependency.load(self._path_lookup)
         assert isinstance(container, WorkflowTaskContainer)
         session_state = CurrentSessionState(
             data_security_mode=container.data_security_mode,
             named_parameters=container.named_parameters,
             spark_conf=container.spark_conf,
             dbr_version=container.runtime_version,
         )
-        graph = DependencyGraph(dependency, None, self._resolver, self._path_lookup, session_state)
+        graph = DependencyGraph(root_dependency, None, self._resolver, self._path_lookup, session_state)
         problems = container.build_dependency_graph(graph)
         if problems:
             for problem in problems:
                 source_path = self._UNKNOWN if problem.is_path_missing() else problem.source_path
                 yield LocatedAdvice(problem.as_advisory(), source_path)
             return
-        for dependency in graph.root_dependencies:
-            root = dependency.path  # since it's a root
-            yield from self._lint_one(task, dependency, graph, root, session_state, linted_paths)
+        walker = LintingWalker(
+            graph, linted_paths, self._path_lookup, task.task_key, session_state, self._migration_index
+        )
+        yield from walker
+
 
-    def _lint_one(
+class LintingWalker(DependencyGraphWalker[LocatedAdvice]):
+
+    def __init__(
         self,
-        task: jobs.Task,
-        dependency: Dependency,
         graph: DependencyGraph,
-        root_path: Path,
-        session_state: CurrentSessionState,
         linted_paths: set[Path],
+        path_lookup: PathLookup,
+        key: str,
+        session_state: CurrentSessionState,
+        migration_index: MigrationIndex,
+    ):
+        super().__init__(graph, linted_paths, path_lookup)
+        self._key = key
+        self._session_state = session_state
+        self._migration_index = migration_index
+
+    def _log_walk_one(self, dependency: Dependency):
+        logger.info(f'Linting {self._key} dependency: {dependency}')
+
+    def _process_dependency(
+        self, dependency: Dependency, path_lookup: PathLookup, inherited_tree: Tree | None
     ) -> Iterable[LocatedAdvice]:
-        if dependency.path in linted_paths:
-            return
-        linted_paths.add(dependency.path)
-        logger.info(f'Linting {task.task_key} dependency: {dependency}')
-        if dependency.path.is_file() or is_a_notebook(dependency.path):
-            inherited_tree = graph.root.build_inherited_tree(root_path, dependency.path)
-            ctx = LinterContext(self._migration_index, session_state)
-            path_lookup = self._path_lookup.change_directory(dependency.path.parent)
-            # FileLinter will determine which file/notebook linter to use
-            linter = FileLinter(ctx, path_lookup, session_state, dependency.path, inherited_tree)
-            for advice in linter.lint():
-                yield LocatedAdvice(advice, dependency.path)
-        maybe_graph = graph.locate_dependency(dependency.path)
-        # problems have already been reported while building the graph
-        if maybe_graph.graph:
-            child_graph = maybe_graph.graph
-            for child_dependency in child_graph.local_dependencies:
-                yield from self._lint_one(task, child_dependency, child_graph, root_path, session_state, linted_paths)
+        ctx = LinterContext(self._migration_index, self._session_state)
+        # FileLinter determines which file/notebook linter to use
+        linter = FileLinter(ctx, path_lookup, self._session_state, dependency.path, inherited_tree)
+        for advice in linter.lint():
+            yield LocatedAdvice(advice, dependency.path)

src/databricks/labs/ucx/source_code/linters/files.py

@@ -7,6 +7,7 @@
 from typing import TextIO
 
 from databricks.labs.ucx.source_code.base import LocatedAdvice, CurrentSessionState, file_language, is_a_notebook
+from databricks.labs.ucx.source_code.linters.python_ast import Tree
 from databricks.labs.ucx.source_code.notebooks.loaders import NotebookLoader
 from databricks.labs.ucx.source_code.notebooks.sources import FileLinter
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
@@ -27,6 +28,7 @@
     SourceContainer,
     DependencyResolver,
     InheritedContext,
+    DependencyGraphWalker,
 )
 
 logger = logging.getLogger(__name__)
@@ -141,40 +143,32 @@ def lint_path(self, path: Path, linted_paths: set[Path] | None = None) -> Iterab
         is_dir = path.is_dir()
         loader = self._folder_loader if is_dir else self._file_loader
         path_lookup = self._path_lookup.change_directory(path if is_dir else path.parent)
-        dependency = Dependency(loader, path, not is_dir)  # don't inherit context when traversing folders
-        graph = DependencyGraph(dependency, None, self._dependency_resolver, path_lookup, self._session_state)
-        container = dependency.load(path_lookup)
+        root_dependency = Dependency(loader, path, not is_dir)  # don't inherit context when traversing folders
+        graph = DependencyGraph(root_dependency, None, self._dependency_resolver, path_lookup, self._session_state)
+        container = root_dependency.load(path_lookup)
         assert container is not None  # because we just created it
         problems = container.build_dependency_graph(graph)
         for problem in problems:
             problem_path = Path('UNKNOWN') if problem.is_path_missing() else problem.source_path.absolute()
             yield problem.as_advisory().for_path(problem_path)
+        context_factory = self._context_factory
+        session_state = self._session_state
+
+        class LintingWalker(DependencyGraphWalker[LocatedAdvice]):
+
+            def _process_dependency(
+                self, dependency: Dependency, path_lookup: PathLookup, inherited_tree: Tree | None
+            ) -> Iterable[LocatedAdvice]:
+                ctx = context_factory()
+                # FileLinter will determine which file/notebook linter to use
+                linter = FileLinter(ctx, path_lookup, session_state, dependency.path, inherited_tree)
+                for advice in linter.lint():
+                    yield LocatedAdvice(advice, dependency.path)
+
         if linted_paths is None:
             linted_paths = set()
-        for dependency in graph.root_dependencies:
-            root = dependency.path  # since it's a root
-            yield from self._lint_one(dependency, graph, root, linted_paths)
-
-    def _lint_one(
-        self, dependency: Dependency, graph: DependencyGraph, root_path: Path, linted_paths: set[Path]
-    ) -> Iterable[LocatedAdvice]:
-        if dependency.path in linted_paths:
-            return
-        linted_paths.add(dependency.path)
-        if dependency.path.is_file() or is_a_notebook(dependency.path):
-            inherited_tree = graph.root.build_inherited_tree(root_path, dependency.path)
-            ctx = self._context_factory()
-            path_lookup = self._path_lookup.change_directory(dependency.path.parent)
-            # FileLinter will determine which file/notebook linter to use
-            linter = FileLinter(ctx, path_lookup, self._session_state, dependency.path, inherited_tree)
-            for advice in linter.lint():
-                yield advice.for_path(dependency.path)
-        maybe_graph = graph.locate_dependency(dependency.path)
-        # problems have already been reported while building the graph
-        if maybe_graph.graph:
-            child_graph = maybe_graph.graph
-            for child_dependency in child_graph.local_dependencies:
-                yield from self._lint_one(child_dependency, child_graph, root_path, linted_paths)
+        walker = LintingWalker(graph, linted_paths, self._path_lookup)
+        yield from walker
 
 
 class LocalFileMigrator: