Support ignoring files and notebooks (#3653)

<!-- REMOVE IRRELEVANT COMMENTS BEFORE CREATING A PULL REQUEST -->
## Changes
Support ignoring files and notebooks

### Linked issues

Progresses #3645

### Functionality

- [x] modified file and notebook loading in the source code related
modules

### Tests

- [x] added unit tests

Author: JCZuurmond

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

@@ -17,6 +17,7 @@
     BaseImportResolver,
     BaseFileResolver,
     MaybeDependency,
+    StubContainer,
 )
 from databricks.labs.ucx.source_code.known import KnownList
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
@@ -70,24 +71,30 @@ def __repr__(self):
         return f"<LocalFile {self._path}>"
 
 
-class StubContainer(SourceContainer):
-
-    def __init__(self, path: Path):
-        super().__init__()
-        self._path = path
+class FileLoader(DependencyLoader):
+    """Loader for a file dependency.
 
-    def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
-        return []
+    Args:
+        exclude_paths (set[Path] | None) : A set of paths to load as
+            class:`StubContainer`. If None, no paths are excluded.
 
+            Note: The exclude paths are loaded as `StubContainer` to
+            signal that the path is found, however, it should not be
+            processed.
+    """
 
-class FileLoader(DependencyLoader):
-    """Loader for a file dependency."""
+    def __init__(self, *, exclude_paths: set[Path] | None = None):
+        self._exclude_paths = exclude_paths or set[Path]()
 
     def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> LocalFile | StubContainer | None:
         """Load the dependency."""
         resolved_path = path_lookup.resolve(dependency.path)
         if not resolved_path:
             return None
+        if resolved_path in self._exclude_paths:
+            # Paths are excluded from further processing by loading them as stub container.
+            # Note we don't return `None`, as the path is found.
+            return StubContainer(resolved_path)
         language = file_language(resolved_path)
         if not language:
             return StubContainer(resolved_path)

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

@@ -600,3 +600,13 @@ def finalize(self) -> InheritedContext:
             return self
         tree = self.tree.renumber(-1)
         return InheritedContext(tree, self.found, [])
+
+
+class StubContainer(SourceContainer):
+
+    def __init__(self, path: Path):
+        super().__init__()
+        self._path = path
+
+    def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]:
+        return []

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

@@ -1,20 +1,18 @@
 from __future__ import annotations
 
-import abc
 import logging
 from pathlib import Path
 from typing import TypeVar
 
-from databricks.sdk.errors import NotFound
 from databricks.sdk.service.workspace import Language
 
-from databricks.labs.ucx.source_code.base import is_a_notebook, file_language
+from databricks.labs.ucx.source_code.base import is_a_notebook, file_language, safe_read_text
 from databricks.labs.ucx.source_code.graph import (
     BaseNotebookResolver,
     Dependency,
     DependencyLoader,
     MaybeDependency,
-    SourceContainer,
+    StubContainer,
 )
 from databricks.labs.ucx.source_code.notebooks.cells import CellLanguage
 from databricks.labs.ucx.source_code.notebooks.sources import Notebook
