Support linting python wheel tasks

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

@@ -3,6 +3,7 @@
 import tempfile
 from collections.abc import Iterable
 from dataclasses import dataclass
+from importlib import metadata
 from pathlib import Path
 
 from databricks.labs.blueprint.parallel import Threads
@@ -139,11 +140,37 @@ def _register_spark_python_task(self, graph: DependencyGraph):
         path = WorkspacePath(self._ws, notebook_path)
         return graph.register_notebook(path)
 
-    def _register_python_wheel_task(self, graph: DependencyGraph):  # pylint: disable=unused-argument
+    @staticmethod
+    def _find_first_matching_distribution(path_lookup: PathLookup, name: str) -> metadata.Distribution | None:
+        # Prepared exists in importlib.metadata.__init__pyi, but is not defined in importlib.metadata.__init__.py
+        normalize_name = metadata.Prepared.normalize  # type: ignore
+        normalized_name = normalize_name(name)
+        for library_root in path_lookup.library_roots:
+            for path in library_root.glob("*.dist-info"):
+                distribution = metadata.Distribution.at(path)
+                if normalize_name(distribution.name) == normalized_name:
+                    return distribution
+        return None
+
+    def _register_python_wheel_task(self, graph: DependencyGraph) -> Iterable[DependencyProblem]:
         if not self._task.python_wheel_task:
-            return
-        # TODO: https://github.com/databrickslabs/ucx/issues/1640
-        yield DependencyProblem('not-yet-implemented', 'Python wheel task is not yet implemented')
+            return []
+
+        distribution_name = self._task.python_wheel_task.package_name
+        distribution = self._find_first_matching_distribution(graph.path_lookup, distribution_name)
+        if distribution is None:
+            return [DependencyProblem("distribution-not-found", f"Could not find distribution for {distribution_name}")]
+        entry_point_name = self._task.python_wheel_task.entry_point
+        try:
+            entry_point = distribution.entry_points[entry_point_name]
+        except KeyError:
+            return [
+                DependencyProblem(
+                    "distribution-entry-point-not-found",
+                    f"Could not find distribution entry point for {distribution_name}.{entry_point_name}",
+                )
+            ]
+        return graph.register_import(entry_point.module)
 
     def _register_spark_jar_task(self, graph: DependencyGraph):  # pylint: disable=unused-argument
         if not self._task.spark_jar_task:

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

@@ -96,7 +96,15 @@ def remove_path(self, index: int):
     @property
     def library_roots(self) -> list[Path]:
         # we may return a combination of WorkspacePath and PosixPath here
-        return [self._cwd] + self._sys_paths
+        library_roots = []
+        for library_root in [self._cwd] + self._sys_paths:
+            try:
+                is_existing_directory = library_root.exists() and library_root.is_dir()
+            except PermissionError:
+                continue
+            if is_existing_directory:
+                library_roots.append(library_root)
+        return library_roots
 
     @property
     def cwd(self):

tests/integration/source_code/test_jobs.py

