Added support for job environments

bundle/artifacts/artifacts.go

@@ -12,7 +12,6 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/artifacts/whl"
 	"github.com/databricks/cli/bundle/config"
-	"github.com/databricks/cli/bundle/libraries"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/libs/filer"
@@ -117,8 +116,6 @@ func (m *basicUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnost
 }
 
 func uploadArtifact(ctx context.Context, b *bundle.Bundle, a *config.Artifact, uploadPath string, client filer.Filer) error {
-	filesToLibraries := libraries.MapFilesToTaskLibraries(ctx, b)
-
 	for i := range a.Files {
  f := &a.Files[i]
 
@@ -133,31 +130,59 @@ func uploadArtifact(ctx context.Context, b *bundle.Bundle, a *config.Artifact, u
  log.Infof(ctx, "Upload succeeded")
  f.RemotePath = path.Join(uploadPath, filepath.Base(f.Source))
 
- // Lookup all tasks that reference this file.
- libs, ok := filesToLibraries[f.Source]
- if !ok {
- log.Debugf(ctx, "No tasks reference %s", f.Source)
- continue
- }
-
- // Update all tasks that reference this file.
- for _, lib := range libs {
- wsfsBase := "/Workspace"
- remotePath := path.Join(wsfsBase, f.RemotePath)
- if lib.Whl != "" {
- lib.Whl = remotePath
- continue
+ // TODO: confirm if we still need to update the remote path to start with /Workspace
+ wsfsBase := "/Workspace"
+ remotePath := path.Join(wsfsBase, f.RemotePath)
+
+ for _, job := range b.Config.Resources.Jobs {
+ for i := range job.Tasks {
+ task := &job.Tasks[i]
+ for j := range task.Libraries {
+ lib := &task.Libraries[j]
+ if lib.Whl != "" && isArtifactMatchLibrary(f, lib.Whl, b) {
+ lib.Whl = remotePath
+ }
+ if lib.Jar != "" && isArtifactMatchLibrary(f, lib.Jar, b) {
+ lib.Jar = remotePath
+ }
+ }
  }
- if lib.Jar != "" {
- lib.Jar = remotePath
- continue
+
+ for i := range job.Environments {
+ env := &job.Environments[i]
+ for j := range env.Spec.Dependencies {
+ lib := env.Spec.Dependencies[j]
+ if isArtifactMatchLibrary(f, lib, b) {
+ env.Spec.Dependencies[j] = remotePath
+ }
+ }
  }
  }
 	}
 
 	return nil
 }
 
+func isArtifactMatchLibrary(f *config.ArtifactFile, libPath string, b *bundle.Bundle) bool {
+	if !filepath.IsAbs(libPath) {
+ libPath = filepath.Join(b.RootPath, libPath)
+	}
+
+	// libPath can be a glob pattern, so do the match first
+	matches, err := filepath.Glob(libPath)
+	if err != nil {
+ return false
+	}
+
+	for _, m := range matches {
+ if m == f.Source {
+ return true
+ }
+	}
+
+	return false
+}
+
 // Function to upload artifact file to Workspace
 func uploadArtifactFile(ctx context.Context, file string, client filer.Filer) error {
 	raw, err := os.ReadFile(file)

bundle/artifacts/artifacts_test.go

@@ -0,0 +1,91 @@
+package artifacts
+
+import (
+	"context"
+	"path/filepath"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/config/resources"
+	mockfiler "github.com/databricks/cli/internal/mocks/libs/filer"
+	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestArtifactUpload(t *testing.T) {
+	tmpDir := t.TempDir()
+	whlFolder := filepath.Join(tmpDir, "whl")
+	testutil.Touch(t, whlFolder, "source.whl")
+	whlLocalPath := filepath.Join(whlFolder, "source.whl")
+
+	b := &bundle.Bundle{
+ RootPath: tmpDir,
+ Config: config.Root{
+ Workspace: config.Workspace{
+ ArtifactPath: "/foo/bar/artifacts",
+ },
+ Artifacts: config.Artifacts{
+ "whl": {
+ Type: config.ArtifactPythonWheel,
+ Files: []config.ArtifactFile{
+ {Source: whlLocalPath},
+ },
+ },
+ },
+ Resources: config.Resources{
+ Jobs: map[string]*resources.Job{
+ "job": {
+ JobSettings: &jobs.JobSettings{
+ Tasks: []jobs.Task{
+ {
+ Libraries: []compute.Library{
+ {
+ Whl: filepath.Join("whl", "*.whl"),
+ },
+ {
+ Whl: "/Workspace/Users/foo@bar.com/mywheel.whl",
+ },
+ },
+ },
+ },
+ Environments: []jobs.JobEnvironment{
+ {
+ Spec: &compute.Environment{
+ Dependencies: []string{
+ filepath.Join("whl", "source.whl"),
+ "/Workspace/Users/foo@bar.com/mywheel.whl",
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+	}
+
+	artifact := b.Config.Artifacts["whl"]
+	mockFiler := mockfiler.NewMockFiler(t)
+	mockFiler.EXPECT().Write(
+ mock.Anything,
+ filepath.Join("source.whl"),
+ mock.AnythingOfType("*bytes.Reader"),
+ filer.OverwriteIfExists,
+ filer.CreateParentDirectories,
+	).Return(nil)
+
+	err := uploadArtifact(context.Background(), b, artifact, "/foo/bar/artifacts", mockFiler)
+	require.NoError(t, err)
+
+	// Test that libraries path is updated
+	require.Equal(t, "/Workspace/foo/bar/artifacts/source.whl", b.Config.Resources.Jobs["job"].JobSettings.Tasks[0].Libraries[0].Whl)
+	require.Equal(t, "/Workspace/Users/foo@bar.com/mywheel.whl", b.Config.Resources.Jobs["job"].JobSettings.Tasks[0].Libraries[1].Whl)
+	require.Equal(t, "/Workspace/foo/bar/artifacts/source.whl", b.Config.Resources.Jobs["job"].JobSettings.Environments[0].Spec.Dependencies[0])
+	require.Equal(t, "/Workspace/Users/foo@bar.com/mywheel.whl", b.Config.Resources.Jobs["job"].JobSettings.Environments[0].Spec.Dependencies[1])
+}

bundle/artifacts/whl/from_libraries.go

@@ -30,28 +30,45 @@ func (*fromLibraries) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnost
 	tasks := libraries.FindAllWheelTasksWithLocalLibraries(b)
 	for _, task := range tasks {
  for _, lib := range task.Libraries {
- matches, err := filepath.Glob(filepath.Join(b.RootPath, lib.Whl))
- // File referenced from libraries section does not exists, skipping
- if err != nil {
- continue
- }
-
- for _, match := range matches {
- name := filepath.Base(match)
- if b.Config.Artifacts == nil {
- b.Config.Artifacts = make(map[string]*config.Artifact)
- }
+ matchAndAdd(ctx, lib.Whl, b)
+ }
+	}
 
- log.Debugf(ctx, "Adding an artifact block for %s", match)
- b.Config.Artifacts[name] = &config.Artifact{
- Files: []config.ArtifactFile{
- {Source: match},
- },
- Type: config.ArtifactPythonWheel,
+	envs := libraries.FindAllEnvironments(b)
+	for _, jobEnvs := range envs {
+ for _, env := range jobEnvs {
+ if env.Spec != nil {
+ for _, dep := range env.Spec.Dependencies {
+ if libraries.IsEnvironmentDependencyLocal(dep) {
+ matchAndAdd(ctx, dep, b)
+ }
  }
  }
  }
 	}
 
 	return nil
 }
+
+func matchAndAdd(ctx context.Context, lib string, b *bundle.Bundle) {
+	matches, err := filepath.Glob(filepath.Join(b.RootPath, lib))
+	// File referenced from libraries section does not exists, skipping
+	if err != nil {
+ return
+	}
+
+	for _, match := range matches {
+ name := filepath.Base(match)
+ if b.Config.Artifacts == nil {
+ b.Config.Artifacts = make(map[string]*config.Artifact)
+ }
+
+ log.Debugf(ctx, "Adding an artifact block for %s", match)
+ b.Config.Artifacts[name] = &config.Artifact{
+ Files: []config.ArtifactFile{
+ {Source: match},
+ },
+ Type: config.ArtifactPythonWheel,
+ }
+	}
+}

bundle/config/mutator/translate_paths.go

@@ -152,6 +152,13 @@ func translateNoOp(literal, localFullPath, localRelPath, remotePath string) (str
 	return localRelPath, nil
 }
 
+func translateNoOpWithPrefix(literal, localFullPath, localRelPath, remotePath string) (string, error) {
+	if !strings.HasPrefix(localRelPath, ".") {
+ localRelPath = "." + string(filepath.Separator) + localRelPath
+	}
+	return localRelPath, nil
+}
+
 func (m *translatePaths) rewriteValue(b *bundle.Bundle, p dyn.Path, v dyn.Value, fn rewriteFunc, dir string) (dyn.Value, error) {
 	out := v.MustString()
 	err := m.rewritePath(dir, b, &out, fn)

bundle/config/mutator/translate_paths_jobs.go

@@ -5,39 +5,51 @@ import (
 	"slices"
 
 	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/libraries"
 	"github.com/databricks/cli/libs/dyn"
 )
 
-type jobTaskRewritePattern struct {
-	pattern dyn.Pattern
-	fn rewriteFunc
+type jobRewritePattern struct {
+	pattern dyn.Pattern
+	fn rewriteFunc
+	skipRewrite func(string) bool
 }
 
-func rewritePatterns(base dyn.Pattern) []jobTaskRewritePattern {
-	return []jobTaskRewritePattern{
+func noSkipRewrite(string) bool {
+	return false
+}
+
+func rewritePatterns(base dyn.Pattern) []jobRewritePattern {
+	return []jobRewritePattern{
  {
  base.Append(dyn.Key("notebook_task"), dyn.Key("notebook_path")),
  translateNotebookPath,
+ noSkipRewrite,
  },
  {
  base.Append(dyn.Key("spark_python_task"), dyn.Key("python_file")),
  translateFilePath,
+ noSkipRewrite,
  },
  {
  base.Append(dyn.Key("dbt_task"), dyn.Key("project_directory")),
  translateDirectoryPath,
+ noSkipRewrite,
  },
  {
  base.Append(dyn.Key("sql_task"), dyn.Key("file"), dyn.Key("path")),
  translateFilePath,
+ noSkipRewrite,
  },
  {
  base.Append(dyn.Key("libraries"), dyn.AnyIndex(), dyn.Key("whl")),
  translateNoOp,
+ noSkipRewrite,
  },
  {
  base.Append(dyn.Key("libraries"), dyn.AnyIndex(), dyn.Key("jar")),
  translateNoOp,
+ noSkipRewrite,
  },
 	}
 }
@@ -73,9 +85,28 @@ func (m *translatePaths) applyJobTranslations(b *bundle.Bundle, v dyn.Value) (dy
 	)
 
 	// Compile list of patterns and their respective rewrite functions.
+	jobEnvironmentsPatterns := []jobRewritePattern{
+ {
+ dyn.NewPattern(
+ dyn.Key("resources"),
+ dyn.Key("jobs"),
+ dyn.AnyKey(),
+ dyn.Key("environments"),
+ dyn.AnyIndex(),
+ dyn.Key("spec"),
+ dyn.Key("dependencies"),
+ dyn.AnyIndex(),
+ ),
+ translateNoOpWithPrefix,
+ func(s string) bool {
+ return !libraries.IsEnvironmentDependencyLocal(s)
+ },
+ },
+	}
 	taskPatterns := rewritePatterns(base)
 	forEachPatterns := rewritePatterns(base.Append(dyn.Key("for_each_task"), dyn.Key("task")))
-	allPatterns := append(taskPatterns, forEachPatterns...)
+	allPatterns := append(taskPatterns, jobEnvironmentsPatterns...)
+	allPatterns = append(allPatterns, forEachPatterns...)
 
 	for _, t := range allPatterns {
  v, err = dyn.MapByPattern(v, t.pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) {
@@ -91,6 +122,10 @@ func (m *translatePaths) applyJobTranslations(b *bundle.Bundle, v dyn.Value) (dy
  return dyn.InvalidValue, fmt.Errorf("unable to determine directory for job %s: %w", key, err)
  }
 
+ sv := v.MustString()
+ if t.skipRewrite(sv) {
+ return v, nil
+ }
  return m.rewriteRelativeTo(b, p, v, t.fn, dir, fallback[key])
  })
  if err != nil {