[internal] Simplify import resolvers (#1746)

## Changes Remove chain of import resolvers in favor of a single resolver. Supersedes #1709 Resolves #1707 --------- Co-authored-by: Eric Vergnaud <eic.vergnaud@databricks.com>
databrickslabs · May 23, 2024 · 4ea2b0b · 4ea2b0b
1 parent f0b3341
commit 4ea2b0b
Show file tree

Hide file tree

Showing 10 changed files with 108 additions and 171 deletions.
diff --git a/src/databricks/labs/ucx/contexts/application.py b/src/databricks/labs/ucx/contexts/application.py
@@ -40,10 +40,10 @@
     NotebookResolver,
     NotebookLoader,
 )
-from databricks.labs.ucx.source_code.files import FileLoader, LocalFileResolver
+from databricks.labs.ucx.source_code.files import FileLoader, ImportFileResolver
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
 from databricks.labs.ucx.source_code.graph import DependencyResolver
-from databricks.labs.ucx.source_code.whitelist import WhitelistResolver, Whitelist
+from databricks.labs.ucx.source_code.whitelist import Whitelist
 from databricks.labs.ucx.source_code.languages import Languages
 from databricks.labs.ucx.source_code.redash import Redash
 from databricks.labs.ucx.workspace_access import generic, redash
@@ -381,19 +381,14 @@ def whitelist(self):
         # TODO: fill in the whitelist
         return Whitelist()
 
-    @cached_property
-    def whitelist_resolver(self):
-        return WhitelistResolver(self.whitelist)
-
     @cached_property
     def file_resolver(self):
-        return LocalFileResolver(self.file_loader)
+        return ImportFileResolver(self.file_loader, self.whitelist)
 
     @cached_property
     def dependency_resolver(self):
-        import_resolvers = [self.file_resolver, self.whitelist_resolver]
         library_resolvers = [self.pip_resolver]
-        return DependencyResolver(library_resolvers, self.notebook_resolver, import_resolvers, self.path_lookup)
+        return DependencyResolver(library_resolvers, self.notebook_resolver, self.file_resolver, self.path_lookup)
 
     @cached_property
     def workflow_linter(self):

diff --git a/src/databricks/labs/ucx/source_code/files.py b/src/databricks/labs/ucx/source_code/files.py
@@ -4,6 +4,7 @@
 from pathlib import Path
 
 from databricks.labs.ucx.source_code.path_lookup import PathLookup
+from databricks.labs.ucx.source_code.whitelist import Whitelist, UCCompatibility
 from databricks.sdk.service.workspace import Language
 
 from databricks.labs.ucx.source_code.languages import Languages
@@ -109,15 +110,13 @@ def __repr__(self):
         return "FileLoader()"
 
 
-class LocalFileResolver(BaseImportResolver, BaseFileResolver):
+class ImportFileResolver(BaseImportResolver, BaseFileResolver):
 
-    def __init__(self, file_loader: FileLoader, next_resolver: BaseImportResolver | None = None):
-        super().__init__(next_resolver)
+    def __init__(self, file_loader: FileLoader, whitelist: Whitelist):
+        super().__init__()
+        self._whitelist = whitelist
         self._file_loader = file_loader
 
-    def with_next_resolver(self, resolver: BaseImportResolver) -> BaseImportResolver:
-        return LocalFileResolver(self._file_loader, resolver)
-
     def resolve_local_file(self, path_lookup, path: Path) -> MaybeDependency:
         absolute_path = path_lookup.resolve(path)
         if absolute_path:
@@ -126,6 +125,29 @@ def resolve_local_file(self, path_lookup, path: Path) -> MaybeDependency:
         return MaybeDependency(None, [problem])
 
     def resolve_import(self, path_lookup: PathLookup, name: str) -> MaybeDependency:
+        maybe = self._resolve_whitelist(name)
+        if maybe is not None:
+            return maybe
+        maybe = self._resolve_import(path_lookup, name)
+        if maybe is not None:
+            return maybe
+        return self._fail('import-not-found', f"Could not locate import: {name}")
+
+    def _resolve_whitelist(self, name: str) -> MaybeDependency | None:
+        # TODO attach compatibility to dependency, see https://github.com/databrickslabs/ucx/issues/1382
+        compatibility = self._whitelist.compatibility(name)
+        if compatibility == UCCompatibility.FULL:
+            return MaybeDependency(None, [])
+        if compatibility == UCCompatibility.NONE:
+            # TODO move to linter, see https://github.com/databrickslabs/ucx/issues/1527
+            problem = DependencyProblem("dependency-check", f"Use of dependency {name} is deprecated")
+            return MaybeDependency(None, [problem])
+        if compatibility == UCCompatibility.PARTIAL:
+            problem = DependencyProblem("dependency-check", f"Package {name} is only partially supported by UC")
+            return MaybeDependency(None, [problem])
+        return None
+
+    def _resolve_import(self, path_lookup: PathLookup, name: str) -> MaybeDependency | None:
         if not name:
             return MaybeDependency(None, [DependencyProblem("ucx-bug", "Import name is empty")])
         parts = []
@@ -148,7 +170,11 @@ def resolve_import(self, path_lookup: PathLookup, name: str) -> MaybeDependency:
                 continue
             dependency = Dependency(self._file_loader, absolute_path)
             return MaybeDependency(dependency, [])
-        return super().resolve_import(path_lookup, name)
+        return None
+
+    @staticmethod
+    def _fail(code: str, message: str):
+        return MaybeDependency(None, [DependencyProblem(code, message)])
 
     def __repr__(self):
         return "LocalFileResolver()"
