Compute dependency routes (#2244)

## Changes
Have `DependencyGraph` compute route from root file/notebook to child
file/notebook. That route will be later used to populate an
`InheritedContext`
Have dependencies convey a property for inheriting context, such that
context will be computed when but only when required
This is in preparation of
#2236

### Linked issues
Progresses #2155
Progresses #2156

### Functionality
None

### Tests
- [x] added unit tests

---------

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

Author: ericvergnaud

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

@@ -47,8 +47,8 @@ def path(self):
     def register_library(self, *libraries: str) -> list[DependencyProblem]:
         return self._resolver.register_library(self.path_lookup, *libraries)
 
-    def register_notebook(self, path: Path) -> list[DependencyProblem]:
-        maybe = self._resolver.resolve_notebook(self.path_lookup, path)
+    def register_notebook(self, path: Path, inherit_context: bool) -> list[DependencyProblem]:
+        maybe = self._resolver.resolve_notebook(self.path_lookup, path, inherit_context)
         if not maybe.dependency:
             return maybe.problems
         maybe_graph = self.register_dependency(maybe.dependency)
@@ -199,6 +199,52 @@ def visit(self, visit_node: Callable[[DependencyGraph], bool | None], visited: s
     def new_dependency_graph_context(self):
         return DependencyGraphContext(parent=self, path_lookup=self._path_lookup, session_state=self._session_state)
 
+    def _compute_route(self, root: Path, leaf: Path, visited: set[Path]) -> list[Dependency]:
+        """given 2 files or notebooks root and leaf, compute the list of dependencies that must be traversed
+        in order to lint the leaf in the context of its parents"""
+        try:
+            route = self._do_compute_route(root, leaf, visited)
+            return self._trim_route(route)
+        except ValueError:
+            # we only compute routes on graphs so _compute_route can't fail
+            # but we return an empty route in case it does
+            return []
+
+    def _do_compute_route(self, root: Path, leaf: Path, visited: set[Path]) -> list[Dependency]:
+        maybe = self.locate_dependency(root)
+        if not maybe.graph:
+            logger.warning(f"Could not compute route because dependency {root} cannot be located")
+            raise ValueError()
+        route: list[Dependency] = []
+
+        def do_compute_route(graph: DependencyGraph) -> bool:
+            route.append(graph.dependency)
+            for dependency in graph.local_dependencies:
+                if dependency.path == leaf:
+                    route.append(dependency)
+                    return True
+            for dependency in graph.local_dependencies:
+                sub_route = self._do_compute_route(dependency.path, leaf, visited)
+                if len(sub_route) > 0:
+                    route.extend(sub_route)
+                    return True
+            route.pop()
+            return False
+
+        maybe.graph.visit(do_compute_route, visited)
+        return route
+
+    def _trim_route(self, dependencies: list[Dependency]) -> list[Dependency]:
+        """don't inherit context if dependency is not a file/notebook, or it is loaded via dbutils.notebook.run or via import"""
+        # filter out intermediate containers
+        dependencies = list(dependency for dependency in dependencies if dependency.path.is_file())
+        # restart when not inheriting context
+        for i, dependency in enumerate(dependencies):
+            if dependency.inherits_context:
+                continue
+            return [dependency] + self._trim_route(dependencies[i + 1 :])
+        return dependencies
+
     def __repr__(self):
         return f"<DependencyGraph {self.path}>"
 
@@ -212,14 +258,19 @@ class DependencyGraphContext:
 
 class Dependency(abc.ABC):
 
-    def __init__(self, loader: DependencyLoader, path: Path):
+    def __init__(self, loader: DependencyLoader, path: Path, inherits_context=True):
         self._loader = loader
         self._path = path
+        self._inherits_context = inherits_context
 
     @property
     def path(self) -> Path:
         return self._path
 
+    @property
+    def inherits_context(self):
+        return self._inherits_context
+
     def __hash__(self):
         return hash(self.path)
 
@@ -236,14 +287,12 @@ def __repr__(self):
 class SourceContainer(abc.ABC):
 
     @abc.abstractmethod
-    def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
-        """builds a dependency graph from the contents of this container"""
+    def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]: ...
 
 
 class DependencyLoader(abc.ABC):
     @abc.abstractmethod
-    def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> SourceContainer | None:
-        """loads a dependency"""
+    def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> SourceContainer | None: ...
 
 
 class WrappingLoader(DependencyLoader):
@@ -267,8 +316,7 @@ def register_library(self, path_lookup: PathLookup, *libraries: str) -> list[Dep
 class BaseNotebookResolver(abc.ABC):
 
     @abc.abstractmethod
-    def resolve_notebook(self, path_lookup: PathLookup, path: Path) -> MaybeDependency:
-        """locates a notebook"""
+    def resolve_notebook(self, path_lookup: PathLookup, path: Path, inherit_context: bool) -> MaybeDependency: ...
 
     @staticmethod
     def _fail(code: str, message: str) -> MaybeDependency:
@@ -315,8 +363,8 @@ def __init__(
         self._import_resolver = import_resolver
         self._path_lookup = path_lookup
 
-    def resolve_notebook(self, path_lookup: PathLookup, path: Path) -> MaybeDependency:
-        return self._notebook_resolver.resolve_notebook(path_lookup, path)
+    def resolve_notebook(self, path_lookup: PathLookup, path: Path, inherit_context: bool) -> MaybeDependency:
+        return self._notebook_resolver.resolve_notebook(path_lookup, path, inherit_context)
 
     def resolve_import(self, path_lookup: PathLookup, name: str) -> MaybeDependency:
         return self._import_resolver.resolve_import(path_lookup, name)
@@ -353,7 +401,8 @@ def _local_file_resolver(self) -> BaseFileResolver | None:
     def build_notebook_dependency_graph(self, path: Path, session_state: CurrentSessionState) -> MaybeGraph:
         """Builds a dependency graph starting from a notebook. This method is mainly intended for testing purposes.
         In case of problems, the paths in the problems will be relative to the starting path lookup."""
-        maybe = self._notebook_resolver.resolve_notebook(self._path_lookup, path)
+        # the notebook is the root of the graph, so there's no context to inherit
+        maybe = self._notebook_resolver.resolve_notebook(self._path_lookup, path, inherit_context=False)
         if not maybe.dependency:
             return MaybeGraph(None, self._make_relative_paths(maybe.problems, path))
         graph = DependencyGraph(maybe.dependency, None, self, self._path_lookup, session_state)

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

@@ -57,7 +57,7 @@ def as_message(self) -> str:
 class WorkflowTask(Dependency):
     def __init__(self, ws: WorkspaceClient, task: jobs.Task, job: jobs.Job):
         loader = WrappingLoader(WorkflowTaskContainer(ws, task, job))
-        super().__init__(loader, Path(f'/jobs/{task.task_key}'))
+        super().__init__(loader, Path(f'/jobs/{task.task_key}'), False)
         self._task = task
         self._job = job
 
@@ -188,7 +188,7 @@ def _register_notebook(self, graph: DependencyGraph) -> Iterable[DependencyProbl
         logger.info(f'Discovering {self._task.task_key} entrypoint: {notebook_path}')
         # Notebooks can't be on DBFS.
         path = WorkspacePath(self._ws, notebook_path)
-        return graph.register_notebook(path)
+        return graph.register_notebook(path, False)
 
     def _register_spark_python_task(self, graph: DependencyGraph):
         if not self._task.spark_python_task:
@@ -197,7 +197,7 @@ def _register_spark_python_task(self, graph: DependencyGraph):
         notebook_path = self._task.spark_python_task.python_file
         logger.info(f'Discovering {self._task.task_key} entrypoint: {notebook_path}')
         path = self._as_path(notebook_path)
-        return graph.register_notebook(path)
+        return graph.register_notebook(path, False)
 
     @staticmethod
     def _find_first_matching_distribution(path_lookup: PathLookup, name: str) -> metadata.Distribution | None:
@@ -263,7 +263,8 @@ def _register_pipeline_task(self, graph: DependencyGraph):
                 notebook_path = library.notebook.path
                 # Notebooks can't be on DBFS.
                 path = WorkspacePath(self._ws, notebook_path)
-                yield from graph.register_notebook(path)
+                # the notebook is the root of the graph, so there's no context to inherit
+                yield from graph.register_notebook(path, inherit_context=False)
             if library.jar:
                 yield from self._register_library(graph, compute.Library(jar=library.jar))
             if library.maven:

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

@@ -38,21 +38,21 @@ def __init__(self, path: Path, source: str, language: Language):
         # using CellLanguage so we can reuse the facilities it provides
         self._language = CellLanguage.of_language(language)
 
+    @property
+    def path(self) -> Path:
+        return self._path
+
     def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
         if self._language is CellLanguage.PYTHON:
             context = parent.new_dependency_graph_context()
-            analyser = PythonCodeAnalyzer(context, self._original_code)
-            return analyser.build_graph()
+            analyzer = PythonCodeAnalyzer(context, self._original_code)
+            return analyzer.build_graph()
         # supported language that does not generate dependencies
         if self._language is CellLanguage.SQL:
             return []
         logger.warning(f"Unsupported language: {self._language.language}")
         return []
 
-    @property
-    def path(self) -> Path:
-        return self._path
-
     def __repr__(self):
         return f"<LocalFile {self._path}>"
 
@@ -76,7 +76,7 @@ def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProb
     def _build_dependency_graph(self, parent: DependencyGraph) -> Iterable[DependencyProblem]:
         for child_path in self._path.iterdir():
             loader = self._folder_loader if child_path.is_dir() else self._file_loader
-            dependency = Dependency(loader, child_path)
+            dependency = Dependency(loader, child_path, False)
             yield from parent.register_dependency(dependency).problems
 
     def __repr__(self):
@@ -310,7 +310,7 @@ def _resolve_import(self, path_lookup: PathLookup, name: str) -> MaybeDependency
             absolute_path = path_lookup.resolve(relative_path)
             if not absolute_path:
                 continue
-            dependency = Dependency(self._file_loader, absolute_path)
+            dependency = Dependency(self._file_loader, absolute_path, False)
             return MaybeDependency(dependency, [])
         return None
 

src/databricks/labs/ucx/source_code/notebooks/cells.py

@@ -70,12 +70,10 @@ def migrated_code(self, value: str):
 
     @property
     @abstractmethod
-    def language(self) -> CellLanguage:
-        """returns the language of this cell"""
+    def language(self) -> CellLanguage: ...
 
     @abstractmethod
-    def is_runnable(self) -> bool:
-        """whether of not this cell can be run"""
+    def is_runnable(self) -> bool: ...
 
     def build_dependency_graph(self, _: DependencyGraph) -> list[DependencyProblem]:
         """Check for any problems with dependencies of this cell.
@@ -173,7 +171,7 @@ def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProb
         path, idx, line = self._read_notebook_path()
         if path is not None:
             start_line = self._original_offset + idx
-            problems = parent.register_notebook(path)
+            problems = parent.register_notebook(path, True)
             return [
                 problem.replace(start_line=start_line, start_col=0, end_line=start_line, end_col=len(line))
                 for problem in problems
@@ -467,7 +465,8 @@ def _register_notebook(self, base_node: NotebookRunCall):
                 f"Can't check dependency from {base_node.node.as_string()} because the expression cannot be computed",
             )
         for path in paths:
-            yield from self._context.parent.register_notebook(Path(path))
+            # notebooks ran via dbutils.notebook.run do not inherit or propagate context
+            yield from self._context.parent.register_notebook(Path(path), False)
 
     def _mutate_path_lookup(self, change: SysPathChange):
         if isinstance(change, UnresolvedPath):
@@ -547,7 +546,7 @@ class RunCommand(MagicCommand):
     def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
         path = self.notebook_path
         if path is not None:
-            problems = parent.register_notebook(path)
+            problems = parent.register_notebook(path, True)
             return [problem.from_node(problem.code, problem.message, self._node) for problem in problems]
         problem = DependencyProblem.from_node('invalid-run-cell', "Missing notebook path in %run command", self._node)
         return [problem]

src/databricks/labs/ucx/source_code/notebooks/loaders.py

@@ -26,11 +26,11 @@ def __init__(self, notebook_loader: NotebookLoader):
         super().__init__()
         self._notebook_loader = notebook_loader
 
-    def resolve_notebook(self, path_lookup: PathLookup, path: Path) -> MaybeDependency:
+    def resolve_notebook(self, path_lookup: PathLookup, path: Path, inherit_context: bool) -> MaybeDependency:
         absolute_path = self._notebook_loader.resolve(path_lookup, path)
         if not absolute_path:
             return self._fail('notebook-not-found', f"Notebook not found: {path.as_posix()}")
-        dependency = Dependency(self._notebook_loader, absolute_path)
+        dependency = Dependency(self._notebook_loader, absolute_path, inherit_context)
         return MaybeDependency(dependency, [])
 
 

tests/unit/source_code/samples/functional/_child_that_uses_missing_value.py

@@ -0,0 +1,6 @@
+# Databricks notebook source
+
+# ucx[cannot-autofix-table-reference:+2:0:+2:33] Can't migrate 'spark.table(f'{some_table_name}')' because its table name argument cannot be computed
+# ucx[default-format-changed-in-dbr8:+1:0:+1:33] The default format changed in Databricks Runtime 8.0, from Parquet to Delta
+spark.table(f"{some_table_name}")
+x = 2

tests/unit/source_code/samples/functional/_child_that_uses_value_from_parent.py

@@ -0,0 +1,4 @@
+# Databricks notebook source
+
+# ucx[default-format-changed-in-dbr8:+1:0:+1:33] The default format changed in Databricks Runtime 8.0, from Parquet to Delta
+spark.table(f"{some_table_name}")

tests/unit/source_code/samples/functional/grand_parent_that_dbutils_runs_parent_that_magic_runs_child.py

@@ -0,0 +1,3 @@
+# Databricks notebook source
+
+dbutils.notebook.run("./parent_that_magic_runs_child_that_uses_value_from_parent.py")

tests/unit/source_code/samples/functional/grand_parent_that_imports_parent_that_magic_runs_child.py

@@ -0,0 +1,3 @@
+# Databricks notebook source
+
+from parent_that_magic_runs_child_that_uses_value_from_parent import other_table_name

tests/unit/source_code/samples/functional/grand_parent_that_magic_runs_parent_that_magic_runs_child.py

@@ -0,0 +1,3 @@
+# Databricks notebook source
+
+%run ./parent_that_magic_runs_child_that_uses_value_from_parent.py