diff --git a/src/databricks/labs/ucx/contexts/application.py b/src/databricks/labs/ucx/contexts/application.py index 4f98c18957..41d4953056 100644 --- 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 index 763f9a94ce..f4a72204e4 100644 --- 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 index f061329f51..213ad6def3 100644 --- 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 index 07c793f4a0..1eaa249f62 100644 --- 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]: diff --git a/tests/unit/source_code/test_dependencies.py b/tests/unit/source_code/test_dependencies.py index c2cdb9c1fd..02dd0ad1a9 100644 --- a/tests/unit/source_code/test_dependencies.py +++ b/tests/unit/source_code/test_dependencies.py @@ -10,10 +10,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.python_libraries import PipResolver -from databricks.labs.ucx.source_code.whitelist import WhitelistResolver, Whitelist +from databricks.labs.ucx.source_code.whitelist import Whitelist from tests.unit import ( locate_site_packages, _samples_path, @@ -46,30 +46,23 @@ def test_dependency_resolver_visits_local_notebook_dependencies(mock_path_lookup def test_dependency_resolver_visits_workspace_file_dependencies(mock_path_lookup): - whi = Whitelist() + whitelist = Whitelist() file_loader = FileLoader() notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [ - WhitelistResolver(whi), - LocalFileResolver(file_loader), - ] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(file_loader, whitelist) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path('./root8.py')) assert not maybe.failed assert maybe.graph.all_relative_names() == {'leaf1.py', 'leaf2.py', 'root8.py'} def test_dependency_resolver_raises_problem_with_unfound_workspace_notebook_dependency(mock_path_lookup): - whi = Whitelist() file_loader = FileLoader() notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [ - WhitelistResolver(whi), - LocalFileResolver(file_loader), - ] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(file_loader, Whitelist()) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_notebook_dependency_graph(Path("root1-no-leaf.run.py")) assert list(maybe.problems) == [ DependencyProblem( @@ -127,8 +120,8 @@ def test_dependency_resolver_raises_problem_with_invalid_run_cell(mock_path_look def test_dependency_resolver_visits_recursive_file_dependencies(mock_path_lookup): notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [LocalFileResolver(FileLoader())] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("root6.py")) assert not maybe.failed assert maybe.graph.all_relative_names() == {"root6.py", "root5.py", "leaf4.py"} @@ -137,8 +130,8 @@ def test_dependency_resolver_visits_recursive_file_dependencies(mock_path_lookup def test_dependency_resolver_raises_problem_with_unresolved_import(mock_path_lookup): notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [LocalFileResolver(FileLoader())] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path('root7.py')) assert list(maybe.problems) == [ DependencyProblem('import-not-found', 'Could not locate import: some_library', Path("root7.py"), 1, 0, 1, 19) @@ -148,8 +141,8 @@ def test_dependency_resolver_raises_problem_with_unresolved_import(mock_path_loo def test_dependency_resolver_raises_problem_with_non_constant_notebook_argument(mock_path_lookup): notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [WhitelistResolver(Whitelist()), LocalFileResolver(FileLoader())] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("run_notebooks.py")) assert list(maybe.problems) == [ DependencyProblem( @@ -167,8 +160,8 @@ def test_dependency_resolver_raises_problem_with_non_constant_notebook_argument( def test_dependency_resolver_visits_file_dependencies(mock_path_lookup): notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [LocalFileResolver(FileLoader())] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("root5.py")) assert not maybe.failed assert maybe.graph.all_relative_names() == {"root5.py", "leaf4.py"} @@ -177,8 +170,8 @@ def test_dependency_resolver_visits_file_dependencies(mock_path_lookup): def test_dependency_resolver_skips_builtin_dependencies(mock_path_lookup): notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [WhitelistResolver(Whitelist()), LocalFileResolver(FileLoader())] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("python_builtins.py")) assert not maybe.failed graph = maybe.graph @@ -191,8 +184,8 @@ def test_dependency_resolver_skips_builtin_dependencies(mock_path_lookup): def test_dependency_resolver_ignores_known_dependencies(mock_path_lookup): notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [WhitelistResolver(Whitelist()), LocalFileResolver(FileLoader())] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("python_builtins.py")) assert maybe.graph graph = maybe.graph @@ -205,11 +198,8 @@ def test_dependency_resolver_visits_site_packages(empty_index, mock_notebook_res lookup = PathLookup.from_pathlike_string(Path.cwd(), _samples_path(SourceContainer)) lookup.append_path(site_packages_path) file_loader = FileLoader() - import_resolvers = [ - WhitelistResolver(Whitelist()), - LocalFileResolver(file_loader), - ] - dependency_resolver = DependencyResolver([], mock_notebook_resolver, import_resolvers, lookup) + import_resolver = ImportFileResolver(file_loader, Whitelist()) + dependency_resolver = DependencyResolver([], mock_notebook_resolver, import_resolver, lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("import-site-package.py")) assert not maybe.failed graph = maybe.graph @@ -229,11 +219,8 @@ def test_dependency_resolver_resolves_sub_site_package(): file_loader = FileLoader() notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [ - LocalFileResolver(file_loader), - WhitelistResolver(whitelist), - ] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, lookup) + import_resolver = ImportFileResolver(file_loader, whitelist) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("import-sub-site-package.py")) assert maybe.graph maybe = maybe.graph.locate_dependency(Path(site_packages_path, "databricks", "labs", "lsql", "core.py")) @@ -241,8 +228,8 @@ def test_dependency_resolver_resolves_sub_site_package(): def test_dependency_resolver_raises_problem_with_unfound_root_file(mock_path_lookup, mock_notebook_resolver): - import_resolvers = [LocalFileResolver(FileLoader())] - dependency_resolver = DependencyResolver([], mock_notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + dependency_resolver = DependencyResolver([], mock_notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("non-existing.py")) assert list(maybe.problems) == [ DependencyProblem('file-not-found', 'File not found: non-existing.py', Path("non-existing.py")) @@ -266,8 +253,8 @@ def load_dependency(self, path_lookup: PathLookup, dependency: Dependency) -> So return None file_loader = FailingFileLoader() - import_resolvers = [LocalFileResolver(file_loader)] - dependency_resolver = DependencyResolver([], mock_notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(file_loader, Whitelist()) + dependency_resolver = DependencyResolver([], mock_notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(Path("import-sub-site-package.py")) assert list(maybe.problems) == [ DependencyProblem( @@ -301,8 +288,13 @@ def test_dependency_resolver_raises_problem_with_missing_file_loader(mock_notebo def test_dependency_resolver_resolves_already_installed_library_dependency(mock_notebook_resolver): path_lookup = PathLookup.from_sys_path(Path.cwd()) + file_loader = FileLoader() + whitelist = Whitelist() dependency_resolver = DependencyResolver( - [PipResolver(FileLoader(), Whitelist())], mock_notebook_resolver, [], path_lookup + [PipResolver(file_loader, whitelist)], + mock_notebook_resolver, + ImportFileResolver(file_loader, whitelist), + path_lookup, ) maybe = dependency_resolver.build_library_dependency_graph(Path("astroid")) library_graph = maybe.graph diff --git a/tests/unit/source_code/test_files.py b/tests/unit/source_code/test_files.py index 2c4aedfdc0..4d40a57f6a 100644 --- a/tests/unit/source_code/test_files.py +++ b/tests/unit/source_code/test_files.py @@ -2,10 +2,11 @@ from pathlib import Path from unittest.mock import Mock, create_autospec +from databricks.labs.ucx.source_code.whitelist import Whitelist from databricks.sdk.service.workspace import Language from databricks.labs.ucx.hive_metastore.migration_status import MigrationIndex -from databricks.labs.ucx.source_code.files import LocalFileMigrator, LocalFileResolver, FileLoader +from databricks.labs.ucx.source_code.files import LocalFileMigrator, ImportFileResolver, FileLoader from databricks.labs.ucx.source_code.languages import Languages from databricks.labs.ucx.source_code.path_lookup import PathLookup @@ -75,8 +76,7 @@ def test_files_walks_directory(): def test_triple_dot_import(): - file_loader = FileLoader() - file_resolver = LocalFileResolver(file_loader) + file_resolver = ImportFileResolver(FileLoader(), Whitelist()) path_lookup = create_autospec(PathLookup) path_lookup.cwd.as_posix.return_value = '/some/path/to/folder' path_lookup.resolve.return_value = Path('/some/path/foo.py') @@ -88,8 +88,7 @@ def test_triple_dot_import(): def test_single_dot_import(): - file_loader = FileLoader() - file_resolver = LocalFileResolver(file_loader) + file_resolver = ImportFileResolver(FileLoader(), Whitelist()) path_lookup = create_autospec(PathLookup) path_lookup.cwd.as_posix.return_value = '/some/path/to/folder' path_lookup.resolve.return_value = Path('/some/path/to/folder/foo.py') diff --git a/tests/unit/source_code/test_graph.py b/tests/unit/source_code/test_graph.py index 8a502698a3..bac0b797ba 100644 --- a/tests/unit/source_code/test_graph.py +++ b/tests/unit/source_code/test_graph.py @@ -1,18 +1,20 @@ from pathlib import Path -from databricks.labs.ucx.source_code.files import FileLoader +from databricks.labs.ucx.source_code.files import FileLoader, ImportFileResolver from databricks.labs.ucx.source_code.graph import Dependency, DependencyGraph, DependencyResolver, StubLibraryResolver from databricks.labs.ucx.source_code.notebooks.loaders import NotebookResolver, NotebookLoader from databricks.labs.ucx.source_code.python_libraries import PipResolver -from databricks.labs.ucx.source_code.whitelist import WhitelistResolver, Whitelist +from databricks.labs.ucx.source_code.whitelist import Whitelist def test_dependency_graph_registers_library(mock_path_lookup): dependency = Dependency(FileLoader(), Path("test")) + file_loader = FileLoader() + whitelist = Whitelist() dependency_resolver = DependencyResolver( - [PipResolver(FileLoader(), Whitelist())], + [PipResolver(file_loader, whitelist)], NotebookResolver(NotebookLoader()), - [WhitelistResolver(Whitelist())], + ImportFileResolver(file_loader, whitelist), mock_path_lookup, ) graph = DependencyGraph(dependency, None, dependency_resolver, mock_path_lookup) diff --git a/tests/unit/source_code/test_jobs.py b/tests/unit/source_code/test_jobs.py index 6342d2cb15..3e157c9dfc 100644 --- a/tests/unit/source_code/test_jobs.py +++ b/tests/unit/source_code/test_jobs.py @@ -8,7 +8,7 @@ from databricks.sdk import WorkspaceClient from databricks.sdk.service import compute, jobs -from databricks.labs.ucx.source_code.files import FileLoader +from databricks.labs.ucx.source_code.files import FileLoader, ImportFileResolver from databricks.labs.ucx.source_code.graph import Dependency, DependencyGraph, DependencyResolver from databricks.labs.ucx.source_code.jobs import JobProblem, WorkflowLinter, WorkflowTaskContainer from databricks.labs.ucx.source_code.notebooks.loaders import NotebookResolver, NotebookLoader @@ -26,8 +26,13 @@ def test_job_problem_as_message(): @pytest.fixture def dependency_resolver(mock_path_lookup) -> DependencyResolver: + file_loader = FileLoader() + whitelist = Whitelist() resolver = DependencyResolver( - [PipResolver(FileLoader(), Whitelist())], NotebookResolver(NotebookLoader()), [], mock_path_lookup + [PipResolver(file_loader, whitelist)], + NotebookResolver(NotebookLoader()), + ImportFileResolver(file_loader, whitelist), + mock_path_lookup, ) return resolver diff --git a/tests/unit/source_code/test_path_lookup_simulation.py b/tests/unit/source_code/test_path_lookup_simulation.py index 1e97bbc2ac..f5959fbc40 100644 --- a/tests/unit/source_code/test_path_lookup_simulation.py +++ b/tests/unit/source_code/test_path_lookup_simulation.py @@ -2,11 +2,11 @@ from tempfile import TemporaryDirectory import pytest -from databricks.labs.ucx.source_code.files import LocalFileResolver, FileLoader +from databricks.labs.ucx.source_code.files import ImportFileResolver, FileLoader from databricks.labs.ucx.source_code.path_lookup import PathLookup from databricks.labs.ucx.source_code.graph import SourceContainer, DependencyResolver from databricks.labs.ucx.source_code.notebooks.loaders import NotebookResolver, NotebookLoader -from databricks.labs.ucx.source_code.whitelist import WhitelistResolver, Whitelist +from databricks.labs.ucx.source_code.whitelist import Whitelist from tests.unit import ( _samples_path, ) @@ -41,11 +41,8 @@ def test_locates_notebooks(source: list[str], expected: int, mock_path_lookup): file_loader = FileLoader() notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [ - WhitelistResolver(Whitelist()), - LocalFileResolver(file_loader), - ] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(file_loader, Whitelist()) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_notebook_dependency_graph(notebook_path) assert not maybe.problems assert maybe.graph is not None @@ -69,11 +66,8 @@ def test_locates_files(source: list[str], expected: int): file_loader = FileLoader() notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [ - WhitelistResolver(whitelist), - LocalFileResolver(file_loader), - ] - resolver = DependencyResolver([], notebook_resolver, import_resolvers, lookup) + import_resolver = ImportFileResolver(file_loader, whitelist) + resolver = DependencyResolver([], notebook_resolver, import_resolver, lookup) maybe = resolver.build_local_file_dependency_graph(file_path) assert not maybe.problems assert maybe.graph is not None @@ -105,16 +99,11 @@ def test_locates_notebooks_with_absolute_path(): """, "utf-8", ) - whitelist = Whitelist() lookup = PathLookup.from_sys_path(Path.cwd()) - file_loader = FileLoader() notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [ - WhitelistResolver(whitelist), - LocalFileResolver(file_loader), - ] - resolver = DependencyResolver([], notebook_resolver, import_resolvers, lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + resolver = DependencyResolver([], notebook_resolver, import_resolver, lookup) maybe = resolver.build_notebook_dependency_graph(parent_file_path) assert not maybe.problems assert maybe.graph is not None @@ -146,16 +135,11 @@ def func(): """, "utf-8", ) - whitelist = Whitelist() lookup = PathLookup.from_sys_path(Path.cwd()) - file_loader = FileLoader() notebook_loader = NotebookLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [ - WhitelistResolver(whitelist), - LocalFileResolver(file_loader), - ] - resolver = DependencyResolver([], notebook_resolver, import_resolvers, lookup) + import_resolver = ImportFileResolver(FileLoader(), Whitelist()) + resolver = DependencyResolver([], notebook_resolver, import_resolver, lookup) maybe = resolver.build_notebook_dependency_graph(parent_file_path) assert not maybe.problems assert maybe.graph is not None diff --git a/tests/unit/source_code/test_s3fs.py b/tests/unit/source_code/test_s3fs.py index bc8a09a02d..93b80918ef 100644 --- a/tests/unit/source_code/test_s3fs.py +++ b/tests/unit/source_code/test_s3fs.py @@ -6,9 +6,9 @@ DependencyResolver, DependencyProblem, ) -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.notebooks.loaders import NotebookLoader, NotebookResolver -from databricks.labs.ucx.source_code.whitelist import WhitelistResolver, Whitelist +from databricks.labs.ucx.source_code.whitelist import Whitelist S3FS_DEPRECATION_MESSAGE = "Use of dependency s3fs is deprecated" @@ -116,8 +116,8 @@ def test_detect_s3fs_import(empty_index, source: str, expected: list[DependencyP notebook_loader = NotebookLoader() file_loader = FileLoader() notebook_resolver = NotebookResolver(notebook_loader) - import_resolvers = [LocalFileResolver(file_loader), WhitelistResolver(whitelist)] - dependency_resolver = DependencyResolver([], notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(file_loader, whitelist) + dependency_resolver = DependencyResolver([], notebook_resolver, import_resolver, mock_path_lookup) maybe = dependency_resolver.build_local_file_dependency_graph(sample) assert maybe.problems == [_.replace(source_path=sample) for _ in expected] @@ -144,8 +144,8 @@ def test_detect_s3fs_import_in_dependencies( yml = mock_path_lookup.cwd / "s3fs-python-compatibility-catalog.yml" file_loader = FileLoader() whitelist = Whitelist.parse(yml.read_text(), use_defaults=True) - import_resolvers = [LocalFileResolver(file_loader), WhitelistResolver(whitelist)] - dependency_resolver = DependencyResolver([], mock_notebook_resolver, import_resolvers, mock_path_lookup) + import_resolver = ImportFileResolver(file_loader, whitelist) + dependency_resolver = DependencyResolver([], mock_notebook_resolver, import_resolver, mock_path_lookup) sample = mock_path_lookup.cwd / "root9.py" maybe = dependency_resolver.build_local_file_dependency_graph(sample) assert maybe.problems == expected