Split artifact cleanup into prepare step before build (#1618)

## Changes
Now prepare stage which does cleanup is execute once before every build,
so artifacts built into the same folder are correctly kept

Fixes workaround 2 from this issue #1602

## Tests
Added unit test

bundle/artifacts/artifacts.go

@@ -13,6 +13,7 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/artifacts/whl"
 	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/config/mutator"
 	"github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/diag"
@@ -29,6 +30,10 @@ var buildMutators map[config.ArtifactType]mutatorFactory = map[config.ArtifactTy
 
 var uploadMutators map[config.ArtifactType]mutatorFactory = map[config.ArtifactType]mutatorFactory{}
 
+var prepareMutators map[config.ArtifactType]mutatorFactory = map[config.ArtifactType]mutatorFactory{
+	config.ArtifactPythonWheel: whl.Prepare,
+}
+
 func getBuildMutator(t config.ArtifactType, name string) bundle.Mutator {
 	mutatorFactory, ok := buildMutators[t]
 	if !ok {
@@ -47,6 +52,17 @@ func getUploadMutator(t config.ArtifactType, name string) bundle.Mutator {
 	return mutatorFactory(name)
 }
 
+func getPrepareMutator(t config.ArtifactType, name string) bundle.Mutator {
+	mutatorFactory, ok := prepareMutators[t]
+	if !ok {
+		mutatorFactory = func(_ string) bundle.Mutator {
+			return mutator.NoOp()
+		}
+	}
+
+	return mutatorFactory(name)
+}
+
 // Basic Build defines a general build mutator which builds artifact based on artifact.BuildCommand
 type basicBuild struct {
 	name string

bundle/artifacts/build.go

@@ -35,15 +35,6 @@ func (m *build) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 		return diag.Errorf("artifact doesn't exist: %s", m.name)
 	}
 
-	// Check if source paths are absolute, if not, make them absolute
-	for k := range artifact.Files {
-		f := &artifact.Files[k]
-		if !filepath.IsAbs(f.Source) {
-			dirPath := filepath.Dir(artifact.ConfigFilePath)
-			f.Source = filepath.Join(dirPath, f.Source)
-		}
-	}
-
 	// Skip building if build command is not specified or infered
 	if artifact.BuildCommand == "" {
 		// If no build command was specified or infered and there is no
@@ -58,16 +49,6 @@ func (m *build) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 		return diags
 	}
 
-	// If artifact path is not provided, use bundle root dir
-	if artifact.Path == "" {
-		artifact.Path = b.RootPath
-	}
-
-	if !filepath.IsAbs(artifact.Path) {
-		dirPath := filepath.Dir(artifact.ConfigFilePath)
-		artifact.Path = filepath.Join(dirPath, artifact.Path)
-	}
-
 	diags := bundle.Apply(ctx, b, getBuildMutator(artifact.Type, m.name))
 	if diags.HasError() {
 		return diags

bundle/artifacts/prepare.go

@@ -0,0 +1,57 @@
+package artifacts
+
+import (
+	"context"
+	"fmt"
+	"path/filepath"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/libs/diag"
+)
+
+func PrepareAll() bundle.Mutator {
+	return &all{
+		name: "Prepare",
+		fn:   prepareArtifactByName,
+	}
+}
+
+type prepare struct {
+	name string
+}
+
+func prepareArtifactByName(name string) (bundle.Mutator, error) {
+	return &prepare{name}, nil
+}
+
+func (m *prepare) Name() string {
+	return fmt.Sprintf("artifacts.Prepare(%s)", m.name)
+}
+
+func (m *prepare) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+	artifact, ok := b.Config.Artifacts[m.name]
+	if !ok {
+		return diag.Errorf("artifact doesn't exist: %s", m.name)
+	}
+
+	// Check if source paths are absolute, if not, make them absolute
+	for k := range artifact.Files {
+		f := &artifact.Files[k]
+		if !filepath.IsAbs(f.Source) {
+			dirPath := filepath.Dir(artifact.ConfigFilePath)
+			f.Source = filepath.Join(dirPath, f.Source)
+		}
+	}
+
+	// If artifact path is not provided, use bundle root dir
+	if artifact.Path == "" {
+		artifact.Path = b.RootPath
+	}
+
+	if !filepath.IsAbs(artifact.Path) {
+		dirPath := filepath.Dir(artifact.ConfigFilePath)
+		artifact.Path = filepath.Join(dirPath, artifact.Path)
+	}
+
+	return bundle.Apply(ctx, b, getPrepareMutator(artifact.Type, m.name))
+}

bundle/artifacts/upload_test.go

@@ -63,7 +63,12 @@ func TestExpandGlobFilesSource(t *testing.T) {
 		return &noop{}
 	}
 
-	diags := bundle.Apply(context.Background(), b, bundle.Seq(bm, u))
+	pm := &prepare{"test"}
+	prepareMutators[config.ArtifactType("custom")] = func(name string) bundle.Mutator {
+		return &noop{}
+	}
+
+	diags := bundle.Apply(context.Background(), b, bundle.Seq(pm, bm, u))
 	require.NoError(t, diags.Error())
 
 	require.Equal(t, 2, len(b.Config.Artifacts["test"].Files))

bundle/artifacts/whl/build.go

@@ -3,7 +3,6 @@ package whl
 import (
 	"context"
 	"fmt"
-	"os"
 	"path/filepath"
 
 	"github.com/databricks/cli/bundle"
@@ -36,18 +35,14 @@ func (m *build) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 
 	cmdio.LogString(ctx, fmt.Sprintf("Building %s...", m.name))
 
-	dir := artifact.Path
-
-	distPath := filepath.Join(dir, "dist")
-	os.RemoveAll(distPath)
-	python.CleanupWheelFolder(dir)
-
 	out, err := artifact.Build(ctx)
 	if err != nil {
 		return diag.Errorf("build failed %s, error: %v, output: %s", m.name, err, out)
 	}
 	log.Infof(ctx, "Build succeeded")
 
+	dir := artifact.Path
+	distPath := filepath.Join(artifact.Path, "dist")
 	wheels := python.FindFilesWithSuffixInPath(distPath, ".whl")
 	if len(wheels) == 0 {
 		return diag.Errorf("cannot find built wheel in %s for package %s", dir, m.name)

bundle/artifacts/whl/prepare.go

@@ -0,0 +1,48 @@
+package whl
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/libs/diag"
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/python"
+)
+
+type prepare struct {
+	name string
+}
+
+func Prepare(name string) bundle.Mutator {
+	return &prepare{
+		name: name,
+	}
+}
+
+func (m *prepare) Name() string {
+	return fmt.Sprintf("artifacts.whl.Prepare(%s)", m.name)
+}
+
+func (m *prepare) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+	artifact, ok := b.Config.Artifacts[m.name]
+	if !ok {
+		return diag.Errorf("artifact doesn't exist: %s", m.name)
+	}
+
+	dir := artifact.Path
+
+	distPath := filepath.Join(dir, "dist")
+
+	// If we have multiple artifacts con figured, prepare will be called multiple times
+	// The first time we will remove the folders, other times will be no-op.
+	err := os.RemoveAll(distPath)
+	if err != nil {
+		log.Infof(ctx, "Failed to remove dist folder: %v", err)
+	}
+	python.CleanupWheelFolder(dir)
+
+	return nil
+}

bundle/phases/build.go

@@ -16,6 +16,7 @@ func Build() bundle.Mutator {
 			scripts.Execute(config.ScriptPreBuild),
 			artifacts.DetectPackages(),
 			artifacts.InferMissingProperties(),
+			artifacts.PrepareAll(),
 			artifacts.BuildAll(),
 			scripts.Execute(config.ScriptPostBuild),
 			mutator.ResolveVariableReferences(

bundle/tests/python_wheel/python_wheel_multiple/.gitignore

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

bundle/tests/python_wheel/python_wheel_multiple/bundle.yml

@@ -0,0 +1,25 @@
+bundle:
+  name: python-wheel
+
+artifacts:
+    my_test_code:
+      type: whl
+      path: "./my_test_code"
+      build: "python3 setup.py bdist_wheel"
+    my_test_code_2:
+      type: whl
+      path: "./my_test_code"
+      build: "python3 setup2.py bdist_wheel"
+
+resources:
+  jobs:
+    test_job:
+      name: "[${bundle.environment}] My Wheel Job"
+      tasks:
+        - task_key: TestTask
+          existing_cluster_id: "0717-132531-5opeqon1"
+          python_wheel_task:
+            package_name: "my_test_code"
+            entry_point: "run"
+          libraries:
+          - whl: ./my_test_code/dist/*.whl

bundle/tests/python_wheel/python_wheel_multiple/my_test_code/setup.py

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

bundle/tests/python_wheel/python_wheel_multiple/my_test_code/setup2.py

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

bundle/tests/python_wheel/python_wheel_multiple/my_test_code/src/__init__.py

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

bundle/tests/python_wheel/python_wheel_multiple/my_test_code/src/__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_test.go

@@ -96,3 +96,20 @@ func TestPythonWheelBuildWithEnvironmentKey(t *testing.T) {
 	diags = bundle.Apply(ctx, b, match)
 	require.NoError(t, diags.Error())
 }
+
+func TestPythonWheelBuildMultiple(t *testing.T) {
+	ctx := context.Background()
+	b, err := bundle.Load(ctx, "./python_wheel/python_wheel_multiple")
+	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_multiple/my_test_code/dist/my_test_code*.whl")
+	require.NoError(t, err)
+	require.Equal(t, 2, len(matches))
+
+	match := libraries.ValidateLocalLibrariesExist()
+	diags = bundle.Apply(ctx, b, match)
+	require.NoError(t, diags.Error())
+}