@@ -40,10 +38,32 @@ def resolve_notebook(self, path_lookup: PathLookup, path: Path, inherit_context:
         return MaybeDependency(dependency, [])
 
 
-class NotebookLoader(DependencyLoader, abc.ABC):
+class NotebookLoader(DependencyLoader):
+    """Load a notebook.
+
+    Args:
+        exclude_paths (set[Path] | None) : A set of paths to load as
+            class:`StubContainer`. If None, no paths are excluded.
+
+            Note: The exclude paths are loaded as `StubContainer` to
+            signal that the path is found, however, it should not be
+            processed.
+
+    TODO:
+        Let `NotebookLoader` inherit from `FileLoader` and reuse the
+        implementation of `load_dependency` to first load the file, then
+        convert it to a `Notebook` if it is a notebook source.
+    """
+
+    def __init__(self, *, exclude_paths: set[Path] | None = None):
+        self._exclude_paths = exclude_paths or set[Path]()
+
     def resolve(self, path_lookup: PathLookup, path: Path) -> Path | None:
-        """If the path is a Python file, return the path to the Python file. If the path is neither,
-        return None."""
+        """Resolve the notebook path.
+
+        If the path is a Python file, return the path to the Python file. If the path is neither,
+        return None.
+        """
         # check current working directory first
         absolute_path = path_lookup.cwd / path
         absolute_path = absolute_path.resolve()
@@ -58,29 +78,27 @@ def resolve(self, path_lookup: PathLookup, path: Path) -> Path | None:
             return a_path
         return None
 
-    def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> SourceContainer | None:
-        absolute_path = self.resolve(path_lookup, dependency.path)
-        if not absolute_path:
-            return None
-        try:
-            content = absolute_path.read_text("utf-8")
-        except NotFound:
-            logger.warning(f"Path not found trying to read notebook from workspace: {absolute_path}")
+    def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> Notebook | StubContainer | None:
+        """Load the notebook dependency."""
+        resolved_path = self.resolve(path_lookup, dependency.path)
+        if not resolved_path:
             return None
-        except PermissionError:
-            logger.warning(
-                f"Permission error while reading notebook from workspace: {absolute_path}",
-                exc_info=True,
-            )
+        if resolved_path in self._exclude_paths:
+            # Paths are excluded from further processing by loading them as stub container.
+            # Note we don't return `None`, as the path is found.
+            return StubContainer(resolved_path)
+        content = safe_read_text(resolved_path)
+        if content is None:
             return None
-        except UnicodeDecodeError:
-            logger.warning(f"Cannot decode non-UTF-8 encoded notebook from workspace: {absolute_path}")
-            return None
-        language = self._detect_language(absolute_path, content)
+        language = self._detect_language(resolved_path, content)
         if not language:
-            logger.warning(f"Could not detect language for {absolute_path}")
+            logger.warning(f"Could not detect language for {resolved_path}")
+            return None
+        try:
+            return Notebook.parse(resolved_path, content, language)
+        except ValueError as e:
+            logger.warning(f"Could not parse notebook for {resolved_path}", exc_info=e)
             return None
-        return Notebook.parse(absolute_path, content, language)
 
     @staticmethod
     def _detect_language(path: Path, content: str) -> Language | None:

tests/unit/source_code/notebooks/test_loaders.py

@@ -1,12 +1,15 @@
+import codecs
 import logging
 from pathlib import Path
 from unittest.mock import create_autospec
 
 import pytest
 from databricks.sdk.service.workspace import Language
 
-from databricks.labs.ucx.source_code.graph import Dependency
+from databricks.labs.ucx.source_code.base import NOTEBOOK_HEADER
+from databricks.labs.ucx.source_code.graph import Dependency, StubContainer
 from databricks.labs.ucx.source_code.notebooks.loaders import NotebookLoader
+from databricks.labs.ucx.source_code.notebooks.sources import Notebook
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
 
 
@@ -26,25 +29,85 @@ def detect_language(cls, path: Path, content: str):
 
 
 @pytest.mark.parametrize(
-    "error, message",
+    "error",
     [
-        (PermissionError("Permission denied"), "Permission error while reading notebook from workspace"),
-        (
-            UnicodeDecodeError("utf-8", b"\x80\x81\x82", 0, 1, "invalid start byte"),
-            "Cannot decode non-UTF-8 encoded notebook from workspace",
-        ),
+        PermissionError("Permission denied"),
+        FileNotFoundError("File not found"),
+        UnicodeDecodeError("utf-8", b"\x80\x81\x82", 0, 1, "invalid start byte"),
     ],
 )
-def test_notebook_loader_loads_dependency_raises_error(caplog, error: Exception, message: str) -> None:
+def test_notebook_loader_loads_dependency_raises_error(caplog, error: Exception) -> None:
     path = create_autospec(Path)
-    path.read_text.side_effect = error
+    path.suffix = ".py"
+    path.open.side_effect = error
     path_lookup = create_autospec(PathLookup)
     path_lookup.resolve.return_value = path
-    dependency = create_autospec(Dependency)
-    dependency.path.return_value = path
+    dependency = Dependency(NotebookLoader(), path)
 
     with caplog.at_level(logging.WARNING, logger="databricks.labs.ucx.source_code.notebooks.loaders"):
-        found = NotebookLoader().load_dependency(path_lookup, dependency)
+        notebook = dependency.load(path_lookup)
 
-    assert f"{message}: {path}" in caplog.text
-    assert found is None
+    assert f"Could not read file: {path}" in caplog.text
+    assert notebook is None
+
+
+def test_notebook_loader_ignores_loading_non_ascii_file(mock_path_lookup) -> None:
+    dependency = Dependency(NotebookLoader(), Path("nonascii.py"))
+
+    notebook = dependency.load(mock_path_lookup)
+
+    # TODO: Test specific error while loading: https://github.com/databrickslabs/ucx/issues/3584
+    assert notebook is None
+    assert Path("nonascii.py") in mock_path_lookup.successfully_resolved_paths
+
+
+@pytest.mark.parametrize(
+    "bom, encoding",
+    [
+        (codecs.BOM_UTF8, "utf-8"),
+        (codecs.BOM_UTF16_LE, "utf-16-le"),
+        (codecs.BOM_UTF16_BE, "utf-16-be"),
+        (codecs.BOM_UTF32_LE, "utf-32-le"),
+        (codecs.BOM_UTF32_BE, "utf-32-be"),
+    ],
+)
+def test_notebook_loader_loads_file_with_bom(tmp_path, bom, encoding) -> None:
+    path = tmp_path / "file.py"
+    path.write_bytes(bom + f"# {NOTEBOOK_HEADER}\na = 12".encode(encoding))
+    dependency = Dependency(NotebookLoader(), path)
+    path_lookup = create_autospec(PathLookup)
+    path_lookup.resolve.return_value = path
+
+    notebook = dependency.load(path_lookup)
+
+    # TODO: Test specific error while loading: https://github.com/databrickslabs/ucx/issues/3584
+    assert isinstance(notebook, Notebook)
+    assert notebook.original_code == f"# {NOTEBOOK_HEADER}\na = 12"
+    path_lookup.resolve.assert_called_once_with(path)
+
+
+def test_notebook_loader_cannot_load_empty_file(tmp_path) -> None:
+    """Empty file fails because it misses the notebook header."""
+    path = tmp_path / "empty.py"
+    path.write_text("")
+    dependency = Dependency(NotebookLoader(), path)
+    path_lookup = create_autospec(PathLookup)
+    path_lookup.resolve.return_value = path
+
+    notebook = dependency.load(path_lookup)
+
+    # TODO: Test specific error while loading: https://github.com/databrickslabs/ucx/issues/3584
+    assert not notebook
+    path_lookup.resolve.assert_called_once_with(path)
+
+
+def test_notebook_loader_ignores_path_by_loading_it_as_stub_container(tmp_path) -> None:
+    path = tmp_path / "path.py"
+    dependency = Dependency(NotebookLoader(exclude_paths={path}), path)
+    path_lookup = create_autospec(PathLookup)
+    path_lookup.resolve.return_value = path
+
+    stub = dependency.load(path_lookup)
+
+    assert isinstance(stub, StubContainer)
+    path_lookup.resolve.assert_called_once_with(path)

tests/unit/source_code/test_dependencies.py

@@ -5,7 +5,6 @@
     Dependency,
     DependencyProblem,
     DependencyResolver,
-    SourceContainer,
 )
 from databricks.labs.ucx.source_code.files import FileLoader, ImportFileResolver
 from databricks.labs.ucx.source_code.notebooks.loaders import (
@@ -191,7 +190,7 @@ def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> No
 def test_dependency_resolver_raises_problem_with_unloadable_root_notebook(mock_path_lookup) -> None:
 
     class FailingNotebookLoader(NotebookLoader):
-        def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> SourceContainer | None:
+        def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> None:
             return None
 
     notebook_loader = FailingNotebookLoader()

tests/unit/source_code/test_files.py

@@ -6,7 +6,7 @@
 from databricks.sdk.service.workspace import Language
 
 from databricks.labs.ucx.source_code.base import CurrentSessionState
-from databricks.labs.ucx.source_code.graph import Dependency, DependencyGraph, DependencyProblem
+from databricks.labs.ucx.source_code.graph import Dependency, DependencyGraph, DependencyProblem, StubContainer
 from databricks.labs.ucx.source_code.files import FileLoader, LocalFile
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
 
@@ -115,13 +115,13 @@ def test_file_loader_loads_file_without_permission() -> None:
     path_lookup.resolve.assert_called_once_with(path)
 
 
-def test_file_loader_loads_non_ascii_file(mock_path_lookup) -> None:
+def test_file_loader_ignores_loading_non_ascii_file(mock_path_lookup) -> None:
     dependency = Dependency(FileLoader(), Path("nonascii.py"))
 
     local_file = dependency.load(mock_path_lookup)
 
     # TODO: Test specific error while loading: https://github.com/databrickslabs/ucx/issues/3584
-    assert local_file is None
+    assert not local_file
     assert Path("nonascii.py") in mock_path_lookup.successfully_resolved_paths
 
 
@@ -177,3 +177,15 @@ def test_file_loader_loads_empty_file(tmp_path) -> None:
     # TODO: Test specific error while loading: https://github.com/databrickslabs/ucx/issues/3584
     assert local_file
     path_lookup.resolve.assert_called_once_with(path)
+
+
+def test_file_loader_ignores_path_by_loading_it_as_stub_container(tmp_path) -> None:
+    path = tmp_path / "path.py"
+    dependency = Dependency(FileLoader(exclude_paths={path}), path)
+    path_lookup = create_autospec(PathLookup)
+    path_lookup.resolve.return_value = path
+
+    stub = dependency.load(path_lookup)
+
+    assert isinstance(stub, StubContainer)
+    path_lookup.resolve.assert_called_once_with(path)