From 463fbb6e7c12d946b3abc6497e72afb6da7efc5d Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Mon, 13 May 2024 10:30:07 +0200
Subject: [PATCH 01/10] Add ApplyPythonMutator

---
 bundle/config/experimental.go                 |  16 +
 bundle/config/mutator/mutator.go              |   2 +
 .../mutator/python/apply_python_mutator.go    | 190 ++++++++
 .../python/apply_python_mutator_test.go       | 460 ++++++++++++++++++
 bundle/config/mutator/python/python_logger.go |  95 ++++
 .../mutator/python/python_logger_test.go      |  65 +++
 bundle/config/root.go                         |   4 +
 bundle/phases/initialize.go                   |   2 +
 libs/process/opts.go                          |  14 +
 9 files changed, 848 insertions(+)
 create mode 100644 bundle/config/mutator/python/apply_python_mutator.go
 create mode 100644 bundle/config/mutator/python/apply_python_mutator_test.go
 create mode 100644 bundle/config/mutator/python/python_logger.go
 create mode 100644 bundle/config/mutator/python/python_logger_test.go

diff --git a/bundle/config/experimental.go b/bundle/config/experimental.go
index 008d7b909f..0c1bcac6ea 100644
--- a/bundle/config/experimental.go
+++ b/bundle/config/experimental.go
@@ -23,6 +23,22 @@ type Experimental struct {
 	// be removed in the future once we have a proper workaround like allowing IS_OWNER
 	// as a top-level permission in the DAB.
 	UseLegacyRunAs bool `json:"use_legacy_run_as,omitempty"`
+
+	// EnablePyDABs if enabled loads the 'databricks-pydabs' package.
+	//
+	// This allows to define bundle configuration using Python. For now, this option has to be used
+	// together with 'venv' configuration.
+	EnablePyDABs bool `json:"enable_pydabs,omitempty"`
+
+	// VEnv is the virtual environment configuration for the bundle.
+	//
+	// When enabled, the bundle will use Python interpreter from the virtual environment.
+	VEnv VEnv `json:"venv,omitempty"`
+}
+
+type VEnv struct {
+	// Path to the virtual environment
+	Path string `json:"path,omitempty"`
 }
 
 type Command string
diff --git a/bundle/config/mutator/mutator.go b/bundle/config/mutator/mutator.go
index ae0d7e5fbd..1e188d1a66 100644
--- a/bundle/config/mutator/mutator.go
+++ b/bundle/config/mutator/mutator.go
@@ -4,6 +4,7 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/loader"
+	pythonmutator "github.com/databricks/cli/bundle/config/mutator/python"
 	"github.com/databricks/cli/bundle/scripts"
 )
 
@@ -24,6 +25,7 @@ func DefaultMutators() []bundle.Mutator {
 		InitializeVariables(),
 		DefineDefaultTarget(),
 		LoadGitDetails(),
+		pythonmutator.ApplyPythonMutator(pythonmutator.ApplyPythonMutatorPhasePreInit),
 	}
 }
 