diff --git a/src/databricks/labs/ucx/source_code/graph.py b/src/databricks/labs/ucx/source_code/graph.py
@@ -277,21 +277,9 @@ def _fail(code: str, message: str) -> MaybeDependency:
 
 class BaseImportResolver(abc.ABC):
 
-    def __init__(self, next_resolver: BaseImportResolver | None):
-        self._next_resolver = next_resolver
-
     @abc.abstractmethod
-    def with_next_resolver(self, resolver: BaseImportResolver) -> BaseImportResolver:
-        """required to create a linked list of resolvers"""
-
-    @property
-    def next_resolver(self):
-        return self._next_resolver
-
     def resolve_import(self, path_lookup: PathLookup, name: str) -> MaybeDependency:
-        # TODO: remove StubResolver and return MaybeDependency(None, [...])
-        assert self._next_resolver is not None
-        return self._next_resolver.resolve_import(path_lookup, name)
+        """resolve a simple or composite import name"""
 
 
 class BaseFileResolver(abc.ABC):
@@ -317,22 +305,6 @@ def _fail(code: str, message: str):
         return MaybeDependency(None, [DependencyProblem(code, message)])
 
 
-class StubImportResolver(BaseImportResolver):
-
-    def __init__(self):
-        super().__init__(None)
-
-    def with_next_resolver(self, resolver: BaseImportResolver) -> BaseImportResolver:
-        raise NotImplementedError("Should never happen!")
-
-    def resolve_import(self, path_lookup: PathLookup, name: str) -> MaybeDependency:
-        return self._fail('import-not-found', f"Could not locate import: {name}")
-
-    @staticmethod
-    def _fail(code: str, message: str):
-        return MaybeDependency(None, [DependencyProblem(code, message)])
-
-
 @dataclass
 class MaybeDependency:
     dependency: Dependency | None
@@ -351,12 +323,12 @@ def __init__(
         self,
         library_resolvers: list[BaseLibraryResolver],
         notebook_resolver: BaseNotebookResolver,
-        import_resolvers: list[BaseImportResolver],
+        import_resolver: BaseImportResolver,
         path_lookup: PathLookup,
     ):
         self._library_resolver = self._chain_library_resolvers(library_resolvers)
         self._notebook_resolver = notebook_resolver
-        self._import_resolver = self._chain_import_resolvers(import_resolvers)
+        self._import_resolver = import_resolver
         self._path_lookup = path_lookup
 
     @staticmethod
@@ -367,14 +339,6 @@ def _chain_library_resolvers(library_resolvers: list[BaseLibraryResolver]) -> Ba
             previous = resolver
         return previous
 
-    @staticmethod
-    def _chain_import_resolvers(import_resolvers: list[BaseImportResolver]) -> BaseImportResolver:
-        previous: BaseImportResolver = StubImportResolver()
-        for resolver in import_resolvers:
-            resolver = resolver.with_next_resolver(previous)
-            previous = resolver
-        return previous
-
     def resolve_notebook(self, path_lookup: PathLookup, path: Path) -> MaybeDependency:
         return self._notebook_resolver.resolve_notebook(path_lookup, path)
 
@@ -406,11 +370,8 @@ def build_local_file_dependency_graph(self, path: Path) -> MaybeGraph:
 
     @property
     def _local_file_resolver(self) -> BaseFileResolver | None:
-        resolver = self._import_resolver
-        while resolver is not None:
-            if isinstance(resolver, BaseFileResolver):
-                return resolver
-            resolver = resolver.next_resolver
+        if isinstance(self._import_resolver, BaseFileResolver):
+            return self._import_resolver
         return None
 
     def build_notebook_dependency_graph(self, path: Path) -> MaybeGraph:

diff --git a/src/databricks/labs/ucx/source_code/whitelist.py b/src/databricks/labs/ucx/source_code/whitelist.py
@@ -9,43 +9,16 @@
 
 from yaml import load_all as load_yaml, Loader
 
-from databricks.labs.ucx.source_code.path_lookup import PathLookup
 
 from databricks.labs.ucx.source_code.graph import (
-    BaseImportResolver,
     DependencyGraph,
     DependencyProblem,
-    MaybeDependency,
     SourceContainer,
 )
 
 logger = logging.getLogger(__name__)
 
 
-class WhitelistResolver(BaseImportResolver):
-
-    def __init__(self, whitelist: Whitelist, next_resolver: BaseImportResolver | None = None):
-        super().__init__(next_resolver)
-        self._whitelist = whitelist
-
-    def with_next_resolver(self, resolver: BaseImportResolver) -> BaseImportResolver:
-        return WhitelistResolver(self._whitelist, resolver)
-
-    def resolve_import(self, path_lookup: PathLookup, name: str) -> MaybeDependency:
-        # TODO attach compatibility to dependency, see https://github.com/databrickslabs/ucx/issues/1382
-        compatibility = self._whitelist.compatibility(name)
-        if compatibility == UCCompatibility.FULL:
-            return MaybeDependency(None, [])
-        if compatibility == UCCompatibility.NONE:
-            # TODO move to linter, see https://github.com/databrickslabs/ucx/issues/1527
-            problem = DependencyProblem("dependency-check", f"Use of dependency {name} is deprecated")
-            return MaybeDependency(None, [problem])
-        if compatibility == UCCompatibility.PARTIAL:
-            problem = DependencyProblem("dependency-check", f"Package {name} is only partially supported by UC")
-            return MaybeDependency(None, [problem])
-        return super().resolve_import(path_lookup, name)
-
-
 class StubContainer(SourceContainer):
 
     def build_dependency_graph(self, parent: DependencyGraph) -> list[DependencyProblem]: