Expand and upload local wheel libraries for all task types (#1649)

## Changes
Fixes #1553 

## Tests
Added regression test

bundle/artifacts/whl/autodetect.go

@@ -27,9 +27,9 @@ func (m *detectPkg) Name() string {
 }
 
 func (m *detectPkg) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
-	wheelTasks := libraries.FindAllWheelTasksWithLocalLibraries(b)
-	if len(wheelTasks) == 0 {
-		log.Infof(ctx, "No local wheel tasks in databricks.yml config, skipping auto detect")
+	tasks := libraries.FindTasksWithLocalLibraries(b)
+	if len(tasks) == 0 {
+		log.Infof(ctx, "No local tasks in databricks.yml config, skipping auto detect")
 		return nil
 	}
 	log.Infof(ctx, "Detecting Python wheel project...")

bundle/artifacts/whl/from_libraries.go

@@ -27,8 +27,13 @@ func (*fromLibraries) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnost
 		return nil
 	}
 
-	tasks := libraries.FindAllWheelTasksWithLocalLibraries(b)
+	tasks := libraries.FindTasksWithLocalLibraries(b)
 	for _, task := range tasks {
+		// Skip tasks that are not PythonWheelTasks for now, we can later support Jars too
+		if task.PythonWheelTask == nil {
+			continue
+		}
+
 		for _, lib := range task.Libraries {
 			matchAndAdd(ctx, lib.Whl, b)
 		}

bundle/libraries/libraries.go

@@ -44,29 +44,25 @@ func isEnvsWithLocalLibraries(envs []jobs.JobEnvironment) bool {
 	return false
 }
 
-func FindAllWheelTasksWithLocalLibraries(b *bundle.Bundle) []*jobs.Task {
+func FindTasksWithLocalLibraries(b *bundle.Bundle) []jobs.Task {
 	tasks := findAllTasks(b)
 	envs := FindAllEnvironments(b)
 
-	wheelTasks := make([]*jobs.Task, 0)
+	allTasks := make([]jobs.Task, 0)
 	for k, jobTasks := range tasks {
 		for i := range jobTasks {
-			task := &jobTasks[i]
-			if task.PythonWheelTask == nil {
-				continue
+			task := jobTasks[i]
+			if isTaskWithLocalLibraries(task) {
+				allTasks = append(allTasks, task)
 			}
+		}
 
-			if isTaskWithLocalLibraries(*task) {
-				wheelTasks = append(wheelTasks, task)
-			}
-
-			if envs[k] != nil && isEnvsWithLocalLibraries(envs[k]) {
-				wheelTasks = append(wheelTasks, task)
-			}
+		if envs[k] != nil && isEnvsWithLocalLibraries(envs[k]) {
+			allTasks = append(allTasks, jobTasks...)
 		}
 	}
 
-	return wheelTasks
+	return allTasks
 }
 
 func isTaskWithLocalLibraries(task jobs.Task) bool {

bundle/python/warning.go

@@ -35,7 +35,7 @@ func isPythonWheelWrapperOn(b *bundle.Bundle) bool {
 }
 
 func hasIncompatibleWheelTasks(ctx context.Context, b *bundle.Bundle) bool {
-	tasks := libraries.FindAllWheelTasksWithLocalLibraries(b)
+	tasks := libraries.FindTasksWithLocalLibraries(b)
 	for _, task := range tasks {
 		if task.NewCluster != nil {
 			if lowerThanExpectedVersion(ctx, task.NewCluster.SparkVersion) {

bundle/tests/python_wheel/python_wheel_no_artifact_notebook/.gitignore

@@ -0,0 +1,3 @@
+build/
+*.egg-info
+.databricks

bundle/tests/python_wheel/python_wheel_no_artifact_notebook/bundle.yml

@@ -0,0 +1,14 @@
+bundle:
+  name: python-wheel-notebook
+
+resources:
+  jobs:
+    test_job:
+      name: "[${bundle.environment}] My Wheel Job"
+      tasks:
+      - task_key: TestTask
+        existing_cluster_id: "0717-aaaaa-bbbbbb"
+        notebook_task:
+          notebook_path: "/notebook.py"
+        libraries:
+        - whl: ./dist/*.whl

bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__init__.py

@@ -0,0 +1,2 @@
+__version__ = "0.0.1"
+__author__ = "Databricks"

bundle/tests/python_wheel/python_wheel_no_artifact_notebook/my_test_code/__main__.py

@@ -0,0 +1,16 @@
+"""
+The entry point of the Python Wheel
+"""
+
+import sys
+
+
+def main():
+    # This method will print the provided arguments
+    print('Hello from my func')
+    print('Got arguments:')
+    print(sys.argv)
+
+
+if __name__ == '__main__':
+    main()

bundle/tests/python_wheel/python_wheel_no_artifact_notebook/notebook.py

@@ -0,0 +1,3 @@
+# Databricks notebook source
+
+print("Hello, World!")

bundle/tests/python_wheel/python_wheel_no_artifact_notebook/setup.py

@@ -0,0 +1,15 @@
+from setuptools import setup, find_packages
+
+import my_test_code
+
+setup(
+    name="my_test_code",
+    version=my_test_code.__version__,
+    author=my_test_code.__author__,
+    url="https://databricks.com",
+    author_email="john.doe@databricks.com",
+    description="my test wheel",
+    packages=find_packages(include=["my_test_code"]),
+    entry_points={"group_1": "run=my_test_code.__main__:main"},
+    install_requires=["setuptools"],
+)

bundle/tests/python_wheel_test.go

@@ -45,6 +45,23 @@ func TestPythonWheelBuildAutoDetect(t *testing.T) {
 	require.NoError(t, diags.Error())
 }
 
+func TestPythonWheelBuildAutoDetectWithNotebookTask(t *testing.T) {
+	ctx := context.Background()
+	b, err := bundle.Load(ctx, "./python_wheel/python_wheel_no_artifact_notebook")
+	require.NoError(t, err)
+
+	diags := bundle.Apply(ctx, b, bundle.Seq(phases.Load(), phases.Build()))
+	require.NoError(t, diags.Error())
+
+	matches, err := filepath.Glob("./python_wheel/python_wheel_no_artifact_notebook/dist/my_test_code-*.whl")
+	require.NoError(t, err)
+	require.Equal(t, 1, len(matches))
+
+	match := libraries.ValidateLocalLibrariesExist()
+	diags = bundle.Apply(ctx, b, match)
+	require.NoError(t, diags.Error())
+}
+
 func TestPythonWheelWithDBFSLib(t *testing.T) {
 	ctx := context.Background()
 	b, err := bundle.Load(ctx, "./python_wheel/python_wheel_dbfs_lib")