Expand and upload local wheel libraries for all task types

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
- }
-
- if isTaskWithLocalLibraries(*task) {
- wheelTasks = append(wheelTasks, task)
+ task := jobTasks[i]
+ if isTaskWithLocalLibraries(task) {
+ allTasks = append(allTasks, 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")