From fd6f11ba18a741de0ee5c64be504b0ce38aa46ee Mon Sep 17 00:00:00 2001
From: Denis Bilenko <denis.bilenko@databricks.com>
Date: Tue, 25 Feb 2025 17:46:27 +0100
Subject: [PATCH] Revert "Simplify whl artifact autodetection code (#2371)"

This reverts commit 14f0292598e84009e786a9606ae5d0d4b99dcb6b.
---
 bundle/artifacts/autodetect.go                | 32 ++++++++++++++++
 bundle/artifacts/whl/autodetect.go            | 38 ++++++++++++++-----
 bundle/artifacts/whl/autodetect_test.go       | 22 +++++++++++
 bundle/phases/build.go                        |  3 +-
 .../testdata/default_python/bundle_deploy.txt |  2 +-
 5 files changed, 84 insertions(+), 13 deletions(-)
 create mode 100644 bundle/artifacts/autodetect.go
 create mode 100644 bundle/artifacts/whl/autodetect_test.go

diff --git a/bundle/artifacts/autodetect.go b/bundle/artifacts/autodetect.go
new file mode 100644
index 000000000..c8d235616
--- /dev/null
+++ b/bundle/artifacts/autodetect.go
@@ -0,0 +1,32 @@
+package artifacts
+
+import (
+	"context"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/artifacts/whl"
+	"github.com/databricks/cli/libs/diag"
+	"github.com/databricks/cli/libs/log"
+)
+
+func DetectPackages() bundle.Mutator {
+	return &autodetect{}
+}
+
+type autodetect struct{}
+
+func (m *autodetect) Name() string {
+	return "artifacts.DetectPackages"
+}
+
+func (m *autodetect) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+	// If artifacts section explicitly defined, do not try to auto detect packages
+	if b.Config.Artifacts != nil {
+		log.Debugf(ctx, "artifacts block is defined, skipping auto-detecting")
+		return nil
+	}
+
+	return bundle.Apply(ctx, b, bundle.Seq(
+		whl.DetectPackage(),
+	))
+}
diff --git a/bundle/artifacts/whl/autodetect.go b/bundle/artifacts/whl/autodetect.go
index 9eead83b7..202ea12bc 100644
--- a/bundle/artifacts/whl/autodetect.go
+++ b/bundle/artifacts/whl/autodetect.go
@@ -2,8 +2,11 @@ package whl
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"path/filepath"
+	"regexp"
+	"time"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
@@ -23,17 +26,11 @@ func (m *detectPkg) Name() string {
 }
 
 func (m *detectPkg) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
-	if b.Config.Artifacts != nil {
-		log.Debugf(ctx, "artifacts block is defined, skipping auto-detecting")
-		return nil
-	}
-
 	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...")
 
 	// checking if there is setup.py in the bundle root
@@ -45,18 +42,39 @@ func (m *detectPkg) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostic
 	}
 
 	log.Infof(ctx, "Found Python wheel project at %s", b.BundleRootPath)
+	module := extractModuleName(setupPy)
+
+	if b.Config.Artifacts == nil {
+		b.Config.Artifacts = make(map[string]*config.Artifact)
+	}
 
 	pkgPath, err := filepath.Abs(b.BundleRootPath)
 	if err != nil {
 		return diag.FromErr(err)
 	}
-
-	b.Config.Artifacts = make(map[string]*config.Artifact)
-	b.Config.Artifacts["python_artifact"] = &config.Artifact{
+	b.Config.Artifacts[module] = &config.Artifact{
 		Path: pkgPath,
 		Type: config.ArtifactPythonWheel,
-		// BuildCommand will be set by bundle/artifacts/whl/infer.go to "python3 setup.py bdist_wheel"
 	}
 
 	return nil
 }
+
+func extractModuleName(setupPy string) string {
+	bytes, err := os.ReadFile(setupPy)
+	if err != nil {
+		return randomName()
+	}
+
+	content := string(bytes)
+	r := regexp.MustCompile(`name=['"](.*)['"]`)
+	matches := r.FindStringSubmatch(content)
+	if len(matches) == 0 {
+		return randomName()
+	}
+	return matches[1]
+}
+
+func randomName() string {
+	return fmt.Sprintf("artifact%d", time.Now().Unix())
+}
diff --git a/bundle/artifacts/whl/autodetect_test.go b/bundle/artifacts/whl/autodetect_test.go
new file mode 100644
index 000000000..b53289b2a
--- /dev/null
+++ b/bundle/artifacts/whl/autodetect_test.go
@@ -0,0 +1,22 @@
+package whl
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExtractModuleName(t *testing.T) {
+	moduleName := extractModuleName("./testdata/setup.py")
+	assert.Equal(t, "my_test_code", moduleName)
+}
+
+func TestExtractModuleNameMinimal(t *testing.T) {
+	moduleName := extractModuleName("./testdata/setup_minimal.py")
+	assert.Equal(t, "my_test_code", moduleName)
+}
+
+func TestExtractModuleNameIncorrect(t *testing.T) {
+	moduleName := extractModuleName("./testdata/setup_incorrect.py")
+	assert.Contains(t, moduleName, "artifact")
+}
diff --git a/bundle/phases/build.go b/bundle/phases/build.go
index cc35983ec..3ddc6b181 100644
--- a/bundle/phases/build.go
+++ b/bundle/phases/build.go
@@ -3,7 +3,6 @@ package phases
 import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/artifacts"
-	"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/scripts"
@@ -15,7 +14,7 @@ func Build() bundle.Mutator {
 		"build",
 		[]bundle.Mutator{
 			scripts.Execute(config.ScriptPreBuild),
-			whl.DetectPackage(),
+			artifacts.DetectPackages(),
 			artifacts.InferMissingProperties(),
 			artifacts.PrepareAll(),
 			artifacts.BuildAll(),
diff --git a/integration/bundle/testdata/default_python/bundle_deploy.txt b/integration/bundle/testdata/default_python/bundle_deploy.txt
index fe1cc4fac..076e7618f 100644
--- a/integration/bundle/testdata/default_python/bundle_deploy.txt
+++ b/integration/bundle/testdata/default_python/bundle_deploy.txt
@@ -1,4 +1,4 @@
-Building python_artifact...
+Building project_name_$UNIQUE_PRJ...
 Uploading project_name_$UNIQUE_PRJ-0.0.1+[NUMID].[NUMID]-py3-none-any.whl...
 Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/project_name_$UNIQUE_PRJ/dev/files...
 Deploying resources...