diff --git a/bundle/config/mutator/python/apply_python_mutator.go b/bundle/config/mutator/python/apply_python_mutator.go
new file mode 100644
index 0000000000..ffa34d4313
--- /dev/null
+++ b/bundle/config/mutator/python/apply_python_mutator.go
@@ -0,0 +1,190 @@
+package python
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/libs/diag"
+	"github.com/databricks/cli/libs/dyn"
+	"github.com/databricks/cli/libs/dyn/convert"
+	"github.com/databricks/cli/libs/dyn/merge"
+	"github.com/databricks/cli/libs/dyn/yamlloader"
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/process"
+	"os"
+	"path"
+	"path/filepath"
+)
+
+type phase string
+
+const (
+	ApplyPythonMutatorPhasePreInit phase = "preinit"
+	ApplyPythonMutatorPhaseInit    phase = "init"
+)
+
+type applyPythonMutator struct {
+	phase phase
+}
+
+func ApplyPythonMutator(phase phase) bundle.Mutator {
+	return &applyPythonMutator{
+		phase: phase,
+	}
+}
+
+func (m *applyPythonMutator) Name() string {
+	return fmt.Sprintf("ApplyPythonMutator(%s)", m.phase)
+}
+
+func getExperimental(b *bundle.Bundle) config.Experimental {
+	if b.Config.Experimental == nil {
+		return config.Experimental{}
+	}
+
+	return *b.Config.Experimental
+}
+
+func (m *applyPythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+	experimental := getExperimental(b)
+
+	if !experimental.EnablePyDABs {
+		log.Debugf(ctx, "'experimental.enable_pydabs' isn't enabled, skipping")
+		return nil
+	}
+
+	if experimental.EnablePyDABs && experimental.VEnv.Path == "" {
+		return diag.Errorf("'experimental.enable_pydabs' can only be used when 'experimental.venv.path' is set")
+	}
+
+	err := b.Config.Mutate(func(leftRoot dyn.Value) (dyn.Value, error) {
+		pythonPath := path.Join(experimental.VEnv.Path, "bin", "python3")
+
+		if _, err := os.Stat(pythonPath); err != nil {
+			if os.IsNotExist(err) {
+				return dyn.InvalidValue, fmt.Errorf("can't find '%s', check if venv is created", pythonPath)
+			} else {
+				return dyn.InvalidValue, fmt.Errorf("can't find '%s': %w", pythonPath, err)
+			}
+		}
+
+		rightRoot, err := m.runPythonMutator(ctx, b.RootPath, pythonPath, leftRoot)
+
+		if err != nil {
+			return dyn.InvalidValue, err
+		}
+
+		return merge.Override(leftRoot, rightRoot, createOverrideVisitor(ctx, m.phase))
+	})
+
+	return diag.FromErr(err)
+}
+
+func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath string, pythonPath string, root dyn.Value) (dyn.Value, error) {
+	args := []string{
+		pythonPath,
+		"-m",
+		"databricks.bundles.build",
+		"--phase",
+		string(m.phase),
+	}
+
+	// we need to marshal dyn.Value instead of bundle.Config to JSON to support
+	// non-string fields assigned with bundle variables
+	rootConfigJson, err := json.Marshal(root.AsAny())
+
+	if err != nil {
+		return dyn.InvalidValue, fmt.Errorf("failed to marshal root config: %w", err)
+	}
+
+	pythonLogger := NewStderrLogger(ctx)
+
+	stdout, err := process.Background(
+		ctx,
+		args,
+		process.WithDir(rootPath),
+		process.WithStderrWriter(pythonLogger),
+		process.WithStdinReader(bytes.NewBuffer(rootConfigJson)),
+	)
+
+	if err != nil {
+		return dyn.InvalidValue, fmt.Errorf("python mutator process failed: %w", err)
+	}
+
+	// we need absolute path, or because later parts of pipeline assume all paths are absolute
+	// and this file will be used as location
+	virtualPath, err := filepath.Abs(filepath.Join(rootPath, "__generated__.yml"))
+
+	if err != nil {
+		return dyn.InvalidValue, fmt.Errorf("failed to get absolute path: %w", err)
+	}
+
+	generated, err := yamlloader.LoadYAML(virtualPath, bytes.NewReader([]byte(stdout)))
+
+	if err != nil {
+		return dyn.InvalidValue, fmt.Errorf("failed to parse Python mutator output: %w", err)
+	}
+
+	normalized, diagnostic := convert.Normalize(config.Root{}, generated)
+
+	if diagnostic.Error() != nil {
+		return dyn.InvalidValue, fmt.Errorf("failed to normalize Python mutator output: %w", diagnostic.Error())
+	}
+
+	return normalized, nil
+}
+
+func createOverrideVisitor(ctx context.Context, phase phase) merge.OverrideVisitor {
+	jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs"))
+
+	return merge.OverrideVisitor{
+		VisitDelete: func(valuePath dyn.Path, left dyn.Value) error {
+			if !valuePath.HasPrefix(jobsPath) {
+				return fmt.Errorf("unexpected change at '%s' (delete)", valuePath.String())
+			}
+
+			deleteResource := len(valuePath) == len(jobsPath)+1
+
+			if phase == ApplyPythonMutatorPhasePreInit {
+				return fmt.Errorf("unexpected change at '%s' (delete)", valuePath.String())
+			} else if deleteResource && phase == ApplyPythonMutatorPhaseInit {
+				return fmt.Errorf("unexpected change at '%s' (delete)", valuePath.String())
+			}
+
+			log.Debugf(ctx, "Delete value at '%s'", valuePath.String())
+
+			return nil
+		},
+		VisitInsert: func(valuePath dyn.Path, right dyn.Value) (dyn.Value, error) {
+			if !valuePath.HasPrefix(jobsPath) {
+				return dyn.InvalidValue, fmt.Errorf("unexpected change at '%s' (insert)", valuePath.String())
+			}
+
+			insertResource := len(valuePath) == len(jobsPath)+1
+
+			if phase == ApplyPythonMutatorPhasePreInit && !insertResource {
+				return dyn.InvalidValue, fmt.Errorf("unexpected change at '%s' (insert)", valuePath.String())
+			}
+
+			log.Debugf(ctx, "Insert value at '%s'", valuePath.String())
+
+			return right, nil
+		},
+		VisitUpdate: func(valuePath dyn.Path, left dyn.Value, right dyn.Value) (dyn.Value, error) {
+			if !valuePath.HasPrefix(jobsPath) {
+				return dyn.InvalidValue, fmt.Errorf("unexpected change at '%s' (update)", valuePath.String())
+			}
+
+			if phase == ApplyPythonMutatorPhasePreInit {
+				return dyn.InvalidValue, fmt.Errorf("unexpected change at '%s' (update)", valuePath.String())
+			}
+
+			log.Debugf(ctx, "Update value at '%s'", valuePath.String())
+
+			return right, nil
+		},
+	}
+}
diff --git a/bundle/config/mutator/python/apply_python_mutator_test.go b/bundle/config/mutator/python/apply_python_mutator_test.go
new file mode 100644
index 0000000000..f892ce513f
--- /dev/null
+++ b/bundle/config/mutator/python/apply_python_mutator_test.go
@@ -0,0 +1,460 @@
+package python
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"github.com/databricks/cli/libs/dyn"
+	"log/slog"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"reflect"
+	"testing"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/config"
+	assert "github.com/databricks/cli/libs/dyn/dynassert"
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/process"
+)
+
+func TestApplyPythonMutator_Name_preinit(t *testing.T) {
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+
+	assert.Equal(t, "ApplyPythonMutator(preinit)", mutator.Name())
+}
+
+func TestApplyPythonMutator_Name_init(t *testing.T) {
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseInit)
+
+	assert.Equal(t, "ApplyPythonMutator(init)", mutator.Name())
+}
+
+func TestApplyPythonMutator_preinit(t *testing.T) {
+	withFakeVEnv(t, ".venv")
+
+	b := loadYaml("bundle.yaml", `
+      experimental:
+        enable_pydabs: true
+        venv:
+          path: .venv
+      resources:
+        jobs:
+          job0:
+            name: job_0`)
+
+	ctx := withProcessStub(
+		[]string{
+			".venv/bin/python3",
+			"-m",
+			"databricks.bundles.build",
+			"--phase",
+			"preinit",
+		},
+		`{
+			"experimental": {
+				"enable_pydabs": true,
+				"venv": { "path": ".venv" }
+			},
+			"resources": {
+				"jobs": {
+					"job0": {
+						name: "job_0"
+					},
+					"job1": {
+						name: "job_1"
+					},
+				}
+			}
+		}`,
+		`{"level": "INFO", "message": "Applying Python mutator"}`,
+	)
+
+	loggerBuf, ctx := withLoggerStub(ctx)
+
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	diag := bundle.Apply(ctx, b, mutator)
+
+	assert.NoError(t, diag.Error())
+
+	assert.ElementsMatch(t, []string{"job0", "job1"}, keys(b.Config.Resources.Jobs))
+
+	if job0, ok := b.Config.Resources.Jobs["job0"]; ok {
+		assert.Equal(t, "job_0", job0.Name)
+	}
+
+	if job1, ok := b.Config.Resources.Jobs["job1"]; ok {
+		assert.Equal(t, "job_1", job1.Name)
+	}
+
+	assert.Equal(t, "level=INFO msg=\"Applying Python mutator\"\n", loggerBuf.String())
+}
+
+func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
+	withFakeVEnv(t, ".venv")
+
+	b := loadYaml("bundle.yaml", `
+      experimental:
+        enable_pydabs: true
+        venv:
+          path: .venv
+      resources:
+        jobs:
+          job0:
+            name: job_0`)
+
+	ctx := withProcessStub(
+		[]string{
+			".venv/bin/python3",
+			"-m",
+			"databricks.bundles.build",
+			"--phase",
+			"preinit",
+		},
+		`{
+			"experimental": {
+				"enable_pydabs": true,
+				"venv": { "path": ".venv" }
+			},
+			"resources": {
+				"jobs": {
+					"job0": {
+						name: "job_0",
+						description: "job description"
+					}
+				}
+			}
+		}`,
+		"",
+	)
+
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	diag := bundle.Apply(ctx, b, mutator)
+
+	assert.EqualError(t, diag.Error(), "unexpected change at 'resources.jobs.job0.description' (insert)")
+}
+
+func TestApplyPythonMutator_init(t *testing.T) {
+	withFakeVEnv(t, ".venv")
+
+	b := loadYaml("bundle.yaml", `
+      experimental:
+        enable_pydabs: true
+        venv:
+          path: .venv
+      resources:
+        jobs:
+          job0:
+            name: job_0`)
+
+	ctx := withProcessStub(
+		[]string{
+			".venv/bin/python3",
+			"-m",
+			"databricks.bundles.build",
+			"--phase",
+			"init",
+		},
+		`{
+			"experimental": {
+				"enable_pydabs": true,
+				"venv": { "path": ".venv" }
+			},
+			"resources": {
+				"jobs": {
+					"job0": {
+						name: "job_0",
+						description: "my job"
+					}
+				}
+			}
+		}`,
+		`{"level": "INFO", "message": "Applying Python mutator"}`,
+	)
+	loggerBuf, ctx := withLoggerStub(ctx)
+
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseInit)
+	diag := bundle.Apply(ctx, b, mutator)
+
+	assert.NoError(t, diag.Error())
+
+	assert.ElementsMatch(t, []string{"job0"}, keys(b.Config.Resources.Jobs))
+	assert.Equal(t, "job_0", b.Config.Resources.Jobs["job0"].Name)
+	assert.Equal(t, "my job", b.Config.Resources.Jobs["job0"].Description)
+
+	assert.Equal(t, "level=INFO msg=\"Applying Python mutator\"\n", loggerBuf.String())
+}
+
+func TestApplyPythonMutator_disabled(t *testing.T) {
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Experimental: &config.Experimental{
+				EnablePyDABs: false,
+				VEnv: config.VEnv{
+					Path: ".venv",
+				},
+			},
+		},
+	}
+
+	ctx := context.Background()
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	diag := bundle.Apply(ctx, b, mutator)
+
+	assert.NoError(t, diag.Error())
+}
+
+func TestApplyPythonMutator_venvRequired(t *testing.T) {
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Experimental: &config.Experimental{
+				EnablePyDABs: true,
+				VEnv:         config.VEnv{},
+			},
+		},
+	}
+
+	ctx := context.Background()
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	diag := bundle.Apply(ctx, b, mutator)
+
+	assert.Error(t, diag.Error(), "'experimental.enable_pydabs' is enabled, but 'experimental.venv.path' is not set")
+}
+
+func TestApplyPythonMutator_venvNotFound(t *testing.T) {
+	withFakeVEnv(t, ".venv")
+
+	b := loadYaml("bundle.yaml", `
+      experimental:
+        enable_pydabs: true
+        venv:
+          path: bad_path`)
+
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseInit)
+	diag := bundle.Apply(context.Background(), b, mutator)
+
+	assert.EqualError(t, diag.Error(), "can't find 'bad_path/bin/python3', check if venv is created")
+}
+
+type createOverrideVisitorTestCase struct {
+	name        string
+	updatePath  dyn.Path
+	deletePath  dyn.Path
+	insertPath  dyn.Path
+	phase       phase
+	updateError error
+	deleteError error
+	insertError error
+}
+
+func TestCreateOverrideVisitor(t *testing.T) {
+	left := dyn.NewValue(42, dyn.Location{})
+	right := dyn.NewValue(1337, dyn.Location{})
+
+	testCases := []createOverrideVisitorTestCase{
+		{
+			name:        "preinit: can't change an existing job",
+			phase:       ApplyPythonMutatorPhasePreInit,
+			updatePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
+			deletePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
+			insertPath:  dyn.MustPathFromString("resources.jobs.job0.name"),
+			deleteError: fmt.Errorf("unexpected change at 'resources.jobs.job0.name' (delete)"),
+			insertError: fmt.Errorf("unexpected change at 'resources.jobs.job0.name' (insert)"),
+			updateError: fmt.Errorf("unexpected change at 'resources.jobs.job0.name' (update)"),
+		},
+		{
+			name:        "preinit: can't delete an existing job",
+			phase:       ApplyPythonMutatorPhasePreInit,
+			deletePath:  dyn.MustPathFromString("resources.jobs.job0"),
+			deleteError: fmt.Errorf("unexpected change at 'resources.jobs.job0' (delete)"),
+		},
+		{
+			name:        "preinit: can insert a job",
+			phase:       ApplyPythonMutatorPhasePreInit,
+			insertPath:  dyn.MustPathFromString("resources.jobs.job0"),
+			insertError: nil,
+		},
+		{
+			name:        "preinit: can't change include",
+			phase:       ApplyPythonMutatorPhasePreInit,
+			deletePath:  dyn.MustPathFromString("include[0]"),
+			insertPath:  dyn.MustPathFromString("include[0]"),
+			updatePath:  dyn.MustPathFromString("include[0]"),
+			deleteError: fmt.Errorf("unexpected change at 'include[0]' (delete)"),
+			insertError: fmt.Errorf("unexpected change at 'include[0]' (insert)"),
+			updateError: fmt.Errorf("unexpected change at 'include[0]' (update)"),
+		},
+		{
+			name:        "preinit: can change an existing job",
+			phase:       ApplyPythonMutatorPhaseInit,
+			updatePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
+			deletePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
+			insertPath:  dyn.MustPathFromString("resources.jobs.job0.name"),
+			deleteError: nil,
+			insertError: nil,
+			updateError: nil,
+		},
+		{
+			name:        "init: can't delete an existing job",
+			phase:       ApplyPythonMutatorPhaseInit,
+			deletePath:  dyn.MustPathFromString("resources.jobs.job0"),
+			deleteError: fmt.Errorf("unexpected change at 'resources.jobs.job0' (delete)"),
+		},
+		{
+			name:        "init: can insert a job",
+			phase:       ApplyPythonMutatorPhaseInit,
+			insertPath:  dyn.MustPathFromString("resources.jobs.job0"),
+			insertError: nil,
+		},
+		{
+			name:        "init: can't change include",
+			phase:       ApplyPythonMutatorPhaseInit,
+			deletePath:  dyn.MustPathFromString("include[0]"),
+			insertPath:  dyn.MustPathFromString("include[0]"),
+			updatePath:  dyn.MustPathFromString("include[0]"),
+			deleteError: fmt.Errorf("unexpected change at 'include[0]' (delete)"),
+			insertError: fmt.Errorf("unexpected change at 'include[0]' (insert)"),
+			updateError: fmt.Errorf("unexpected change at 'include[0]' (update)"),
+		},
+	}
+
+	for _, tc := range testCases {
+		visitor := createOverrideVisitor(context.Background(), tc.phase)
+
+		if tc.updatePath != nil {
+			t.Run(tc.name+"-update", func(t *testing.T) {
+				out, err := visitor.VisitUpdate(tc.updatePath, left, right)
+
+				if tc.updateError != nil {
+					assert.Equal(t, tc.updateError, err)
+				} else {
+					assert.NoError(t, err)
+					assert.Equal(t, right, out)
+				}
+			})
+		}
+
+		if tc.deletePath != nil {
+			t.Run(tc.name+"-delete", func(t *testing.T) {
+				err := visitor.VisitDelete(tc.deletePath, left)
+
+				if tc.deleteError != nil {
+					assert.Equal(t, tc.deleteError, err)
+				} else {
+					assert.NoError(t, err)
+				}
+			})
+		}
+
+		if tc.insertPath != nil {
+			t.Run(tc.name+"-insert", func(t *testing.T) {
+				out, err := visitor.VisitInsert(tc.insertPath, right)
+
+				if tc.insertError != nil {
+					assert.Equal(t, tc.insertError, err)
+				} else {
+					assert.NoError(t, err)
+					assert.Equal(t, right, out)
+				}
+			})
+		}
+	}
+}
+
+func withProcessStub(args []string, stdout string, stderr string) context.Context {
+	ctx := context.Background()
+	ctx, stub := process.WithStub(ctx)
+
+	stub.WithCallback(func(actual *exec.Cmd) error {
+		if reflect.DeepEqual(actual.Args, args) {
+			_, err := actual.Stdout.Write([]byte(stdout))
+
+			if err != nil {
+				return err
+			}
+
+			_, err = actual.Stderr.Write([]byte(stderr))
+
+			return err
+		} else {
+			return fmt.Errorf("unexpected command: %v", actual.Args)
+		}
+	})
+
+	return ctx
+}
+
+func withLoggerStub(ctx context.Context) (*bytes.Buffer, context.Context) {
+	var buf = bytes.Buffer{}
+
+	opts := slog.HandlerOptions{
+		Level: slog.LevelInfo,
+		ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr {
+			// remove everything except 'msg' and 'level' for less output
+			if a.Key == slog.MessageKey || a.Key == slog.LevelKey {
+				return a
+			} else {
+				return slog.Attr{}
+			}
+		},
+	}
+	handler := slog.NewTextHandler(&buf, &opts)
+	logger := slog.New(handler)
+
+	ctx = log.NewContext(ctx, logger)
+
+	return &buf, ctx
+}
+
+func keys[T any](value map[string]T) []string {
+	keys := make([]string, 0, len(value))
+
+	for k := range value {
+		keys = append(keys, k)
+	}
+
+	return keys
+}
+
+func loadYaml(name string, content string) *bundle.Bundle {
+	v, diag := config.LoadFromBytes(name, []byte(content))
+
+	if diag.Error() != nil {
+		panic(diag.Error())
+	}
+
+	return &bundle.Bundle{
+		Config: *v,
+	}
+}
+
+func withFakeVEnv(t *testing.T, path string) {
+	cwd, err := os.Getwd()
+
+	if err != nil {
+		panic(err)
+	}
+
+	if err := os.Chdir(t.TempDir()); err != nil {
+		panic(err)
+	}
+
+	err = os.MkdirAll(filepath.Join(path, "bin"), 0755)
+
+	if err != nil {
+		panic(err)
+	}
+
+	err = os.WriteFile(filepath.Join(".venv", "bin/python3"), []byte(""), 0755)
+
+	if err != nil {
+		panic(err)
+	}
+
+	t.Cleanup(func() {
+		if err := os.Chdir(cwd); err != nil {
+			panic(err)
+		}
+	})
+}
diff --git a/bundle/config/mutator/python/python_logger.go b/bundle/config/mutator/python/python_logger.go
new file mode 100644
index 0000000000..9c9e1f3620
--- /dev/null
+++ b/bundle/config/mutator/python/python_logger.go
@@ -0,0 +1,95 @@
+package python
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"strings"
+
+	"github.com/databricks/cli/libs/log"
+)
+
+type stderrLogger struct {
+	ctx context.Context
+	buf bytes.Buffer
+}
+
+type level string
+
+const levelError level = "ERROR"
+const levelInfo level = "INFO"
+const levelDebug level = "DEBUG"
+const levelWarning level = "WARNING"
+
+type logEntry struct {
+	Level   level  `json:"level,omitempty"`
+	Message string `json:"message"`
+}
+
+// NewStderrLogger creates a new io.Writer that parses structured logs from stderr
+// and logs them using CLI logger with appropriate log level.
+func NewStderrLogger(ctx context.Context) io.Writer {
+	return &stderrLogger{
+		ctx: ctx,
+		buf: bytes.Buffer{},
+	}
+}
+
+func (p *stderrLogger) Write(bytes []byte) (n int, err error) {
+	p.buf.Write(bytes)
+
+	scanner := bufio.NewScanner(&p.buf)
+
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		if parsed, ok := parseLogEntry(line); ok {
+			p.writeLogEntry(parsed)
+		} else {
+			log.Debugf(p.ctx, "stderr: %s", line)
+		}
+	}
+
+	remaining := p.buf.Bytes()
+	p.buf.Reset()
+	p.buf.Write(remaining)
+
+	return len(bytes), nil
+}
+
+func (p *stderrLogger) writeLogEntry(entry logEntry) {
+	switch entry.Level {
+	case levelInfo:
+		log.Infof(p.ctx, "%s", entry.Message)
+	case levelError:
+		log.Errorf(p.ctx, "%s", entry.Message)
+	case levelWarning:
+		log.Warnf(p.ctx, "%s", entry.Message)
+	case levelDebug:
+		log.Debugf(p.ctx, "%s", entry.Message)
+	default:
+		log.Debugf(p.ctx, "%s", entry.Message)
+	}
+}
+
+func parseLogEntry(line string) (logEntry, bool) {
+	if !strings.HasPrefix(line, "{") {
+		return logEntry{}, false
+	}
+
+	if !strings.HasSuffix(line, "}") {
+		return logEntry{}, false
+	}
+
+	var out logEntry
+
+	err := json.Unmarshal([]byte(line), &out)
+
+	if err != nil {
+		return logEntry{}, false
+	}
+
+	return out, true
+}
diff --git a/bundle/config/mutator/python/python_logger_test.go b/bundle/config/mutator/python/python_logger_test.go
new file mode 100644
index 0000000000..1d23c28c46
--- /dev/null
+++ b/bundle/config/mutator/python/python_logger_test.go
@@ -0,0 +1,65 @@
+package python
+
+import (
+	"testing"
+
+	assert "github.com/databricks/cli/libs/dyn/dynassert"
+)
+
+type parseLogEntryTestCase struct {
+	text     string
+	ok       bool
+	expected logEntry
+}
+
+func TestParseLogEntry(t *testing.T) {
+
+	testCases := []parseLogEntryTestCase{
+		{
+			text: "{\"level\":\"INFO\",\"message\":\"hello\"}",
+			ok:   true,
+			expected: logEntry{
+				Level:   "INFO",
+				Message: "hello",
+			},
+		},
+		{
+			text: "{\"level\":\"DEBUG\",\"message\":\"hello\"}",
+			ok:   true,
+			expected: logEntry{
+				Level:   "DEBUG",
+				Message: "hello",
+			},
+		},
+		{
+			text: "{\"level\":\"WARNING\",\"message\":\"hello\"}",
+			ok:   true,
+			expected: logEntry{
+				Level:   "WARNING",
+				Message: "hello",
+			},
+		},
+		{
+			text: "{\"level\":\"ERROR\",\"message\":\"hello\"}",
+			ok:   true,
+			expected: logEntry{
+				Level:   "ERROR",
+				Message: "hello",
+			},
+		},
+		{
+			text:     "hi {} there",
+			ok:       false,
+			expected: logEntry{},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.text, func(t *testing.T) {
+			actual, ok := parseLogEntry(tc.text)
+
+			assert.Equal(t, tc.ok, ok)
+			assert.Equal(t, tc.expected, actual)
+		})
+	}
+}
diff --git a/bundle/config/root.go b/bundle/config/root.go
index 88197c2b87..6f530661af 100644
--- a/bundle/config/root.go
+++ b/bundle/config/root.go
@@ -74,6 +74,10 @@ func Load(path string) (*Root, diag.Diagnostics) {
 		return nil, diag.FromErr(err)
 	}
 
+	return LoadFromBytes(path, raw)
+}
+
+func LoadFromBytes(path string, raw []byte) (*Root, diag.Diagnostics) {
 	r := Root{}
 
 	// Load configuration tree from YAML.
diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go
index ded2e19808..be5723d2fd 100644
--- a/bundle/phases/initialize.go
+++ b/bundle/phases/initialize.go
@@ -4,6 +4,7 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/mutator"
+	pythonmutator "github.com/databricks/cli/bundle/config/mutator/python"
 	"github.com/databricks/cli/bundle/deploy/metadata"
 	"github.com/databricks/cli/bundle/deploy/terraform"
 	"github.com/databricks/cli/bundle/permissions"
@@ -28,6 +29,7 @@ func Initialize() bundle.Mutator {
 			mutator.ExpandWorkspaceRoot(),
 			mutator.DefineDefaultWorkspacePaths(),
 			mutator.SetVariables(),
+			pythonmutator.ApplyPythonMutator(pythonmutator.ApplyPythonMutatorPhaseInit),
 			mutator.ResolveVariableReferencesInLookup(),
 			mutator.ResolveResourceReferences(),
 			mutator.ResolveVariableReferences(
diff --git a/libs/process/opts.go b/libs/process/opts.go
index e201c66684..a860dc8147 100644
--- a/libs/process/opts.go
+++ b/libs/process/opts.go
@@ -48,6 +48,20 @@ func WithStdoutPipe(dst *io.ReadCloser) execOption {
 	}
 }
 
+func WithStdinReader(src io.Reader) execOption {
+	return func(_ context.Context, c *exec.Cmd) error {
+		c.Stdin = src
+		return nil
+	}
+}
+
+func WithStderrWriter(writer io.Writer) execOption {
+	return func(_ context.Context, c *exec.Cmd) error {
+		c.Stderr = writer
+		return nil
+	}
+}
+
 func WithCombinedOutput(buf *bytes.Buffer) execOption {
 	return func(_ context.Context, c *exec.Cmd) error {
 		c.Stdout = io.MultiWriter(buf, c.Stdout)

From 16fc994b9d2a81782cee9998d560b5672557bbe7 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Fri, 17 May 2024 16:20:17 +0200
Subject: [PATCH 02/10] Address CR comments

---
 bundle/config/experimental.go                 | 11 ++-
 .../mutator/python/apply_python_mutator.go    | 20 ++--
 .../python/apply_python_mutator_test.go       | 98 +++++--------------
 bundle/config/mutator/python/log_writer.go    | 42 ++++++++
 bundle/config/mutator/python/python_logger.go | 95 ------------------
 .../mutator/python/python_logger_test.go      | 65 ------------
 libs/process/opts.go                          | 11 ++-
 7 files changed, 91 insertions(+), 251 deletions(-)
 create mode 100644 bundle/config/mutator/python/log_writer.go
 delete mode 100644 bundle/config/mutator/python/python_logger.go
 delete mode 100644 bundle/config/mutator/python/python_logger_test.go

diff --git a/bundle/config/experimental.go b/bundle/config/experimental.go
index 0c1bcac6ea..b38b97e647 100644
--- a/bundle/config/experimental.go
+++ b/bundle/config/experimental.go
@@ -24,11 +24,11 @@ type Experimental struct {
 	// as a top-level permission in the DAB.
 	UseLegacyRunAs bool `json:"use_legacy_run_as,omitempty"`
 
-	// EnablePyDABs if enabled loads the 'databricks-pydabs' package.
+	// PyDABs if enabled loads the 'databricks-pydabs' package.
 	//
 	// This allows to define bundle configuration using Python. For now, this option has to be used
 	// together with 'venv' configuration.
-	EnablePyDABs bool `json:"enable_pydabs,omitempty"`
+	PyDABs PyDABs `json:"pydabs,omitempty"`
 
 	// VEnv is the virtual environment configuration for the bundle.
 	//
@@ -37,10 +37,15 @@ type Experimental struct {
 }
 
 type VEnv struct {
-	// Path to the virtual environment
+	// Path to the virtual environment.
 	Path string `json:"path,omitempty"`
 }
 
+type PyDABs struct {
+	// Enable is a flag to enable the feature.
+	Enable bool `json:"enable,omitempty"`
+}
+
 type Command string
 type ScriptHook string
 
diff --git a/bundle/config/mutator/python/apply_python_mutator.go b/bundle/config/mutator/python/apply_python_mutator.go
index ffa34d4313..89191f14c6 100644
--- a/bundle/config/mutator/python/apply_python_mutator.go
+++ b/bundle/config/mutator/python/apply_python_mutator.go
@@ -5,6 +5,10 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"os"
+	"path"
+	"path/filepath"
+
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/libs/diag"
@@ -14,9 +18,6 @@ import (
 	"github.com/databricks/cli/libs/dyn/yamlloader"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/process"
-	"os"
-	"path"
-	"path/filepath"
 )
 
 type phase string
@@ -51,13 +52,13 @@ func getExperimental(b *bundle.Bundle) config.Experimental {
 func (m *applyPythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 	experimental := getExperimental(b)
 
-	if !experimental.EnablePyDABs {
-		log.Debugf(ctx, "'experimental.enable_pydabs' isn't enabled, skipping")
+	if !experimental.PyDABs.Enable {
+		log.Debugf(ctx, "'experimental.pydabs.enable' isn't enabled, skipping")
 		return nil
 	}
 
-	if experimental.EnablePyDABs && experimental.VEnv.Path == "" {
-		return diag.Errorf("'experimental.enable_pydabs' can only be used when 'experimental.venv.path' is set")
+	if experimental.PyDABs.Enable && experimental.VEnv.Path == "" {
+		return diag.Errorf("'experimental.pydabs.enable' can only be used when 'experimental.venv.path' is set")
 	}
 
 	err := b.Config.Mutate(func(leftRoot dyn.Value) (dyn.Value, error) {
@@ -100,13 +101,13 @@ func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath stri
 		return dyn.InvalidValue, fmt.Errorf("failed to marshal root config: %w", err)
 	}
 
-	pythonLogger := NewStderrLogger(ctx)
+	logWriter := newLogWriter(ctx, "stderr: ")
 
 	stdout, err := process.Background(
 		ctx,
 		args,
 		process.WithDir(rootPath),
-		process.WithStderrWriter(pythonLogger),
+		process.WithStderrWriter(logWriter),
 		process.WithStdinReader(bytes.NewBuffer(rootConfigJson)),
 	)
 
@@ -136,7 +137,6 @@ func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath stri
 
 	return normalized, nil
 }
-
 func createOverrideVisitor(ctx context.Context, phase phase) merge.OverrideVisitor {
 	jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs"))
 
diff --git a/bundle/config/mutator/python/apply_python_mutator_test.go b/bundle/config/mutator/python/apply_python_mutator_test.go
index f892ce513f..90d9cb24ce 100644
--- a/bundle/config/mutator/python/apply_python_mutator_test.go
+++ b/bundle/config/mutator/python/apply_python_mutator_test.go
@@ -1,21 +1,19 @@
 package python
 
 import (
-	"bytes"
 	"context"
 	"fmt"
-	"github.com/databricks/cli/libs/dyn"
-	"log/slog"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"reflect"
 	"testing"
 
+	"github.com/databricks/cli/libs/dyn"
+
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	assert "github.com/databricks/cli/libs/dyn/dynassert"
-	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/process"
 )
 
@@ -36,7 +34,8 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 
 	b := loadYaml("bundle.yaml", `
       experimental:
-        enable_pydabs: true
+        pydabs:
+          enable: true
         venv:
           path: .venv
       resources:
@@ -54,7 +53,7 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 		},
 		`{
 			"experimental": {
-				"enable_pydabs": true,
+				"pydabs": { "enable": true },
 				"venv": { "path": ".venv" }
 			},
 			"resources": {
@@ -67,11 +66,7 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 					},
 				}
 			}
-		}`,
-		`{"level": "INFO", "message": "Applying Python mutator"}`,
-	)
-
-	loggerBuf, ctx := withLoggerStub(ctx)
+		}`)
 
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
 	diag := bundle.Apply(ctx, b, mutator)
@@ -87,8 +82,6 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 	if job1, ok := b.Config.Resources.Jobs["job1"]; ok {
 		assert.Equal(t, "job_1", job1.Name)
 	}
-
-	assert.Equal(t, "level=INFO msg=\"Applying Python mutator\"\n", loggerBuf.String())
 }
 
 func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
@@ -96,7 +89,8 @@ func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
 
 	b := loadYaml("bundle.yaml", `
       experimental:
-        enable_pydabs: true
+        pydabs:
+          enable: true
         venv:
           path: .venv
       resources:
@@ -114,7 +108,7 @@ func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
 		},
 		`{
 			"experimental": {
-				"enable_pydabs": true,
+				"pydabs": { "enable": true },
 				"venv": { "path": ".venv" }
 			},
 			"resources": {
@@ -125,9 +119,7 @@ func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
 					}
 				}
 			}
-		}`,
-		"",
-	)
+		}`)
 
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
 	diag := bundle.Apply(ctx, b, mutator)
@@ -140,7 +132,8 @@ func TestApplyPythonMutator_init(t *testing.T) {
 
 	b := loadYaml("bundle.yaml", `
       experimental:
-        enable_pydabs: true
+        pydabs:
+          enable: true
         venv:
           path: .venv
       resources:
@@ -158,7 +151,7 @@ func TestApplyPythonMutator_init(t *testing.T) {
 		},
 		`{
 			"experimental": {
-				"enable_pydabs": true,
+				"pydabs": { "enable": true },
 				"venv": { "path": ".venv" }
 			},
 			"resources": {
@@ -169,10 +162,7 @@ func TestApplyPythonMutator_init(t *testing.T) {
 					}
 				}
 			}
-		}`,
-		`{"level": "INFO", "message": "Applying Python mutator"}`,
-	)
-	loggerBuf, ctx := withLoggerStub(ctx)
+		}`)
 
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseInit)
 	diag := bundle.Apply(ctx, b, mutator)
@@ -182,21 +172,10 @@ func TestApplyPythonMutator_init(t *testing.T) {
 	assert.ElementsMatch(t, []string{"job0"}, keys(b.Config.Resources.Jobs))
 	assert.Equal(t, "job_0", b.Config.Resources.Jobs["job0"].Name)
 	assert.Equal(t, "my job", b.Config.Resources.Jobs["job0"].Description)
-
-	assert.Equal(t, "level=INFO msg=\"Applying Python mutator\"\n", loggerBuf.String())
 }
 
 func TestApplyPythonMutator_disabled(t *testing.T) {
-	b := &bundle.Bundle{
-		Config: config.Root{
-			Experimental: &config.Experimental{
-				EnablePyDABs: false,
-				VEnv: config.VEnv{
-					Path: ".venv",
-				},
-			},
-		},
-	}
+	b := loadYaml("bundle.yaml", ``)
 
 	ctx := context.Background()
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
@@ -206,14 +185,10 @@ func TestApplyPythonMutator_disabled(t *testing.T) {
 }
 
 func TestApplyPythonMutator_venvRequired(t *testing.T) {
-	b := &bundle.Bundle{
-		Config: config.Root{
-			Experimental: &config.Experimental{
-				EnablePyDABs: true,
-				VEnv:         config.VEnv{},
-			},
-		},
-	}
+	b := loadYaml("bundle.yaml", `
+      experimental:
+        pydabs:
+          enable: true`)
 
 	ctx := context.Background()
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
@@ -223,11 +198,10 @@ func TestApplyPythonMutator_venvRequired(t *testing.T) {
 }
 
 func TestApplyPythonMutator_venvNotFound(t *testing.T) {
-	withFakeVEnv(t, ".venv")
-
 	b := loadYaml("bundle.yaml", `
       experimental:
-        enable_pydabs: true
+        pydabs:
+          enable: true
         venv:
           path: bad_path`)
 
@@ -362,7 +336,7 @@ func TestCreateOverrideVisitor(t *testing.T) {
 	}
 }
 
-func withProcessStub(args []string, stdout string, stderr string) context.Context {
+func withProcessStub(args []string, stdout string) context.Context {
 	ctx := context.Background()
 	ctx, stub := process.WithStub(ctx)
 
@@ -370,12 +344,6 @@ func withProcessStub(args []string, stdout string, stderr string) context.Contex
 		if reflect.DeepEqual(actual.Args, args) {
 			_, err := actual.Stdout.Write([]byte(stdout))
 
-			if err != nil {
-				return err
-			}
-
-			_, err = actual.Stderr.Write([]byte(stderr))
-
 			return err
 		} else {
 			return fmt.Errorf("unexpected command: %v", actual.Args)
@@ -385,28 +353,6 @@ func withProcessStub(args []string, stdout string, stderr string) context.Contex
 	return ctx
 }
 
-func withLoggerStub(ctx context.Context) (*bytes.Buffer, context.Context) {
-	var buf = bytes.Buffer{}
-
-	opts := slog.HandlerOptions{
-		Level: slog.LevelInfo,
-		ReplaceAttr: func(groups []string, a slog.Attr) slog.Attr {
-			// remove everything except 'msg' and 'level' for less output
-			if a.Key == slog.MessageKey || a.Key == slog.LevelKey {
-				return a
-			} else {
-				return slog.Attr{}
-			}
-		},
-	}
-	handler := slog.NewTextHandler(&buf, &opts)
-	logger := slog.New(handler)
-
-	ctx = log.NewContext(ctx, logger)
-
-	return &buf, ctx
-}
-
 func keys[T any](value map[string]T) []string {
 	keys := make([]string, 0, len(value))
 
diff --git a/bundle/config/mutator/python/log_writer.go b/bundle/config/mutator/python/log_writer.go
new file mode 100644
index 0000000000..aa3db05719
--- /dev/null
+++ b/bundle/config/mutator/python/log_writer.go
@@ -0,0 +1,42 @@
+package python
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"io"
+
+	"github.com/databricks/cli/libs/log"
+)
+
+type logWriter struct {
+	ctx    context.Context
+	prefix string
+	buf    bytes.Buffer
+}
+
+// newLogWriter creates a new io.Writer that writes to log with specified prefix.
+func newLogWriter(ctx context.Context, prefix string) io.Writer {
+	return &logWriter{
+		ctx:    ctx,
+		prefix: prefix,
+	}
+}
+
+func (p *logWriter) Write(bytes []byte) (n int, err error) {
+	p.buf.Write(bytes)
+
+	scanner := bufio.NewScanner(&p.buf)
+
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		log.Debugf(p.ctx, "%s%s", p.prefix, line)
+	}
+
+	remaining := p.buf.Bytes()
+	p.buf.Reset()
+	p.buf.Write(remaining)
+
+	return len(bytes), nil
+}
diff --git a/bundle/config/mutator/python/python_logger.go b/bundle/config/mutator/python/python_logger.go
deleted file mode 100644
index 9c9e1f3620..0000000000
--- a/bundle/config/mutator/python/python_logger.go
+++ /dev/null
@@ -1,95 +0,0 @@
-package python
-
-import (
-	"bufio"
-	"bytes"
-	"context"
-	"encoding/json"
-	"io"
-	"strings"
-
-	"github.com/databricks/cli/libs/log"
-)
-
-type stderrLogger struct {
-	ctx context.Context
-	buf bytes.Buffer
-}
-
-type level string
-
-const levelError level = "ERROR"
-const levelInfo level = "INFO"
-const levelDebug level = "DEBUG"
-const levelWarning level = "WARNING"
-
-type logEntry struct {
-	Level   level  `json:"level,omitempty"`
-	Message string `json:"message"`
-}
-
-// NewStderrLogger creates a new io.Writer that parses structured logs from stderr
-// and logs them using CLI logger with appropriate log level.
-func NewStderrLogger(ctx context.Context) io.Writer {
-	return &stderrLogger{
-		ctx: ctx,
-		buf: bytes.Buffer{},
-	}
-}
-
-func (p *stderrLogger) Write(bytes []byte) (n int, err error) {
-	p.buf.Write(bytes)
-
-	scanner := bufio.NewScanner(&p.buf)
-
-	for scanner.Scan() {
-		line := scanner.Text()
-
-		if parsed, ok := parseLogEntry(line); ok {
-			p.writeLogEntry(parsed)
-		} else {
-			log.Debugf(p.ctx, "stderr: %s", line)
-		}
-	}
-
-	remaining := p.buf.Bytes()
-	p.buf.Reset()
-	p.buf.Write(remaining)
-
-	return len(bytes), nil
-}
-
-func (p *stderrLogger) writeLogEntry(entry logEntry) {
-	switch entry.Level {
-	case levelInfo:
-		log.Infof(p.ctx, "%s", entry.Message)
-	case levelError:
-		log.Errorf(p.ctx, "%s", entry.Message)
-	case levelWarning:
-		log.Warnf(p.ctx, "%s", entry.Message)
-	case levelDebug:
-		log.Debugf(p.ctx, "%s", entry.Message)
-	default:
-		log.Debugf(p.ctx, "%s", entry.Message)
-	}
-}
-
-func parseLogEntry(line string) (logEntry, bool) {
-	if !strings.HasPrefix(line, "{") {
-		return logEntry{}, false
-	}
-
-	if !strings.HasSuffix(line, "}") {
-		return logEntry{}, false
-	}
-
-	var out logEntry
-
-	err := json.Unmarshal([]byte(line), &out)
-
-	if err != nil {
-		return logEntry{}, false
-	}
-
-	return out, true
-}
diff --git a/bundle/config/mutator/python/python_logger_test.go b/bundle/config/mutator/python/python_logger_test.go
deleted file mode 100644
index 1d23c28c46..0000000000
--- a/bundle/config/mutator/python/python_logger_test.go
+++ /dev/null
@@ -1,65 +0,0 @@
-package python
-
-import (
-	"testing"
-
-	assert "github.com/databricks/cli/libs/dyn/dynassert"
-)
-
-type parseLogEntryTestCase struct {
-	text     string
-	ok       bool
-	expected logEntry
-}
-
-func TestParseLogEntry(t *testing.T) {
-
-	testCases := []parseLogEntryTestCase{
-		{
-			text: "{\"level\":\"INFO\",\"message\":\"hello\"}",
-			ok:   true,
-			expected: logEntry{
-				Level:   "INFO",
-				Message: "hello",
-			},
-		},
-		{
-			text: "{\"level\":\"DEBUG\",\"message\":\"hello\"}",
-			ok:   true,
-			expected: logEntry{
-				Level:   "DEBUG",
-				Message: "hello",
-			},
-		},
-		{
-			text: "{\"level\":\"WARNING\",\"message\":\"hello\"}",
-			ok:   true,
-			expected: logEntry{
-				Level:   "WARNING",
-				Message: "hello",
-			},
-		},
-		{
-			text: "{\"level\":\"ERROR\",\"message\":\"hello\"}",
-			ok:   true,
-			expected: logEntry{
-				Level:   "ERROR",
-				Message: "hello",
-			},
-		},
-		{
-			text:     "hi {} there",
-			ok:       false,
-			expected: logEntry{},
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.text, func(t *testing.T) {
-			actual, ok := parseLogEntry(tc.text)
-
-			assert.Equal(t, tc.ok, ok)
-			assert.Equal(t, tc.expected, actual)
-		})
-	}
-}
diff --git a/libs/process/opts.go b/libs/process/opts.go
index a860dc8147..9516e49ba5 100644
--- a/libs/process/opts.go
+++ b/libs/process/opts.go
@@ -55,9 +55,16 @@ func WithStdinReader(src io.Reader) execOption {
 	}
 }
 
-func WithStderrWriter(writer io.Writer) execOption {
+func WithStderrWriter(dst io.Writer) execOption {
 	return func(_ context.Context, c *exec.Cmd) error {
-		c.Stderr = writer
+		c.Stderr = dst
+		return nil
+	}
+}
+
+func WithStdoutWriter(dst io.Writer) execOption {
+	return func(_ context.Context, c *exec.Cmd) error {
+		c.Stdout = dst
 		return nil
 	}
 }

From 05c661dd4ea174c8e062990e4a7a238b086221fb Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Mon, 17 Jun 2024 16:12:32 +0200
Subject: [PATCH 03/10] Address CR comments

---
 bundle/config/experimental.go                 |   6 +-
 .../mutator/python/apply_python_mutator.go    | 145 ++++++++++++++----
 .../python/apply_python_mutator_test.go       |  99 +++++++++---
 bundle/phases/initialize.go                   |   2 +
 4 files changed, 196 insertions(+), 56 deletions(-)

diff --git a/bundle/config/experimental.go b/bundle/config/experimental.go
index b38b97e647..2eaf965bc4 100644
--- a/bundle/config/experimental.go
+++ b/bundle/config/experimental.go
@@ -24,10 +24,10 @@ type Experimental struct {
 	// as a top-level permission in the DAB.
 	UseLegacyRunAs bool `json:"use_legacy_run_as,omitempty"`
 
-	// PyDABs if enabled loads the 'databricks-pydabs' package.
+	// PyDABs determines whether to load the 'databricks-pydabs' package.
 	//
-	// This allows to define bundle configuration using Python. For now, this option has to be used
-	// together with 'venv' configuration.
+	// PyDABs allows to define bundle configuration using Python. For now, this option has to
+	// be used together with 'venv' configuration.
 	PyDABs PyDABs `json:"pydabs,omitempty"`
 
 	// VEnv is the virtual environment configuration for the bundle.
diff --git a/bundle/config/mutator/python/apply_python_mutator.go b/bundle/config/mutator/python/apply_python_mutator.go
index 89191f14c6..bb6a42de1a 100644
--- a/bundle/config/mutator/python/apply_python_mutator.go
+++ b/bundle/config/mutator/python/apply_python_mutator.go
@@ -6,8 +6,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
-	"path"
 	"path/filepath"
+	"runtime"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
@@ -23,8 +23,34 @@ import (
 type phase string
 
 const (
+	// ApplyPythonMutatorPhasePreInit is the phase that while bundle configuration is loading.
+	//
+	// At this stage, PyDABs adds statically defined resources to the bundle configuration.
+	// Which resources are added should be deterministic and not depend on the bundle configuration.
+	//
+	// We also open for possibility of appending other sections of bundle configuration,
+	// for example, adding new variables. However, this is not supported yet, and CLI rejects
+	// such changes.
 	ApplyPythonMutatorPhasePreInit phase = "preinit"
-	ApplyPythonMutatorPhaseInit    phase = "init"
+
+	// ApplyPythonMutatorPhaseInit is the phase after bundle configuration was loaded, and
+	// the list of statically declared resources is known.
+	//
+	// At this stage, PyDABs adds resources defined using generators, or mutates existing resources,
+	// including the ones defined using YAML.
+	//
+	// During this process, within generator and mutators, PyDABs can access:
+	// - selected deployment target
+	// - bundle variables values
+	//
+	// The following is not available:
+	// - variables referencing other variables are in unresolved format
+	// - variables provided through CLI arguments
+	//
+	// PyDABs can output YAML containing references to variables, and CLI should resolve them.
+	//
+	// Existing resources can't be removed, and CLI rejects such changes.
+	ApplyPythonMutatorPhaseInit phase = "init"
 )
 
 type applyPythonMutator struct {
@@ -53,22 +79,21 @@ func (m *applyPythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.D
 	experimental := getExperimental(b)
 
 	if !experimental.PyDABs.Enable {
-		log.Debugf(ctx, "'experimental.pydabs.enable' isn't enabled, skipping")
 		return nil
 	}
 
 	if experimental.PyDABs.Enable && experimental.VEnv.Path == "" {
-		return diag.Errorf("'experimental.pydabs.enable' can only be used when 'experimental.venv.path' is set")
+		return diag.Errorf("\"experimental.pydabs.enable\" can only be used when \"experimental.venv.path\" is set")
 	}
 
 	err := b.Config.Mutate(func(leftRoot dyn.Value) (dyn.Value, error) {
-		pythonPath := path.Join(experimental.VEnv.Path, "bin", "python3")
+		pythonPath := interpreterPath(experimental.VEnv.Path)
 
 		if _, err := os.Stat(pythonPath); err != nil {
 			if os.IsNotExist(err) {
-				return dyn.InvalidValue, fmt.Errorf("can't find '%s', check if venv is created", pythonPath)
+				return dyn.InvalidValue, fmt.Errorf("can't find %q, check if venv is created", pythonPath)
 			} else {
-				return dyn.InvalidValue, fmt.Errorf("can't find '%s': %w", pythonPath, err)
+				return dyn.InvalidValue, fmt.Errorf("can't find %q: %w", pythonPath, err)
 			}
 		}
 
@@ -78,7 +103,13 @@ func (m *applyPythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.D
 			return dyn.InvalidValue, err
 		}
 
-		return merge.Override(leftRoot, rightRoot, createOverrideVisitor(ctx, m.phase))
+		visitor, err := createOverrideVisitor(ctx, m.phase)
+
+		if err != nil {
+			return dyn.InvalidValue, err
+		}
+
+		return merge.Override(leftRoot, rightRoot, visitor)
 	})
 
 	return diag.FromErr(err)
@@ -117,7 +148,7 @@ func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath stri
 
 	// we need absolute path, or because later parts of pipeline assume all paths are absolute
 	// and this file will be used as location
-	virtualPath, err := filepath.Abs(filepath.Join(rootPath, "__generated__.yml"))
+	virtualPath, err := filepath.Abs(filepath.Join(rootPath, "__generated_by_pydabs__.yml"))
 
 	if err != nil {
 		return dyn.InvalidValue, fmt.Errorf("failed to get absolute path: %w", err)
@@ -135,56 +166,110 @@ func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath stri
 		return dyn.InvalidValue, fmt.Errorf("failed to normalize Python mutator output: %w", diagnostic.Error())
 	}
 
+	// warnings shouldn't happen because output should be already normalized
+	// when it happens, it's a bug in the mutator, and should be treated as an error
+
+	for _, d := range diagnostic.Filter(diag.Warning) {
+		return dyn.InvalidValue, fmt.Errorf("failed to normalize Python mutator output: %s", d.Summary)
+	}
+
 	return normalized, nil
 }
-func createOverrideVisitor(ctx context.Context, phase phase) merge.OverrideVisitor {
+
+func createOverrideVisitor(ctx context.Context, phase phase) (merge.OverrideVisitor, error) {
+	switch phase {
+	case ApplyPythonMutatorPhasePreInit:
+		return createPreInitOverrideVisitor(ctx), nil
+	case ApplyPythonMutatorPhaseInit:
+		return createInitOverrideVisitor(ctx), nil
+	default:
+		return merge.OverrideVisitor{}, fmt.Errorf("unknown phase: %s", phase)
+	}
+}
+
+// createPreInitOverrideVisitor creates an override visitor for the preinit phase.
+//
+// During pre-init it's only possible to create new resources, and not modify or
+// delete existing ones.
+func createPreInitOverrideVisitor(ctx context.Context) merge.OverrideVisitor {
 	jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs"))
 
 	return merge.OverrideVisitor{
 		VisitDelete: func(valuePath dyn.Path, left dyn.Value) error {
+			return fmt.Errorf("unexpected change at %q (delete)", valuePath.String())
+		},
+		VisitInsert: func(valuePath dyn.Path, right dyn.Value) (dyn.Value, error) {
 			if !valuePath.HasPrefix(jobsPath) {
-				return fmt.Errorf("unexpected change at '%s' (delete)", valuePath.String())
+				return dyn.InvalidValue, fmt.Errorf("unexpected change at %q (insert)", valuePath.String())
 			}
 
-			deleteResource := len(valuePath) == len(jobsPath)+1
+			insertResource := len(valuePath) == len(jobsPath)+1
 
-			if phase == ApplyPythonMutatorPhasePreInit {
-				return fmt.Errorf("unexpected change at '%s' (delete)", valuePath.String())
-			} else if deleteResource && phase == ApplyPythonMutatorPhaseInit {
-				return fmt.Errorf("unexpected change at '%s' (delete)", valuePath.String())
+			// adding a property into an existing resource is not allowed, because it changes it
+			if !insertResource {
+				return dyn.InvalidValue, fmt.Errorf("unexpected change at %q (insert)", valuePath.String())
 			}
 
-			log.Debugf(ctx, "Delete value at '%s'", valuePath.String())
+			log.Debugf(ctx, "Insert value at %q", valuePath.String())
 
-			return nil
+			return right, nil
 		},
-		VisitInsert: func(valuePath dyn.Path, right dyn.Value) (dyn.Value, error) {
+		VisitUpdate: func(valuePath dyn.Path, left dyn.Value, right dyn.Value) (dyn.Value, error) {
+			return dyn.InvalidValue, fmt.Errorf("unexpected change at %q (update)", valuePath.String())
+		},
+	}
+}
+
+// createInitOverrideVisitor creates an override visitor for the init phase.
+//
+// During the init phase it's possible to create new resources, modify existing
+// resources, but not delete existing resources.
+func createInitOverrideVisitor(ctx context.Context) merge.OverrideVisitor {
+	jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs"))
+
+	return merge.OverrideVisitor{
+		VisitDelete: func(valuePath dyn.Path, left dyn.Value) error {
 			if !valuePath.HasPrefix(jobsPath) {
-				return dyn.InvalidValue, fmt.Errorf("unexpected change at '%s' (insert)", valuePath.String())
+				return fmt.Errorf("unexpected change at %q (delete)", valuePath.String())
 			}
 
-			insertResource := len(valuePath) == len(jobsPath)+1
+			deleteResource := len(valuePath) == len(jobsPath)+1
 
-			if phase == ApplyPythonMutatorPhasePreInit && !insertResource {
-				return dyn.InvalidValue, fmt.Errorf("unexpected change at '%s' (insert)", valuePath.String())
+			if deleteResource {
+				return fmt.Errorf("unexpected change at %q (delete)", valuePath.String())
 			}
 
-			log.Debugf(ctx, "Insert value at '%s'", valuePath.String())
+			// deleting properties is allowed because it only changes an existing resource
+			log.Debugf(ctx, "Delete value at %q", valuePath.String())
 
-			return right, nil
+			return nil
 		},
-		VisitUpdate: func(valuePath dyn.Path, left dyn.Value, right dyn.Value) (dyn.Value, error) {
+		VisitInsert: func(valuePath dyn.Path, right dyn.Value) (dyn.Value, error) {
 			if !valuePath.HasPrefix(jobsPath) {
-				return dyn.InvalidValue, fmt.Errorf("unexpected change at '%s' (update)", valuePath.String())
+				return dyn.InvalidValue, fmt.Errorf("unexpected change at %q (insert)", valuePath.String())
 			}
 
-			if phase == ApplyPythonMutatorPhasePreInit {
-				return dyn.InvalidValue, fmt.Errorf("unexpected change at '%s' (update)", valuePath.String())
+			log.Debugf(ctx, "Insert value at %q", valuePath.String())
+
+			return right, nil
+		},
+		VisitUpdate: func(valuePath dyn.Path, left dyn.Value, right dyn.Value) (dyn.Value, error) {
+			if !valuePath.HasPrefix(jobsPath) {
+				return dyn.InvalidValue, fmt.Errorf("unexpected change at %q (update)", valuePath.String())
 			}
 
-			log.Debugf(ctx, "Update value at '%s'", valuePath.String())
+			log.Debugf(ctx, "Update value at %q", valuePath.String())
 
 			return right, nil
 		},
 	}
 }
+
+// interpreterPath returns platform-specific path to Python interpreter in the virtual environment.
+func interpreterPath(venvPath string) string {
+	if runtime.GOOS == "windows" {
+		return filepath.Join(venvPath, "Scripts", "python3.exe")
+	} else {
+		return filepath.Join(venvPath, "bin", "python3")
+	}
+}
diff --git a/bundle/config/mutator/python/apply_python_mutator_test.go b/bundle/config/mutator/python/apply_python_mutator_test.go
index 90d9cb24ce..c2d7a83748 100644
--- a/bundle/config/mutator/python/apply_python_mutator_test.go
+++ b/bundle/config/mutator/python/apply_python_mutator_test.go
@@ -7,6 +7,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"reflect"
+	"runtime"
 	"testing"
 
 	"github.com/databricks/cli/libs/dyn"
@@ -32,7 +33,7 @@ func TestApplyPythonMutator_Name_init(t *testing.T) {
 func TestApplyPythonMutator_preinit(t *testing.T) {
 	withFakeVEnv(t, ".venv")
 
-	b := loadYaml("bundle.yaml", `
+	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
           enable: true
@@ -45,7 +46,7 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 
 	ctx := withProcessStub(
 		[]string{
-			".venv/bin/python3",
+			interpreterPath(".venv"),
 			"-m",
 			"databricks.bundles.build",
 			"--phase",
@@ -87,7 +88,7 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
 	withFakeVEnv(t, ".venv")
 
-	b := loadYaml("bundle.yaml", `
+	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
           enable: true
@@ -124,13 +125,13 @@ func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
 	diag := bundle.Apply(ctx, b, mutator)
 
-	assert.EqualError(t, diag.Error(), "unexpected change at 'resources.jobs.job0.description' (insert)")
+	assert.EqualError(t, diag.Error(), "unexpected change at \"resources.jobs.job0.description\" (insert)")
 }
 
 func TestApplyPythonMutator_init(t *testing.T) {
 	withFakeVEnv(t, ".venv")
 
-	b := loadYaml("bundle.yaml", `
+	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
           enable: true
@@ -143,7 +144,7 @@ func TestApplyPythonMutator_init(t *testing.T) {
 
 	ctx := withProcessStub(
 		[]string{
-			".venv/bin/python3",
+			interpreterPath(".venv"),
 			"-m",
 			"databricks.bundles.build",
 			"--phase",
@@ -174,8 +175,46 @@ func TestApplyPythonMutator_init(t *testing.T) {
 	assert.Equal(t, "my job", b.Config.Resources.Jobs["job0"].Description)
 }
 
+func TestApplyPythonMutator_badOutput(t *testing.T) {
+	withFakeVEnv(t, ".venv")
+
+	b := loadYaml("databricks.yml", `
+      experimental:
+        pydabs:
+          enable: true
+        venv:
+          path: .venv
+      resources:
+        jobs:
+          job0:
+            name: job_0`)
+
+	ctx := withProcessStub(
+		[]string{
+			interpreterPath(".venv"),
+			"-m",
+			"databricks.bundles.build",
+			"--phase",
+			"preinit",
+		},
+		`{
+			"resources": {
+				"jobs": {
+					"job0": {
+						unknown_property: "my job"
+					}
+				}
+			}
+		}`)
+
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	diag := bundle.Apply(ctx, b, mutator)
+
+	assert.EqualError(t, diag.Error(), "failed to normalize Python mutator output: unknown field: unknown_property")
+}
+
 func TestApplyPythonMutator_disabled(t *testing.T) {
-	b := loadYaml("bundle.yaml", ``)
+	b := loadYaml("databricks.yml", ``)
 
 	ctx := context.Background()
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
@@ -185,7 +224,7 @@ func TestApplyPythonMutator_disabled(t *testing.T) {
 }
 
 func TestApplyPythonMutator_venvRequired(t *testing.T) {
-	b := loadYaml("bundle.yaml", `
+	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
           enable: true`)
@@ -194,11 +233,13 @@ func TestApplyPythonMutator_venvRequired(t *testing.T) {
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
 	diag := bundle.Apply(ctx, b, mutator)
 
-	assert.Error(t, diag.Error(), "'experimental.enable_pydabs' is enabled, but 'experimental.venv.path' is not set")
+	assert.Error(t, diag.Error(), "\"experimental.enable_pydabs\" is enabled, but \"experimental.venv.path\" is not set")
 }
 
 func TestApplyPythonMutator_venvNotFound(t *testing.T) {
-	b := loadYaml("bundle.yaml", `
+	expectedError := fmt.Sprintf("can't find %q, check if venv is created", interpreterPath("bad_path"))
+
+	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
           enable: true
@@ -208,7 +249,7 @@ func TestApplyPythonMutator_venvNotFound(t *testing.T) {
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseInit)
 	diag := bundle.Apply(context.Background(), b, mutator)
 
-	assert.EqualError(t, diag.Error(), "can't find 'bad_path/bin/python3', check if venv is created")
+	assert.EqualError(t, diag.Error(), expectedError)
 }
 
 type createOverrideVisitorTestCase struct {
@@ -233,15 +274,15 @@ func TestCreateOverrideVisitor(t *testing.T) {
 			updatePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
 			deletePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
 			insertPath:  dyn.MustPathFromString("resources.jobs.job0.name"),
-			deleteError: fmt.Errorf("unexpected change at 'resources.jobs.job0.name' (delete)"),
-			insertError: fmt.Errorf("unexpected change at 'resources.jobs.job0.name' (insert)"),
-			updateError: fmt.Errorf("unexpected change at 'resources.jobs.job0.name' (update)"),
+			deleteError: fmt.Errorf("unexpected change at \"resources.jobs.job0.name\" (delete)"),
+			insertError: fmt.Errorf("unexpected change at \"resources.jobs.job0.name\" (insert)"),
+			updateError: fmt.Errorf("unexpected change at \"resources.jobs.job0.name\" (update)"),
 		},
 		{
 			name:        "preinit: can't delete an existing job",
 			phase:       ApplyPythonMutatorPhasePreInit,
 			deletePath:  dyn.MustPathFromString("resources.jobs.job0"),
-			deleteError: fmt.Errorf("unexpected change at 'resources.jobs.job0' (delete)"),
+			deleteError: fmt.Errorf("unexpected change at \"resources.jobs.job0\" (delete)"),
 		},
 		{
 			name:        "preinit: can insert a job",
@@ -255,9 +296,9 @@ func TestCreateOverrideVisitor(t *testing.T) {
 			deletePath:  dyn.MustPathFromString("include[0]"),
 			insertPath:  dyn.MustPathFromString("include[0]"),
 			updatePath:  dyn.MustPathFromString("include[0]"),
-			deleteError: fmt.Errorf("unexpected change at 'include[0]' (delete)"),
-			insertError: fmt.Errorf("unexpected change at 'include[0]' (insert)"),
-			updateError: fmt.Errorf("unexpected change at 'include[0]' (update)"),
+			deleteError: fmt.Errorf("unexpected change at \"include[0]\" (delete)"),
+			insertError: fmt.Errorf("unexpected change at \"include[0]\" (insert)"),
+			updateError: fmt.Errorf("unexpected change at \"include[0]\" (update)"),
 		},
 		{
 			name:        "preinit: can change an existing job",
@@ -273,7 +314,7 @@ func TestCreateOverrideVisitor(t *testing.T) {
 			name:        "init: can't delete an existing job",
 			phase:       ApplyPythonMutatorPhaseInit,
 			deletePath:  dyn.MustPathFromString("resources.jobs.job0"),
-			deleteError: fmt.Errorf("unexpected change at 'resources.jobs.job0' (delete)"),
+			deleteError: fmt.Errorf("unexpected change at \"resources.jobs.job0\" (delete)"),
 		},
 		{
 			name:        "init: can insert a job",
@@ -287,14 +328,18 @@ func TestCreateOverrideVisitor(t *testing.T) {
 			deletePath:  dyn.MustPathFromString("include[0]"),
 			insertPath:  dyn.MustPathFromString("include[0]"),
 			updatePath:  dyn.MustPathFromString("include[0]"),
-			deleteError: fmt.Errorf("unexpected change at 'include[0]' (delete)"),
-			insertError: fmt.Errorf("unexpected change at 'include[0]' (insert)"),
-			updateError: fmt.Errorf("unexpected change at 'include[0]' (update)"),
+			deleteError: fmt.Errorf("unexpected change at \"include[0]\" (delete)"),
+			insertError: fmt.Errorf("unexpected change at \"include[0]\" (insert)"),
+			updateError: fmt.Errorf("unexpected change at \"include[0]\" (update)"),
 		},
 	}
 
 	for _, tc := range testCases {
-		visitor := createOverrideVisitor(context.Background(), tc.phase)
+		visitor, err := createOverrideVisitor(context.Background(), tc.phase)
+
+		if err != nil {
+			t.Fatalf("create visitor failed: %v", err)
+		}
 
 		if tc.updatePath != nil {
 			t.Run(tc.name+"-update", func(t *testing.T) {
@@ -336,6 +381,14 @@ func TestCreateOverrideVisitor(t *testing.T) {
 	}
 }
 
+func TestInterpreterPath(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		assert.Equal(t, "venv\\Scripts\\python3.exe", interpreterPath("venv"))
+	} else {
+		assert.Equal(t, "venv/bin/python3", interpreterPath("venv"))
+	}
+}
+
 func withProcessStub(args []string, stdout string) context.Context {
 	ctx := context.Background()
 	ctx, stub := process.WithStub(ctx)
diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go
index be5723d2fd..d96c8d3b34 100644
--- a/bundle/phases/initialize.go
+++ b/bundle/phases/initialize.go
@@ -29,6 +29,8 @@ func Initialize() bundle.Mutator {
 			mutator.ExpandWorkspaceRoot(),
 			mutator.DefineDefaultWorkspacePaths(),
 			mutator.SetVariables(),
+			// Intentionally placed before ResolveVariableReferencesInLookup, ResolveResourceReferences
+			// and ResolveVariableReferences. See what is expected in ApplyPythonMutatorPhaseInit doc
 			pythonmutator.ApplyPythonMutator(pythonmutator.ApplyPythonMutatorPhaseInit),
 			mutator.ResolveVariableReferencesInLookup(),
 			mutator.ResolveResourceReferences(),

From 48d9328f773463dcca980724da213d9141167c05 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Mon, 17 Jun 2024 16:44:09 +0200
Subject: [PATCH 04/10] Fix tests on Windows

---
 bundle/config/mutator/python/apply_python_mutator_test.go | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/bundle/config/mutator/python/apply_python_mutator_test.go b/bundle/config/mutator/python/apply_python_mutator_test.go
index c2d7a83748..f580305206 100644
--- a/bundle/config/mutator/python/apply_python_mutator_test.go
+++ b/bundle/config/mutator/python/apply_python_mutator_test.go
@@ -430,6 +430,7 @@ func loadYaml(name string, content string) *bundle.Bundle {
 
 func withFakeVEnv(t *testing.T, path string) {
 	cwd, err := os.Getwd()
+	interpreterPath := interpreterPath(path)
 
 	if err != nil {
 		panic(err)
@@ -439,13 +440,13 @@ func withFakeVEnv(t *testing.T, path string) {
 		panic(err)
 	}
 
-	err = os.MkdirAll(filepath.Join(path, "bin"), 0755)
+	err = os.MkdirAll(filepath.Dir(interpreterPath), 0755)
 
 	if err != nil {
 		panic(err)
 	}
 
-	err = os.WriteFile(filepath.Join(".venv", "bin/python3"), []byte(""), 0755)
+	err = os.WriteFile(interpreterPath, []byte(""), 0755)
 
 	if err != nil {
 		panic(err)

From ebca9279f2351422eb557ec17e629318da93612c Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Tue, 18 Jun 2024 08:58:04 +0200
Subject: [PATCH 05/10] Fix tests on Windows

---
 bundle/config/mutator/python/apply_python_mutator_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bundle/config/mutator/python/apply_python_mutator_test.go b/bundle/config/mutator/python/apply_python_mutator_test.go
index f580305206..76a9e643e9 100644
--- a/bundle/config/mutator/python/apply_python_mutator_test.go
+++ b/bundle/config/mutator/python/apply_python_mutator_test.go
@@ -101,7 +101,7 @@ func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
 
 	ctx := withProcessStub(
 		[]string{
-			".venv/bin/python3",
+			interpreterPath(".venv"),
 			"-m",
 			"databricks.bundles.build",
 			"--phase",

From 3a49549b3f4a4201071411d85606a300ef12a46f Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Wed, 19 Jun 2024 15:28:42 +0200
Subject: [PATCH 06/10] Address CR comments

---
 .../mutator/python/apply_python_mutator.go    |  31 ++---
 .../python/apply_python_mutator_test.go       | 109 ++++++++----------
 2 files changed, 61 insertions(+), 79 deletions(-)

diff --git a/bundle/config/mutator/python/apply_python_mutator.go b/bundle/config/mutator/python/apply_python_mutator.go
index bb6a42de1a..298ffb5769 100644
--- a/bundle/config/mutator/python/apply_python_mutator.go
+++ b/bundle/config/mutator/python/apply_python_mutator.go
@@ -23,7 +23,7 @@ import (
 type phase string
 
 const (
-	// ApplyPythonMutatorPhasePreInit is the phase that while bundle configuration is loading.
+	// ApplyPythonMutatorPhaseLoad is the phase in which bundle configuration is loaded.
 	//
 	// At this stage, PyDABs adds statically defined resources to the bundle configuration.
 	// Which resources are added should be deterministic and not depend on the bundle configuration.
@@ -31,7 +31,7 @@ const (
 	// We also open for possibility of appending other sections of bundle configuration,
 	// for example, adding new variables. However, this is not supported yet, and CLI rejects
 	// such changes.
-	ApplyPythonMutatorPhasePreInit phase = "preinit"
+	ApplyPythonMutatorPhaseLoad phase = "load"
 
 	// ApplyPythonMutatorPhaseInit is the phase after bundle configuration was loaded, and
 	// the list of statically declared resources is known.
@@ -42,10 +42,10 @@ const (
 	// During this process, within generator and mutators, PyDABs can access:
 	// - selected deployment target
 	// - bundle variables values
+	// - variables provided through CLI arguments or environment variables
 	//
 	// The following is not available:
 	// - variables referencing other variables are in unresolved format
-	// - variables provided through CLI arguments
 	//
 	// PyDABs can output YAML containing references to variables, and CLI should resolve them.
 	//
@@ -78,16 +78,16 @@ func getExperimental(b *bundle.Bundle) config.Experimental {
 func (m *applyPythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
 	experimental := getExperimental(b)
 
-	if !experimental.PyDABs.Enable {
+	if !experimental.PyDABs.Enabled {
 		return nil
 	}
 
-	if experimental.PyDABs.Enable && experimental.VEnv.Path == "" {
-		return diag.Errorf("\"experimental.pydabs.enable\" can only be used when \"experimental.venv.path\" is set")
+	if experimental.PyDABs.VEnvPath == "" {
+		return diag.Errorf("\"experimental.pydabs.enabled\" can only be used when \"experimental.pydabs.venv_path\" is set")
 	}
 
 	err := b.Config.Mutate(func(leftRoot dyn.Value) (dyn.Value, error) {
-		pythonPath := interpreterPath(experimental.VEnv.Path)
+		pythonPath := interpreterPath(experimental.PyDABs.VEnvPath)
 
 		if _, err := os.Stat(pythonPath); err != nil {
 			if os.IsNotExist(err) {
@@ -98,13 +98,11 @@ func (m *applyPythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.D
 		}
 
 		rightRoot, err := m.runPythonMutator(ctx, b.RootPath, pythonPath, leftRoot)
-
 		if err != nil {
 			return dyn.InvalidValue, err
 		}
 
 		visitor, err := createOverrideVisitor(ctx, m.phase)
-
 		if err != nil {
 			return dyn.InvalidValue, err
 		}
@@ -127,7 +125,6 @@ func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath stri
 	// we need to marshal dyn.Value instead of bundle.Config to JSON to support
 	// non-string fields assigned with bundle variables
 	rootConfigJson, err := json.Marshal(root.AsAny())
-
 	if err != nil {
 		return dyn.InvalidValue, fmt.Errorf("failed to marshal root config: %w", err)
 	}
@@ -141,7 +138,6 @@ func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath stri
 		process.WithStderrWriter(logWriter),
 		process.WithStdinReader(bytes.NewBuffer(rootConfigJson)),
 	)
-
 	if err != nil {
 		return dyn.InvalidValue, fmt.Errorf("python mutator process failed: %w", err)
 	}
@@ -149,19 +145,16 @@ func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath stri
 	// we need absolute path, or because later parts of pipeline assume all paths are absolute
 	// and this file will be used as location
 	virtualPath, err := filepath.Abs(filepath.Join(rootPath, "__generated_by_pydabs__.yml"))
-
 	if err != nil {
 		return dyn.InvalidValue, fmt.Errorf("failed to get absolute path: %w", err)
 	}
 
 	generated, err := yamlloader.LoadYAML(virtualPath, bytes.NewReader([]byte(stdout)))
-
 	if err != nil {
 		return dyn.InvalidValue, fmt.Errorf("failed to parse Python mutator output: %w", err)
 	}
 
 	normalized, diagnostic := convert.Normalize(config.Root{}, generated)
-
 	if diagnostic.Error() != nil {
 		return dyn.InvalidValue, fmt.Errorf("failed to normalize Python mutator output: %w", diagnostic.Error())
 	}
@@ -178,8 +171,8 @@ func (m *applyPythonMutator) runPythonMutator(ctx context.Context, rootPath stri
 
 func createOverrideVisitor(ctx context.Context, phase phase) (merge.OverrideVisitor, error) {
 	switch phase {
-	case ApplyPythonMutatorPhasePreInit:
-		return createPreInitOverrideVisitor(ctx), nil
+	case ApplyPythonMutatorPhaseLoad:
+		return createLoadOverrideVisitor(ctx), nil
 	case ApplyPythonMutatorPhaseInit:
 		return createInitOverrideVisitor(ctx), nil
 	default:
@@ -187,11 +180,11 @@ func createOverrideVisitor(ctx context.Context, phase phase) (merge.OverrideVisi
 	}
 }
 
-// createPreInitOverrideVisitor creates an override visitor for the preinit phase.
+// createLoadOverrideVisitor creates an override visitor for the load phase.
 //
-// During pre-init it's only possible to create new resources, and not modify or
+// During load, it's only possible to create new resources, and not modify or
 // delete existing ones.
-func createPreInitOverrideVisitor(ctx context.Context) merge.OverrideVisitor {
+func createLoadOverrideVisitor(ctx context.Context) merge.OverrideVisitor {
 	jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs"))
 
 	return merge.OverrideVisitor{
diff --git a/bundle/config/mutator/python/apply_python_mutator_test.go b/bundle/config/mutator/python/apply_python_mutator_test.go
index 76a9e643e9..4c8084848e 100644
--- a/bundle/config/mutator/python/apply_python_mutator_test.go
+++ b/bundle/config/mutator/python/apply_python_mutator_test.go
@@ -3,6 +3,7 @@ package python
 import (
 	"context"
 	"fmt"
+	"golang.org/x/exp/maps"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -18,10 +19,10 @@ import (
 	"github.com/databricks/cli/libs/process"
 )
 
-func TestApplyPythonMutator_Name_preinit(t *testing.T) {
-	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+func TestApplyPythonMutator_Name_load(t *testing.T) {
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseLoad)
 
-	assert.Equal(t, "ApplyPythonMutator(preinit)", mutator.Name())
+	assert.Equal(t, "ApplyPythonMutator(load)", mutator.Name())
 }
 
 func TestApplyPythonMutator_Name_init(t *testing.T) {
@@ -30,15 +31,14 @@ func TestApplyPythonMutator_Name_init(t *testing.T) {
 	assert.Equal(t, "ApplyPythonMutator(init)", mutator.Name())
 }
 
-func TestApplyPythonMutator_preinit(t *testing.T) {
+func TestApplyPythonMutator_load(t *testing.T) {
 	withFakeVEnv(t, ".venv")
 
 	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
-          enable: true
-        venv:
-          path: .venv
+          enabled: true
+          venv_path: .venv
       resources:
         jobs:
           job0:
@@ -50,12 +50,14 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 			"-m",
 			"databricks.bundles.build",
 			"--phase",
-			"preinit",
+			"load",
 		},
 		`{
 			"experimental": {
-				"pydabs": { "enable": true },
-				"venv": { "path": ".venv" }
+				"pydabs": {
+					"enabled": true,
+					"venv_path": ".venv"
+				}
 			},
 			"resources": {
 				"jobs": {
@@ -69,12 +71,12 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 			}
 		}`)
 
-	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseLoad)
 	diag := bundle.Apply(ctx, b, mutator)
 
 	assert.NoError(t, diag.Error())
 
-	assert.ElementsMatch(t, []string{"job0", "job1"}, keys(b.Config.Resources.Jobs))
+	assert.ElementsMatch(t, []string{"job0", "job1"}, maps.Keys(b.Config.Resources.Jobs))
 
 	if job0, ok := b.Config.Resources.Jobs["job0"]; ok {
 		assert.Equal(t, "job_0", job0.Name)
@@ -85,15 +87,14 @@ func TestApplyPythonMutator_preinit(t *testing.T) {
 	}
 }
 
-func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
+func TestApplyPythonMutator_load_disallowed(t *testing.T) {
 	withFakeVEnv(t, ".venv")
 
 	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
-          enable: true
-        venv:
-          path: .venv
+          enabled: true
+          venv_path: .venv
       resources:
         jobs:
           job0:
@@ -105,12 +106,14 @@ func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
 			"-m",
 			"databricks.bundles.build",
 			"--phase",
-			"preinit",
+			"load",
 		},
 		`{
 			"experimental": {
-				"pydabs": { "enable": true },
-				"venv": { "path": ".venv" }
+				"pydabs": {
+					"enabled": true,
+					"venv_path": ".venv"
+				}
 			},
 			"resources": {
 				"jobs": {
@@ -122,7 +125,7 @@ func TestApplyPythonMutator_preinit_disallowed(t *testing.T) {
 			}
 		}`)
 
-	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseLoad)
 	diag := bundle.Apply(ctx, b, mutator)
 
 	assert.EqualError(t, diag.Error(), "unexpected change at \"resources.jobs.job0.description\" (insert)")
@@ -134,9 +137,8 @@ func TestApplyPythonMutator_init(t *testing.T) {
 	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
-          enable: true
-        venv:
-          path: .venv
+          enabled: true
+          venv_path: .venv
       resources:
         jobs:
           job0:
@@ -152,8 +154,10 @@ func TestApplyPythonMutator_init(t *testing.T) {
 		},
 		`{
 			"experimental": {
-				"pydabs": { "enable": true },
-				"venv": { "path": ".venv" }
+				"pydabs": {
+					"enabled": true,
+					"venv_path": ".venv"
+				}
 			},
 			"resources": {
 				"jobs": {
@@ -170,7 +174,7 @@ func TestApplyPythonMutator_init(t *testing.T) {
 
 	assert.NoError(t, diag.Error())
 
-	assert.ElementsMatch(t, []string{"job0"}, keys(b.Config.Resources.Jobs))
+	assert.ElementsMatch(t, []string{"job0"}, maps.Keys(b.Config.Resources.Jobs))
 	assert.Equal(t, "job_0", b.Config.Resources.Jobs["job0"].Name)
 	assert.Equal(t, "my job", b.Config.Resources.Jobs["job0"].Description)
 }
@@ -181,9 +185,8 @@ func TestApplyPythonMutator_badOutput(t *testing.T) {
 	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
-          enable: true
-        venv:
-          path: .venv
+          enabled: true
+          venv_path: .venv
       resources:
         jobs:
           job0:
@@ -195,7 +198,7 @@ func TestApplyPythonMutator_badOutput(t *testing.T) {
 			"-m",
 			"databricks.bundles.build",
 			"--phase",
-			"preinit",
+			"load",
 		},
 		`{
 			"resources": {
@@ -207,7 +210,7 @@ func TestApplyPythonMutator_badOutput(t *testing.T) {
 			}
 		}`)
 
-	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseLoad)
 	diag := bundle.Apply(ctx, b, mutator)
 
 	assert.EqualError(t, diag.Error(), "failed to normalize Python mutator output: unknown field: unknown_property")
@@ -217,7 +220,7 @@ func TestApplyPythonMutator_disabled(t *testing.T) {
 	b := loadYaml("databricks.yml", ``)
 
 	ctx := context.Background()
-	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseLoad)
 	diag := bundle.Apply(ctx, b, mutator)
 
 	assert.NoError(t, diag.Error())
@@ -227,10 +230,10 @@ func TestApplyPythonMutator_venvRequired(t *testing.T) {
 	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
-          enable: true`)
+          enabled: true`)
 
 	ctx := context.Background()
-	mutator := ApplyPythonMutator(ApplyPythonMutatorPhasePreInit)
+	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseLoad)
 	diag := bundle.Apply(ctx, b, mutator)
 
 	assert.Error(t, diag.Error(), "\"experimental.enable_pydabs\" is enabled, but \"experimental.venv.path\" is not set")
@@ -242,9 +245,8 @@ func TestApplyPythonMutator_venvNotFound(t *testing.T) {
 	b := loadYaml("databricks.yml", `
       experimental:
         pydabs:
-          enable: true
-        venv:
-          path: bad_path`)
+          enabled: true
+          venv_path: bad_path`)
 
 	mutator := ApplyPythonMutator(ApplyPythonMutatorPhaseInit)
 	diag := bundle.Apply(context.Background(), b, mutator)
@@ -269,8 +271,8 @@ func TestCreateOverrideVisitor(t *testing.T) {
 
 	testCases := []createOverrideVisitorTestCase{
 		{
-			name:        "preinit: can't change an existing job",
-			phase:       ApplyPythonMutatorPhasePreInit,
+			name:        "load: can't change an existing job",
+			phase:       ApplyPythonMutatorPhaseLoad,
 			updatePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
 			deletePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
 			insertPath:  dyn.MustPathFromString("resources.jobs.job0.name"),
@@ -279,20 +281,20 @@ func TestCreateOverrideVisitor(t *testing.T) {
 			updateError: fmt.Errorf("unexpected change at \"resources.jobs.job0.name\" (update)"),
 		},
 		{
-			name:        "preinit: can't delete an existing job",
-			phase:       ApplyPythonMutatorPhasePreInit,
+			name:        "load: can't delete an existing job",
+			phase:       ApplyPythonMutatorPhaseLoad,
 			deletePath:  dyn.MustPathFromString("resources.jobs.job0"),
 			deleteError: fmt.Errorf("unexpected change at \"resources.jobs.job0\" (delete)"),
 		},
 		{
-			name:        "preinit: can insert a job",
-			phase:       ApplyPythonMutatorPhasePreInit,
+			name:        "load: can insert a job",
+			phase:       ApplyPythonMutatorPhaseLoad,
 			insertPath:  dyn.MustPathFromString("resources.jobs.job0"),
 			insertError: nil,
 		},
 		{
-			name:        "preinit: can't change include",
-			phase:       ApplyPythonMutatorPhasePreInit,
+			name:        "load: can't change include",
+			phase:       ApplyPythonMutatorPhaseLoad,
 			deletePath:  dyn.MustPathFromString("include[0]"),
 			insertPath:  dyn.MustPathFromString("include[0]"),
 			updatePath:  dyn.MustPathFromString("include[0]"),
@@ -301,7 +303,7 @@ func TestCreateOverrideVisitor(t *testing.T) {
 			updateError: fmt.Errorf("unexpected change at \"include[0]\" (update)"),
 		},
 		{
-			name:        "preinit: can change an existing job",
+			name:        "init: can change an existing job",
 			phase:       ApplyPythonMutatorPhaseInit,
 			updatePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
 			deletePath:  dyn.MustPathFromString("resources.jobs.job0.name"),
@@ -336,7 +338,6 @@ func TestCreateOverrideVisitor(t *testing.T) {
 
 	for _, tc := range testCases {
 		visitor, err := createOverrideVisitor(context.Background(), tc.phase)
-
 		if err != nil {
 			t.Fatalf("create visitor failed: %v", err)
 		}
@@ -406,16 +407,6 @@ func withProcessStub(args []string, stdout string) context.Context {
 	return ctx
 }
 
-func keys[T any](value map[string]T) []string {
-	keys := make([]string, 0, len(value))
-
-	for k := range value {
-		keys = append(keys, k)
-	}
-
-	return keys
-}
-
 func loadYaml(name string, content string) *bundle.Bundle {
 	v, diag := config.LoadFromBytes(name, []byte(content))
 
@@ -429,9 +420,9 @@ func loadYaml(name string, content string) *bundle.Bundle {
 }
 
 func withFakeVEnv(t *testing.T, path string) {
-	cwd, err := os.Getwd()
 	interpreterPath := interpreterPath(path)
 
+	cwd, err := os.Getwd()
 	if err != nil {
 		panic(err)
 	}
@@ -441,13 +432,11 @@ func withFakeVEnv(t *testing.T, path string) {
 	}
 
 	err = os.MkdirAll(filepath.Dir(interpreterPath), 0755)
-
 	if err != nil {
 		panic(err)
 	}
 
 	err = os.WriteFile(interpreterPath, []byte(""), 0755)
-
 	if err != nil {
 		panic(err)
 	}

From 74a41f4bec014a96a6e67066b832470573fb4d3c Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Wed, 19 Jun 2024 15:37:01 +0200
Subject: [PATCH 07/10] Address CR comments

---
 bundle/config/experimental.go | 20 ++++++++------------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/bundle/config/experimental.go b/bundle/config/experimental.go
index 2eaf965bc4..a77362fd49 100644
--- a/bundle/config/experimental.go
+++ b/bundle/config/experimental.go
@@ -29,21 +29,17 @@ type Experimental struct {
 	// PyDABs allows to define bundle configuration using Python. For now, this option has to
 	// be used together with 'venv' configuration.
 	PyDABs PyDABs `json:"pydabs,omitempty"`
-
-	// VEnv is the virtual environment configuration for the bundle.
-	//
-	// When enabled, the bundle will use Python interpreter from the virtual environment.
-	VEnv VEnv `json:"venv,omitempty"`
-}
-
-type VEnv struct {
-	// Path to the virtual environment.
-	Path string `json:"path,omitempty"`
 }
 
 type PyDABs struct {
-	// Enable is a flag to enable the feature.
-	Enable bool `json:"enable,omitempty"`
+	// Enabled is a flag to enable the feature.
+	Enabled bool `json:"enabled,omitempty"`
+
+	// VEnvPath is path to the virtual environment.
+	//
+	// Required if PyDABs is enabled. PyDABs will load the code in the specified
+	// environment.
+	VEnvPath string `json:"venv_path,omitempty"`
 }
 
 type Command string

From c5de423c011699362cb0f3b56f05605e3b9bc251 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Wed, 19 Jun 2024 15:42:40 +0200
Subject: [PATCH 08/10] Fix fmt

---
 bundle/config/mutator/python/apply_python_mutator_test.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/bundle/config/mutator/python/apply_python_mutator_test.go b/bundle/config/mutator/python/apply_python_mutator_test.go
index 4c8084848e..8759ab8016 100644
--- a/bundle/config/mutator/python/apply_python_mutator_test.go
+++ b/bundle/config/mutator/python/apply_python_mutator_test.go
@@ -3,7 +3,6 @@ package python
 import (
 	"context"
 	"fmt"
-	"golang.org/x/exp/maps"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -11,6 +10,8 @@ import (
 	"runtime"
 	"testing"
 
+	"golang.org/x/exp/maps"
+
 	"github.com/databricks/cli/libs/dyn"
 
 	"github.com/databricks/cli/bundle"

From ab50f40b9b067a4a6e77fd862969ea05d29dbf2c Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Wed, 19 Jun 2024 15:44:17 +0200
Subject: [PATCH 09/10] Fix build

---
 bundle/config/mutator/mutator.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bundle/config/mutator/mutator.go b/bundle/config/mutator/mutator.go
index 1e188d1a66..42be7081cf 100644
--- a/bundle/config/mutator/mutator.go
+++ b/bundle/config/mutator/mutator.go
@@ -25,7 +25,7 @@ func DefaultMutators() []bundle.Mutator {
 		InitializeVariables(),
 		DefineDefaultTarget(),
 		LoadGitDetails(),
-		pythonmutator.ApplyPythonMutator(pythonmutator.ApplyPythonMutatorPhasePreInit),
+		pythonmutator.ApplyPythonMutator(pythonmutator.ApplyPythonMutatorPhaseLoad),
 	}
 }
 

From 8b31ad891b5227db866da4afd8b559bb4d9f0a60 Mon Sep 17 00:00:00 2001
From: Gleb Kanterov <gleb.kanterov@databricks.com>
Date: Thu, 20 Jun 2024 09:17:43 +0200
Subject: [PATCH 10/10] Fix comment

---
 bundle/config/experimental.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/bundle/config/experimental.go b/bundle/config/experimental.go
index a77362fd49..12048a3227 100644
--- a/bundle/config/experimental.go
+++ b/bundle/config/experimental.go
@@ -26,8 +26,7 @@ type Experimental struct {
 
 	// PyDABs determines whether to load the 'databricks-pydabs' package.
 	//
-	// PyDABs allows to define bundle configuration using Python. For now, this option has to
-	// be used together with 'venv' configuration.
+	// PyDABs allows to define bundle configuration using Python.
 	PyDABs PyDABs `json:"pydabs,omitempty"`
 }