@@ -341,3 +341,38 @@ def test_job_spark_python_task_linter_unhappy_path(
 
     problems = simple_ctx.workflow_linter.lint_job(j.job_id)
     assert len([problem for problem in problems if problem.message == "Could not locate import: greenlet"]) == 1
+
+
+def test_workflow_linter_lints_python_wheel_task(simple_ctx, ws, make_job, make_random):
+    whitelist = create_autospec(Whitelist)  # databricks is in default list
+    whitelist.module_compatibility.return_value = UNKNOWN
+    whitelist.distribution_compatibility.return_value = UNKNOWN
+
+    simple_ctx = simple_ctx.replace(
+        whitelist=whitelist,
+        path_lookup=PathLookup(Path("/non/existing/path"), []),  # Avoid finding current project
+    )
+
+    simple_ctx.workspace_installation.run()  # Creates ucx wheel
+    wheels = [file for file in simple_ctx.installation.files() if file.path.endswith(".whl")]
+    library = compute.Library(whl=wheels[0].path)
+
+    python_wheel_task = jobs.PythonWheelTask("databricks_labs_ucx", "runtime")
+    task = jobs.Task(
+        task_key=make_random(4),
+        python_wheel_task=python_wheel_task,
+        new_cluster=compute.ClusterSpec(
+            num_workers=1,
+            node_type_id=ws.clusters.select_node_type(local_disk=True, min_memory_gb=16),
+            spark_version=ws.clusters.select_spark_version(latest=True),
+        ),
+        timeout_seconds=0,
+        libraries=[library],
+    )
+    job_with_ucx_library = make_job(tasks=[task])
+
+    problems = simple_ctx.workflow_linter.lint_job(job_with_ucx_library.job_id)
+
+    assert len([problem for problem in problems if problem.code == "library-dist-info-not-found"]) == 0
+    assert len([problem for problem in problems if problem.code == "library-entrypoint-not-found"]) == 0
+    whitelist.distribution_compatibility.assert_called_once_with(Path(wheels[0].path).name)

tests/unit/source_code/test_jobs.py

@@ -257,3 +257,58 @@ def test_workflow_task_container_builds_dependency_graph_with_known_egg_library(
     assert len(problems) == 0
     assert graph.path_lookup.resolve(Path("thingy")) is not None
     ws.workspace.download.assert_called_once_with(egg_file.as_posix(), format=ExportFormat.AUTO)
+
+
+def test_workflow_task_container_builds_dependency_graph_with_missing_distribution_in_python_wheel_task(
+    mock_path_lookup,
+    graph,
+):
+    ws = create_autospec(WorkspaceClient)
+    python_wheel_task = jobs.PythonWheelTask(package_name="databricks_labs_ucx", entry_point="runtime")
+    task = jobs.Task(task_key="test", python_wheel_task=python_wheel_task)
+
+    workflow_task_container = WorkflowTaskContainer(ws, task)
+    problems = workflow_task_container.build_dependency_graph(graph)
+
+    assert len(problems) == 1
+    assert problems[0].code == "distribution-not-found"
+    assert problems[0].message == "Could not find distribution for databricks_labs_ucx"
+    ws.assert_not_called()
+
+
+def test_workflow_task_container_builds_dependency_graph_with_missing_entrypoint_in_python_wheel_task(graph):
+    ws = create_autospec(WorkspaceClient)
+
+    whl_file = Path(__file__).parent / "samples/distribution/dist/thingy-0.0.1-py2.py3-none-any.whl"
+    with whl_file.open("rb") as f:
+        ws.workspace.download.return_value = io.BytesIO(f.read())
+
+    python_wheel_task = jobs.PythonWheelTask(package_name="thingy", entry_point="non_existing_entrypoint")
+    libraries = [compute.Library(whl=whl_file.as_posix())]
+    task = jobs.Task(task_key="test", libraries=libraries, python_wheel_task=python_wheel_task)
+
+    workflow_task_container = WorkflowTaskContainer(ws, task)
+    problems = workflow_task_container.build_dependency_graph(graph)
+
+    assert len(problems) == 1
+    assert problems[0].code == "distribution-entry-point-not-found"
+    assert problems[0].message == "Could not find distribution entry point for thingy.non_existing_entrypoint"
+    ws.workspace.download.assert_called_once_with(whl_file.as_posix(), format=ExportFormat.AUTO)
+
+
+def test_workflow_task_container_builds_dependency_graph_for_python_wheel_task(graph):
+    ws = create_autospec(WorkspaceClient)
+
+    whl_file = Path(__file__).parent / "samples/distribution/dist/thingy-0.0.1-py2.py3-none-any.whl"
+    with whl_file.open("rb") as f:
+        ws.workspace.download.return_value = io.BytesIO(f.read())
+
+    python_wheel_task = jobs.PythonWheelTask(package_name="thingy", entry_point="runtime")
+    libraries = [compute.Library(whl=whl_file.as_posix())]
+    task = jobs.Task(task_key="test", libraries=libraries, python_wheel_task=python_wheel_task)
+
+    workflow_task_container = WorkflowTaskContainer(ws, task)
+    problems = workflow_task_container.build_dependency_graph(graph)
+
+    assert len(problems) == 0
+    ws.workspace.download.assert_called_once_with(whl_file.as_posix(), format=ExportFormat.AUTO)

tests/unit/source_code/test_path_lookup.py

@@ -11,22 +11,43 @@ def test_lookup_is_initialized_with_syspath():
     assert len(filtered) > 0
 
 
-def test_lookup_is_initialized_with_handmade_string():
-    provider = PathLookup.from_pathlike_string(Path.cwd(), "what:on:earth")
+def test_lookup_is_initialized_with_handmade_string(tmp_path):
+    directories, sys_paths = ("what", "on", "earth"), []
+    for directory in directories:
+        path = tmp_path / directory
+        path.mkdir()
+        sys_paths.append(path)
+
+    provider = PathLookup.from_pathlike_string(Path.cwd(), ":".join([p.as_posix() for p in sys_paths]))
+
     assert provider is not None
-    paths = provider.library_roots[1:]
-    assert paths == [Path("what"), Path("on"), Path("earth")]
+    assert provider.library_roots[1:] == sys_paths
+
+
+def test_lookup_inserts_path(tmp_path):
+    directories, sys_paths = ("what", "on", "earth"), []
+    for directory in directories:
+        path = tmp_path / directory
+        path.mkdir()
+        sys_paths.append(path)
+
+    provider = PathLookup.from_pathlike_string(Path.cwd(), ":".join([p.as_posix() for p in sys_paths]))
+
+    new_sys_path = tmp_path / Path("is")
+    new_sys_path.mkdir()
+    provider.insert_path(1, new_sys_path)
 
+    assert provider.library_roots[1:] == [sys_paths[0]] + [new_sys_path] + sys_paths[1:]
 
-def test_lookup_inserts_path():
-    provider = PathLookup.from_pathlike_string(Path.cwd(), "what:on:earth")
-    provider.insert_path(1, Path("is"))
-    paths = provider.library_roots[1:]
-    assert paths == [Path("what"), Path("is"), Path("on"), Path("earth")]
 
+def test_lookup_removes_path(tmp_path):
+    directories, sys_paths = ("what", "is", "on", "earth"), []
+    for directory in directories:
+        path = tmp_path / directory
+        path.mkdir()
+        sys_paths.append(path)
 
-def test_lookup_removes_path():
-    provider = PathLookup.from_pathlike_string(Path.cwd(), "what:is:on:earth")
+    provider = PathLookup.from_pathlike_string(Path.cwd(), ":".join([p.as_posix() for p in sys_paths]))
     provider.remove_path(1)
-    paths = provider.library_roots[1:]
-    assert paths == [Path("what"), Path("on"), Path("earth")]
+    sys_paths.pop(1)
+    assert provider.library_roots[1:] == sys_paths