Merge 96a6cef0d6 into abc2f3c825

## Changes
The ML production team modified mlops-stack to use `mode: development`
for their development target here:
https://github.com/databricks/mlops-stacks/pull/174

This PR makes the integration test assertion agnostic of the prefix to
make it pass again.

## Tests
The test passes now

.codegen.json

@@ -5,8 +5,7 @@
     },
     "batch": {
         ".codegen/cmds-workspace.go.tmpl": "cmd/workspace/cmd.go",
-        ".codegen/cmds-account.go.tmpl": "cmd/account/cmd.go",
-        ".codegen/lookup.go.tmpl": "bundle/config/variable/lookup.go"
+        ".codegen/cmds-account.go.tmpl": "cmd/account/cmd.go"
     },
     "toolchain": {
         "required": ["go"],

.codegen/lookup.go.tmpl

@@ -1,134 +0,0 @@
-// Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
-
-package variable
-
-{{ $allowlist :=
-	list
-		"alerts"
-		"clusters"
-		"cluster-policies"
-		"clusters"
-		"dashboards"
-		"instance-pools"
-		"jobs"
-		"metastores"
-		"pipelines"
-		"service-principals"
-		"queries"
-		"warehouses"
-}}
-
-{{	$customField :=
-		dict
-			"service-principals" "ApplicationId"
-}}
-
-import (
-	"context"
-	"fmt"
-
-	"github.com/databricks/databricks-sdk-go"
-)
-
-type Lookup struct {
-	{{range .Services -}}
-	{{- if in $allowlist .KebabName -}}
-	{{.Singular.PascalName}} string `json:"{{.Singular.SnakeName}},omitempty"`
-
-	{{end}}
-	{{- end}}
-}
-
-func LookupFromMap(m map[string]any) *Lookup {
-	l := &Lookup{}
-	{{range .Services -}}
-	{{- if in $allowlist .KebabName -}}
-	if v, ok := m["{{.Singular.SnakeName}}"]; ok {
-		l.{{.Singular.PascalName}} = v.(string)
-	}
-	{{end -}}
-	{{- end}}
-	return l
-}
-
-func (l *Lookup) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
-	if err := l.validate(); err != nil {
-		return "", err
-	}
-
-	r := allResolvers()
-	{{range .Services -}}
-	{{- if in $allowlist .KebabName -}}
-	if l.{{.Singular.PascalName}} != "" {
-		return r.{{.Singular.PascalName}}(ctx, w, l.{{.Singular.PascalName}})
-	}
-	{{end -}}
-	{{- end}}
-
-	return "", fmt.Errorf("no valid lookup fields provided")
-}
-
-func (l *Lookup) String() string {
-	{{range .Services -}}
-	{{- if in $allowlist .KebabName -}}
-	if l.{{.Singular.PascalName}} != "" {
-		return fmt.Sprintf("{{.Singular.KebabName}}: %s", l.{{.Singular.PascalName}})
-	}
-	{{end -}}
-	{{- end}}
-	return ""
-}
-
-func (l *Lookup) validate() error {
-	// Validate that only one field is set
-	count := 0
-	{{range .Services -}}
-	{{- if in $allowlist .KebabName -}}
-	if l.{{.Singular.PascalName}} != "" {
-		count++
-	}
-	{{end -}}
-	{{- end}}
-
-	if count != 1 {
-		return fmt.Errorf("exactly one lookup field must be provided")
-	}
-
-	if strings.Contains(l.String(), "${var") {
-		return fmt.Errorf("lookup fields cannot contain variable references")
-	}
-
-	return nil
-}
-
-
-type resolverFunc func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error)
-type resolvers struct {
-	{{range .Services -}}
-	{{- if in $allowlist .KebabName -}}
-	{{.Singular.PascalName}} resolverFunc
-	{{end -}}
-	{{- end}}
-}
-
-func allResolvers() *resolvers {
-	r := &resolvers{}
-	{{range .Services -}}
-	{{- if in $allowlist .KebabName -}}
-	r.{{.Singular.PascalName}} = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["{{.Singular.PascalName}}"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.{{.PascalName}}.GetBy{{range .NamedIdMap.NamePath}}{{.PascalName}}{{end}}(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.{{ getOrDefault $customField .KebabName ((index .NamedIdMap.IdPath 0).PascalName) }}), nil
-	}
-	{{end -}}
-	{{- end}}
-
-	return r
-}

.gitattributes

@@ -1,4 +1,3 @@
-bundle/config/variable/lookup.go linguist-generated=true
 cmd/account/access-control/access-control.go linguist-generated=true
 cmd/account/billable-usage/billable-usage.go linguist-generated=true
 cmd/account/budgets/budgets.go linguist-generated=true

bundle/config/experimental.go

@@ -45,6 +45,12 @@ type PyDABs struct {
 	// These packages are imported to discover resources, resource generators, and mutators.
 	// This list can include namespace packages, which causes the import of nested packages.
 	Import []string `json:"import,omitempty"`
+
+	// LoadLocations is a flag to enable loading Python source locations from the PyDABs.
+	//
+	// Locations are only supported since PyDABs 0.6.0, and because of that,
+	// this flag is disabled by default.
+	LoadLocations bool `json:"load_locations,omitempty"`
 }
 
 type Command string

bundle/config/mutator/python/python_diagnostics.go

@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/cli/libs/dyn"
 )
 
+// pythonDiagnostic is a single entry in diagnostics.json
 type pythonDiagnostic struct {
 	Severity pythonSeverity           `json:"severity"`
 	Summary  string                   `json:"summary"`

bundle/config/mutator/python/python_locations.go

@@ -0,0 +1,181 @@
+package python
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"path/filepath"
+
+	"github.com/databricks/cli/libs/dyn"
+)
+
+// generatedFileName is used as the virtual file name for YAML generated by PyDABs.
+//
+// mergePythonLocations replaces dyn.Location with generatedFileName with locations loaded
+// from locations.json
+const generatedFileName = "__generated_by_pydabs__.yml"
+
+// pythonLocations is data structure for efficient location lookup for a given path
+//
+// Locations form a tree, and we assign locations of the closest ancestor to each dyn.Value based on its path.
+// We implement it as a trie (prefix tree) where keys are components of the path. With that, lookups are O(n)
+// where n is the number of components in the path.
+//
+// For example, with locations.json:
+//
+//	{"path": "resources.jobs.job_0", "file": "src/examples/job_0.py", "line": 3, "column": 5}
+//	{"path": "resources.jobs.job_0.tasks[0].task_key", "file": "src/examples/job_0.py", "line": 10, "column": 5}
+//	{"path": "resources.jobs.job_1", "file": "src/examples/job_1.py", "line": 5, "column": 7}
+//
+// - resources.jobs.job_0.tasks[0].task_key is located at job_0.py:10:5
+//
+//   - resources.jobs.job_0.tasks[0].email_notifications is located at job_0.py:3:5,
+//     because we use the location of the job as the most precise approximation.
+type pythonLocations struct {
+	// descendants referenced by index, e.g. '.foo'
+	keys map[string]*pythonLocations
+
+	// descendants referenced by key, e.g. '[0]'
+	indexes map[int]*pythonLocations
+
+	// location for the current node if it exists
+	location dyn.Location
+
+	// if true, location is present
+	exists bool
+}
+
+// pythonLocationEntry is a single entry in locations.json
+type pythonLocationEntry struct {
+	Path   string `json:"path"`
+	File   string `json:"file"`
+	Line   int    `json:"line"`
+	Column int    `json:"column"`
+}
+
+// mergePythonLocations applies locations from Python mutator into given dyn.Value
+//
+// The primary use-case is to merge locations.json with output.json, so that any
+// validation errors will point to Python source code instead of generated YAML.
+func mergePythonLocations(value dyn.Value, locations *pythonLocations) (dyn.Value, error) {
+	return dyn.Walk(value, func(path dyn.Path, value dyn.Value) (dyn.Value, error) {
+		newLocation, ok := findPythonLocation(locations, path)
+		if !ok {
+			return value, nil
+		}
+
+		var newLocations []dyn.Location
+
+		// the first item in the list is the "last" location used for error reporting
+		newLocations = append(newLocations, newLocation)
+
+		for _, location := range value.Locations() {
+			// When loaded, dyn.Value created by PyDABs use the virtual file path as their location,
+			// we replace it with newLocation.
+			if filepath.Base(location.File) == generatedFileName {
+				continue
+			}
+
+			newLocations = append(newLocations, location)
+		}
+
+		return value.WithLocations(newLocations), nil
+	})
+}
+
+// parsePythonLocations parses locations.json from the Python mutator.
+//
+// locations file is newline-separated JSON objects with pythonLocationEntry structure.
+func parsePythonLocations(input io.Reader) (*pythonLocations, error) {
+	decoder := json.NewDecoder(input)
+	locations := newPythonLocations()
+
+	for decoder.More() {
+		var entry pythonLocationEntry
+
+		err := decoder.Decode(&entry)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse python location: %s", err)
+		}
+
+		path, err := dyn.NewPathFromString(entry.Path)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse python location: %s", err)
+		}
+
+		location := dyn.Location{
+			File:   entry.File,
+			Line:   entry.Line,
+			Column: entry.Column,
+		}
+
+		putPythonLocation(locations, path, location)
+	}
+
+	return locations, nil
+}
+
+// putPythonLocation puts the location to the trie for the given path
+func putPythonLocation(trie *pythonLocations, path dyn.Path, location dyn.Location) {
+	var currentNode = trie
+
+	for _, component := range path {
+		if key := component.Key(); key != "" {
+			if _, ok := currentNode.keys[key]; !ok {
+				currentNode.keys[key] = newPythonLocations()
+			}
+
+			currentNode = currentNode.keys[key]
+		} else {
+			index := component.Index()
+			if _, ok := currentNode.indexes[index]; !ok {
+				currentNode.indexes[index] = newPythonLocations()
+			}
+
+			currentNode = currentNode.indexes[index]
+		}
+	}
+
+	currentNode.location = location
+	currentNode.exists = true
+}
+
+// newPythonLocations creates a new trie node
+func newPythonLocations() *pythonLocations {
+	return &pythonLocations{
+		keys:    make(map[string]*pythonLocations),
+		indexes: make(map[int]*pythonLocations),
+	}
+}
+
+// findPythonLocation finds the location or closest ancestor location in the trie for the given path
+// if no ancestor or exact location is found, false is returned.
+func findPythonLocation(locations *pythonLocations, path dyn.Path) (dyn.Location, bool) {
+	var currentNode = locations
+	var lastLocation = locations.location
+	var exists = locations.exists
+
+	for _, component := range path {
+		if key := component.Key(); key != "" {
+			if _, ok := currentNode.keys[key]; !ok {
+				break
+			}
+
+			currentNode = currentNode.keys[key]
+		} else {
+			index := component.Index()
+			if _, ok := currentNode.indexes[index]; !ok {
+				break
+			}
+
+			currentNode = currentNode.indexes[index]
+		}
+
+		if currentNode.exists {
+			lastLocation = currentNode.location
+			exists = true
+		}
+	}
+
+	return lastLocation, exists
+}

bundle/config/mutator/python/python_locations_test.go

@@ -0,0 +1,124 @@
+package python
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/databricks/cli/libs/dyn"
+	assert "github.com/databricks/cli/libs/dyn/dynassert"
+)
+
+func TestMergeLocations(t *testing.T) {
+	pythonLocation := dyn.Location{File: "foo.py", Line: 1, Column: 1}
+	generatedLocation := dyn.Location{File: generatedFileName, Line: 1, Column: 1}
+	yamlLocation := dyn.Location{File: "foo.yml", Line: 1, Column: 1}
+
+	locations := newPythonLocations()
+	putPythonLocation(locations, dyn.MustPathFromString("foo"), pythonLocation)
+
+	input := dyn.NewValue(
+		map[string]dyn.Value{
+			"foo": dyn.NewValue(
+				map[string]dyn.Value{
+					"baz": dyn.NewValue("baz", []dyn.Location{yamlLocation}),
+					"qux": dyn.NewValue("baz", []dyn.Location{generatedLocation, yamlLocation}),
+				},
+				[]dyn.Location{},
+			),
+			"bar": dyn.NewValue("baz", []dyn.Location{generatedLocation}),
+		},
+		[]dyn.Location{yamlLocation},
+	)
+
+	expected := dyn.NewValue(
+		map[string]dyn.Value{
+			"foo": dyn.NewValue(
+				map[string]dyn.Value{
+					// pythonLocation is appended to the beginning of the list if absent
+					"baz": dyn.NewValue("baz", []dyn.Location{pythonLocation, yamlLocation}),
+					// generatedLocation is replaced by pythonLocation
+					"qux": dyn.NewValue("baz", []dyn.Location{pythonLocation, yamlLocation}),
+				},
+				[]dyn.Location{pythonLocation},
+			),
+			// if location is unknown, we keep it as-is
+			"bar": dyn.NewValue("baz", []dyn.Location{generatedLocation}),
+		},
+		[]dyn.Location{yamlLocation},
+	)
+
+	actual, err := mergePythonLocations(input, locations)
+
+	assert.NoError(t, err)
+	assert.Equal(t, expected, actual)
+}
+
+func TestFindLocation(t *testing.T) {
+	location0 := dyn.Location{File: "foo.py", Line: 1, Column: 1}
+	location1 := dyn.Location{File: "foo.py", Line: 2, Column: 1}
+
+	locations := newPythonLocations()
+	putPythonLocation(locations, dyn.MustPathFromString("foo"), location0)
+	putPythonLocation(locations, dyn.MustPathFromString("foo.bar"), location1)
+
+	actual, exists := findPythonLocation(locations, dyn.MustPathFromString("foo.bar"))
+
+	assert.True(t, exists)
+	assert.Equal(t, location1, actual)
+}
+
+func TestFindLocation_indexPathComponent(t *testing.T) {
+	location0 := dyn.Location{File: "foo.py", Line: 1, Column: 1}
+	location1 := dyn.Location{File: "foo.py", Line: 2, Column: 1}
+	location2 := dyn.Location{File: "foo.py", Line: 3, Column: 1}
+
+	locations := newPythonLocations()
+	putPythonLocation(locations, dyn.MustPathFromString("foo"), location0)
+	putPythonLocation(locations, dyn.MustPathFromString("foo.bar"), location1)
+	putPythonLocation(locations, dyn.MustPathFromString("foo.bar[0]"), location2)
+
+	actual, exists := findPythonLocation(locations, dyn.MustPathFromString("foo.bar[0]"))
+
+	assert.True(t, exists)
+	assert.Equal(t, location2, actual)
+}
+
+func TestFindLocation_closestAncestorLocation(t *testing.T) {
+	location0 := dyn.Location{File: "foo.py", Line: 1, Column: 1}
+	location1 := dyn.Location{File: "foo.py", Line: 2, Column: 1}
+
+	locations := newPythonLocations()
+	putPythonLocation(locations, dyn.MustPathFromString("foo"), location0)
+	putPythonLocation(locations, dyn.MustPathFromString("foo.bar"), location1)
+
+	actual, exists := findPythonLocation(locations, dyn.MustPathFromString("foo.bar.baz"))
+
+	assert.True(t, exists)
+	assert.Equal(t, location1, actual)
+}
+
+func TestFindLocation_unknownLocation(t *testing.T) {
+	location0 := dyn.Location{File: "foo.py", Line: 1, Column: 1}
+	location1 := dyn.Location{File: "foo.py", Line: 2, Column: 1}
+
+	locations := newPythonLocations()
+	putPythonLocation(locations, dyn.MustPathFromString("foo"), location0)
+	putPythonLocation(locations, dyn.MustPathFromString("foo.bar"), location1)
+
+	_, exists := findPythonLocation(locations, dyn.MustPathFromString("bar"))
+
+	assert.False(t, exists)
+}
+
+func TestParsePythonLocations(t *testing.T) {
+	expected := dyn.Location{File: "foo.py", Line: 1, Column: 2}
+
+	input := `{"path": "foo", "file": "foo.py", "line": 1, "column": 2}`
+	reader := bytes.NewReader([]byte(input))
+	locations, err := parsePythonLocations(reader)
+
+	assert.NoError(t, err)
+
+	assert.True(t, locations.keys["foo"].exists)
+	assert.Equal(t, expected, locations.keys["foo"].location)
+}

bundle/config/mutator/python/python_mutator.go

@@ -7,9 +7,12 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/fs"
 	"os"
 	"path/filepath"
 
+	"github.com/databricks/cli/bundle/config/mutator/paths"
+
 	"github.com/databricks/databricks-sdk-go/logger"
 	"github.com/fatih/color"
 
@@ -108,7 +111,12 @@ func (m *pythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno
 			return dyn.InvalidValue, fmt.Errorf("failed to create cache dir: %w", err)
 		}
 
-		rightRoot, diags := m.runPythonMutator(ctx, cacheDir, b.BundleRootPath, pythonPath, leftRoot)
+		rightRoot, diags := m.runPythonMutator(ctx, leftRoot, runPythonMutatorOpts{
+			cacheDir:       cacheDir,
+			bundleRootPath: b.BundleRootPath,
+			pythonPath:     pythonPath,
+			loadLocations:  experimental.PyDABs.LoadLocations,
+		})
 		mutateDiags = diags
 		if diags.HasError() {
 			return dyn.InvalidValue, mutateDiagsHasError
@@ -152,13 +160,21 @@ func createCacheDir(ctx context.Context) (string, error) {
 	return os.MkdirTemp("", "-pydabs")
 }
 
-func (m *pythonMutator) runPythonMutator(ctx context.Context, cacheDir string, rootPath string, pythonPath string, root dyn.Value) (dyn.Value, diag.Diagnostics) {
-	inputPath := filepath.Join(cacheDir, "input.json")
-	outputPath := filepath.Join(cacheDir, "output.json")
-	diagnosticsPath := filepath.Join(cacheDir, "diagnostics.json")
+type runPythonMutatorOpts struct {
+	cacheDir       string
+	bundleRootPath string
+	pythonPath     string
+	loadLocations  bool
+}
+
+func (m *pythonMutator) runPythonMutator(ctx context.Context, root dyn.Value, opts runPythonMutatorOpts) (dyn.Value, diag.Diagnostics) {
+	inputPath := filepath.Join(opts.cacheDir, "input.json")
+	outputPath := filepath.Join(opts.cacheDir, "output.json")
+	diagnosticsPath := filepath.Join(opts.cacheDir, "diagnostics.json")
+	locationsPath := filepath.Join(opts.cacheDir, "locations.json")
 
 	args := []string{
-		pythonPath,
+		opts.pythonPath,
 		"-m",
 		"databricks.bundles.build",
 		"--phase",
@@ -171,6 +187,10 @@ func (m *pythonMutator) runPythonMutator(ctx context.Context, cacheDir string, r
 		diagnosticsPath,
 	}
 
+	if opts.loadLocations {
+		args = append(args, "--locations", locationsPath)
+	}
+
 	if err := writeInputFile(inputPath, root); err != nil {
 		return dyn.InvalidValue, diag.Errorf("failed to write input file: %s", err)
 	}
@@ -185,7 +205,7 @@ func (m *pythonMutator) runPythonMutator(ctx context.Context, cacheDir string, r
 	_, processErr := process.Background(
 		ctx,
 		args,
-		process.WithDir(rootPath),
+		process.WithDir(opts.bundleRootPath),
 		process.WithStderrWriter(stderrWriter),
 		process.WithStdoutWriter(stdoutWriter),
 	)
@@ -221,7 +241,12 @@ func (m *pythonMutator) runPythonMutator(ctx context.Context, cacheDir string, r
 		return dyn.InvalidValue, diag.Errorf("failed to load diagnostics: %s", pythonDiagnosticsErr)
 	}
 
-	output, outputDiags := loadOutputFile(rootPath, outputPath)
+	locations, err := loadLocationsFile(locationsPath)
+	if err != nil {
+		return dyn.InvalidValue, diag.Errorf("failed to load locations: %s", err)
+	}
+
+	output, outputDiags := loadOutputFile(opts.bundleRootPath, outputPath, locations)
 	pythonDiagnostics = pythonDiagnostics.Extend(outputDiags)
 
 	// we pass through pythonDiagnostic because it contains warnings
@@ -266,7 +291,21 @@ func writeInputFile(inputPath string, input dyn.Value) error {
 	return os.WriteFile(inputPath, rootConfigJson, 0600)
 }
 
-func loadOutputFile(rootPath string, outputPath string) (dyn.Value, diag.Diagnostics) {
+// loadLocationsFile loads locations.json containing source locations for generated YAML.
+func loadLocationsFile(locationsPath string) (*pythonLocations, error) {
+	locationsFile, err := os.Open(locationsPath)
+	if errors.Is(err, fs.ErrNotExist) {
+		return newPythonLocations(), nil
+	} else if err != nil {
+		return nil, fmt.Errorf("failed to open locations file: %w", err)
+	}
+
+	defer locationsFile.Close()
+
+	return parsePythonLocations(locationsFile)
+}
+
+func loadOutputFile(rootPath string, outputPath string, locations *pythonLocations) (dyn.Value, diag.Diagnostics) {
 	outputFile, err := os.Open(outputPath)
 	if err != nil {
 		return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to open output file: %w", err))
@@ -277,12 +316,12 @@ func loadOutputFile(rootPath string, outputPath string) (dyn.Value, diag.Diagnos
 	// we need absolute path because later parts of pipeline assume all paths are absolute
 	// and this file will be used as location to resolve relative paths.
 	//
-	// virtualPath has to stay in rootPath, because locations outside root path are not allowed:
+	// virtualPath has to stay in bundleRootPath, because locations outside root path are not allowed:
 	//
 	//   Error: path /var/folders/.../pydabs/dist/*.whl is not contained in bundle root path
 	//
 	// for that, we pass virtualPath instead of outputPath as file location
-	virtualPath, err := filepath.Abs(filepath.Join(rootPath, "__generated_by_pydabs__.yml"))
+	virtualPath, err := filepath.Abs(filepath.Join(rootPath, generatedFileName))
 	if err != nil {
 		return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to get absolute path: %w", err))
 	}
@@ -292,7 +331,25 @@ func loadOutputFile(rootPath string, outputPath string) (dyn.Value, diag.Diagnos
 		return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to parse output file: %w", err))
 	}
 
-	return strictNormalize(config.Root{}, generated)
+	// paths are resolved relative to locations of their values, if we change location
+	// we have to update each path, until we simplify that, we don't update locations
+	// for such values, so we don't change how paths are resolved
+	_, err = paths.VisitJobPaths(generated, func(p dyn.Path, kind paths.PathKind, v dyn.Value) (dyn.Value, error) {
+		putPythonLocation(locations, p, v.Location())
+		return v, nil
+	})
+	if err != nil {
+		return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to update locations: %w", err))
+	}
+
+	// generated has dyn.Location as if it comes from generated YAML file
+	// earlier we loaded locations.json with source locations in Python code
+	generatedWithLocations, err := mergePythonLocations(generated, locations)
+	if err != nil {
+		return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to update locations: %w", err))
+	}
+
+	return strictNormalize(config.Root{}, generatedWithLocations)
 }
 
 func strictNormalize(dst any, generated dyn.Value) (dyn.Value, diag.Diagnostics) {

bundle/config/mutator/python/python_mutator_test.go

@@ -6,7 +6,6 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
-	"reflect"
 	"runtime"
 	"testing"
 
@@ -47,6 +46,7 @@ func TestPythonMutator_load(t *testing.T) {
         pydabs:
           enabled: true
           venv_path: .venv
+          load_locations: true
       resources:
         jobs:
           job0:
@@ -65,7 +65,8 @@ func TestPythonMutator_load(t *testing.T) {
 			"experimental": {
 				"pydabs": {
 					"enabled": true,
-					"venv_path": ".venv"
+					"venv_path": ".venv",
+					"load_locations": true
 				}
 			},
 			"resources": {
@@ -80,6 +81,8 @@ func TestPythonMutator_load(t *testing.T) {
 			}
 		}`,
 		`{"severity": "warning", "summary": "job doesn't have any tasks", "location": {"file": "src/examples/file.py", "line": 10, "column": 5}}`,
+		`{"path": "resources.jobs.job0", "file": "src/examples/job0.py", "line": 3, "column": 5}
+		{"path": "resources.jobs.job1", "file": "src/examples/job1.py", "line": 5, "column": 7}`,
 	)
 
 	mutator := PythonMutator(PythonMutatorPhaseLoad)
@@ -97,6 +100,25 @@ func TestPythonMutator_load(t *testing.T) {
 		assert.Equal(t, "job_1", job1.Name)
 	}
 
+	// output of locations.json should be applied to underlying dyn.Value
+	err := b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) {
+		name1, err := dyn.GetByPath(v, dyn.MustPathFromString("resources.jobs.job1.name"))
+		if err != nil {
+			return dyn.InvalidValue, err
+		}
+
+		assert.Equal(t, []dyn.Location{
+			{
+				File:   "src/examples/job1.py",
+				Line:   5,
+				Column: 7,
+			},
+		}, name1.Locations())
+
+		return v, nil
+	})
+	assert.NoError(t, err)
+
 	assert.Equal(t, 1, len(diags))
 	assert.Equal(t, "job doesn't have any tasks", diags[0].Summary)
 	assert.Equal(t, []dyn.Location{
@@ -106,7 +128,6 @@ func TestPythonMutator_load(t *testing.T) {
 			Column: 5,
 		},
 	}, diags[0].Locations)
-
 }
 
 func TestPythonMutator_load_disallowed(t *testing.T) {
@@ -146,7 +167,7 @@ func TestPythonMutator_load_disallowed(t *testing.T) {
 					}
 				}
 			}
-		}`, "")
+		}`, "", "")
 
 	mutator := PythonMutator(PythonMutatorPhaseLoad)
 	diag := bundle.Apply(ctx, b, mutator)
@@ -191,7 +212,7 @@ func TestPythonMutator_init(t *testing.T) {
 					}
 				}
 			}
-		}`, "")
+		}`, "", "")
 
 	mutator := PythonMutator(PythonMutatorPhaseInit)
 	diag := bundle.Apply(ctx, b, mutator)
@@ -252,7 +273,7 @@ func TestPythonMutator_badOutput(t *testing.T) {
 					}
 				}
 			}
-		}`, "")
+		}`, "", "")
 
 	mutator := PythonMutator(PythonMutatorPhaseLoad)
 	diag := bundle.Apply(ctx, b, mutator)
@@ -588,7 +609,7 @@ or activate the environment before running CLI commands:
 	assert.Equal(t, expected, out)
 }
 
-func withProcessStub(t *testing.T, args []string, output string, diagnostics string) context.Context {
+func withProcessStub(t *testing.T, args []string, output string, diagnostics string, locations string) context.Context {
 	ctx := context.Background()
 	ctx, stub := process.WithStub(ctx)
 
@@ -600,32 +621,51 @@ func withProcessStub(t *testing.T, args []string, output string, diagnostics str
 
 	inputPath := filepath.Join(cacheDir, "input.json")
 	outputPath := filepath.Join(cacheDir, "output.json")
+	locationsPath := filepath.Join(cacheDir, "locations.json")
 	diagnosticsPath := filepath.Join(cacheDir, "diagnostics.json")
 
-	args = append(args, "--input", inputPath)
-	args = append(args, "--output", outputPath)
-	args = append(args, "--diagnostics", diagnosticsPath)
-
 	stub.WithCallback(func(actual *exec.Cmd) error {
 		_, err := os.Stat(inputPath)
 		assert.NoError(t, err)
 
-		if reflect.DeepEqual(actual.Args, args) {
-			err := os.WriteFile(outputPath, []byte(output), 0600)
+		actualInputPath := getArg(actual.Args, "--input")
+		actualOutputPath := getArg(actual.Args, "--output")
+		actualDiagnosticsPath := getArg(actual.Args, "--diagnostics")
+		actualLocationsPath := getArg(actual.Args, "--locations")
+
+		require.Equal(t, inputPath, actualInputPath)
+		require.Equal(t, outputPath, actualOutputPath)
+		require.Equal(t, diagnosticsPath, actualDiagnosticsPath)
+
+		// locations is an optional argument
+		if locations != "" {
+			require.Equal(t, locationsPath, actualLocationsPath)
+
+			err = os.WriteFile(locationsPath, []byte(locations), 0600)
+			require.NoError(t, err)
+		}
+
+		err = os.WriteFile(outputPath, []byte(output), 0600)
 		require.NoError(t, err)
 
 		err = os.WriteFile(diagnosticsPath, []byte(diagnostics), 0600)
 		require.NoError(t, err)
 
 		return nil
-		} else {
-			return fmt.Errorf("unexpected command: %v", actual.Args)
-		}
 	})
 
 	return ctx
 }
 
+func getArg(args []string, name string) string {
+	for i := 0; i < len(args); i++ {
+		if args[i] == name {
+			return args[i+1]
+		}
+	}
+	return ""
+}
+
 func loadYaml(name string, content string) *bundle.Bundle {
 	v, diag := config.LoadFromBytes(name, []byte(content))
 

bundle/config/mutator/translate_paths.go

@@ -126,8 +126,34 @@ func (t *translateContext) rewritePath(
 func (t *translateContext) translateNotebookPath(literal, localFullPath, localRelPath, remotePath string) (string, error) {
 	nb, _, err := notebook.DetectWithFS(t.b.SyncRoot, filepath.ToSlash(localRelPath))
 	if errors.Is(err, fs.ErrNotExist) {
+		if filepath.Ext(localFullPath) != notebook.ExtensionNone {
 			return "", fmt.Errorf("notebook %s not found", literal)
 		}
+
+		extensions := []string{
+			notebook.ExtensionPython,
+			notebook.ExtensionR,
+			notebook.ExtensionScala,
+			notebook.ExtensionSql,
+			notebook.ExtensionJupyter,
+		}
+
+		// Check whether a file with a notebook extension already exists. This
+		// way we can provide a more targeted error message.
+		for _, ext := range extensions {
+			literalWithExt := literal + ext
+			localRelPathWithExt := filepath.ToSlash(localRelPath + ext)
+			if _, err := fs.Stat(t.b.SyncRoot, localRelPathWithExt); err == nil {
+				return "", fmt.Errorf(`notebook %s not found. Did you mean %s?
+Local notebook references are expected to contain one of the following
+file extensions: [%s]`, literal, literalWithExt, strings.Join(extensions, ", "))
+			}
+		}
+
+		// Return a generic error message if no matching possible file is found.
+		return "", fmt.Errorf(`notebook %s not found. Local notebook references are expected
+to contain one of the following file extensions: [%s]`, literal, strings.Join(extensions, ", "))
+	}
 	if err != nil {
 		return "", fmt.Errorf("unable to determine if %s is a notebook: %w", localFullPath, err)
 	}

bundle/config/mutator/translate_paths_test.go

@@ -2,6 +2,7 @@ package mutator_test
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -508,6 +509,59 @@ func TestPipelineNotebookDoesNotExistError(t *testing.T) {
 	assert.EqualError(t, diags.Error(), "notebook ./doesnt_exist.py not found")
 }
 
+func TestPipelineNotebookDoesNotExistErrorWithoutExtension(t *testing.T) {
+	for _, ext := range []string{
+		".py",
+		".r",
+		".scala",
+		".sql",
+		".ipynb",
+		"",
+	} {
+		t.Run("case_"+ext, func(t *testing.T) {
+			dir := t.TempDir()
+
+			if ext != "" {
+				touchEmptyFile(t, filepath.Join(dir, "foo"+ext))
+			}
+
+			b := &bundle.Bundle{
+				SyncRootPath: dir,
+				SyncRoot:     vfs.MustNew(dir),
+				Config: config.Root{
+					Resources: config.Resources{
+						Pipelines: map[string]*resources.Pipeline{
+							"pipeline": {
+								PipelineSpec: &pipelines.PipelineSpec{
+									Libraries: []pipelines.PipelineLibrary{
+										{
+											Notebook: &pipelines.NotebookLibrary{
+												Path: "./foo",
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			}
+
+			bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "fake.yml")}})
+			diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
+
+			if ext == "" {
+				assert.EqualError(t, diags.Error(), `notebook ./foo not found. Local notebook references are expected
+to contain one of the following file extensions: [.py, .r, .scala, .sql, .ipynb]`)
+			} else {
+				assert.EqualError(t, diags.Error(), fmt.Sprintf(`notebook ./foo not found. Did you mean ./foo%s?
+Local notebook references are expected to contain one of the following
+file extensions: [.py, .r, .scala, .sql, .ipynb]`, ext))
+			}
+		})
+	}
+}
+
 func TestPipelineFileDoesNotExistError(t *testing.T) {
 	dir := t.TempDir()
 

bundle/config/variable/lookup.go

@@ -1,11 +1,8 @@
-// Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT.
-
 package variable
 
 import (
 	"context"
 	"fmt"
-	"strings"
 
 	"github.com/databricks/databricks-sdk-go"
 )
@@ -34,323 +31,75 @@ type Lookup struct {
 	Warehouse string `json:"warehouse,omitempty"`
 }
 
-func LookupFromMap(m map[string]any) *Lookup {
-	l := &Lookup{}
-	if v, ok := m["alert"]; ok {
-		l.Alert = v.(string)
+type resolver interface {
+	// Resolve resolves the underlying entity's ID.
+	Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error)
+
+	// String returns a human-readable representation of the resolver.
+	String() string
+}
+
+func (l *Lookup) constructResolver() (resolver, error) {
+	var resolvers []resolver
+
+	if l.Alert != "" {
+		resolvers = append(resolvers, resolveAlert{name: l.Alert})
 	}
-	if v, ok := m["cluster_policy"]; ok {
-		l.ClusterPolicy = v.(string)
+	if l.ClusterPolicy != "" {
+		resolvers = append(resolvers, resolveClusterPolicy{name: l.ClusterPolicy})
 	}
-	if v, ok := m["cluster"]; ok {
-		l.Cluster = v.(string)
+	if l.Cluster != "" {
+		resolvers = append(resolvers, resolveCluster{name: l.Cluster})
 	}
-	if v, ok := m["dashboard"]; ok {
-		l.Dashboard = v.(string)
+	if l.Dashboard != "" {
+		resolvers = append(resolvers, resolveDashboard{name: l.Dashboard})
 	}
-	if v, ok := m["instance_pool"]; ok {
-		l.InstancePool = v.(string)
+	if l.InstancePool != "" {
+		resolvers = append(resolvers, resolveInstancePool{name: l.InstancePool})
 	}
-	if v, ok := m["job"]; ok {
-		l.Job = v.(string)
+	if l.Job != "" {
+		resolvers = append(resolvers, resolveJob{name: l.Job})
 	}
-	if v, ok := m["metastore"]; ok {
-		l.Metastore = v.(string)
+	if l.Metastore != "" {
+		resolvers = append(resolvers, resolveMetastore{name: l.Metastore})
 	}
-	if v, ok := m["pipeline"]; ok {
-		l.Pipeline = v.(string)
+	if l.Pipeline != "" {
+		resolvers = append(resolvers, resolvePipeline{name: l.Pipeline})
 	}
-	if v, ok := m["query"]; ok {
-		l.Query = v.(string)
+	if l.Query != "" {
+		resolvers = append(resolvers, resolveQuery{name: l.Query})
 	}
-	if v, ok := m["service_principal"]; ok {
-		l.ServicePrincipal = v.(string)
+	if l.ServicePrincipal != "" {
+		resolvers = append(resolvers, resolveServicePrincipal{name: l.ServicePrincipal})
 	}
-	if v, ok := m["warehouse"]; ok {
-		l.Warehouse = v.(string)
+	if l.Warehouse != "" {
+		resolvers = append(resolvers, resolveWarehouse{name: l.Warehouse})
 	}
 
-	return l
+	switch len(resolvers) {
+	case 0:
+		return nil, fmt.Errorf("no valid lookup fields provided")
+	case 1:
+		return resolvers[0], nil
+	default:
+		return nil, fmt.Errorf("exactly one lookup field must be provided")
+	}
 }
 
 func (l *Lookup) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
-	if err := l.validate(); err != nil {
+	r, err := l.constructResolver()
+	if err != nil {
 		return "", err
 	}
 
-	r := allResolvers()
-	if l.Alert != "" {
-		return r.Alert(ctx, w, l.Alert)
-	}
-	if l.ClusterPolicy != "" {
-		return r.ClusterPolicy(ctx, w, l.ClusterPolicy)
-	}
-	if l.Cluster != "" {
-		return r.Cluster(ctx, w, l.Cluster)
-	}
-	if l.Dashboard != "" {
-		return r.Dashboard(ctx, w, l.Dashboard)
-	}
-	if l.InstancePool != "" {
-		return r.InstancePool(ctx, w, l.InstancePool)
-	}
-	if l.Job != "" {
-		return r.Job(ctx, w, l.Job)
-	}
-	if l.Metastore != "" {
-		return r.Metastore(ctx, w, l.Metastore)
-	}
-	if l.Pipeline != "" {
-		return r.Pipeline(ctx, w, l.Pipeline)
-	}
-	if l.Query != "" {
-		return r.Query(ctx, w, l.Query)
-	}
-	if l.ServicePrincipal != "" {
-		return r.ServicePrincipal(ctx, w, l.ServicePrincipal)
-	}
-	if l.Warehouse != "" {
-		return r.Warehouse(ctx, w, l.Warehouse)
-	}
-
-	return "", fmt.Errorf("no valid lookup fields provided")
+	return r.Resolve(ctx, w)
 }
 
 func (l *Lookup) String() string {
-	if l.Alert != "" {
-		return fmt.Sprintf("alert: %s", l.Alert)
-	}
-	if l.ClusterPolicy != "" {
-		return fmt.Sprintf("cluster-policy: %s", l.ClusterPolicy)
-	}
-	if l.Cluster != "" {
-		return fmt.Sprintf("cluster: %s", l.Cluster)
-	}
-	if l.Dashboard != "" {
-		return fmt.Sprintf("dashboard: %s", l.Dashboard)
-	}
-	if l.InstancePool != "" {
-		return fmt.Sprintf("instance-pool: %s", l.InstancePool)
-	}
-	if l.Job != "" {
-		return fmt.Sprintf("job: %s", l.Job)
-	}
-	if l.Metastore != "" {
-		return fmt.Sprintf("metastore: %s", l.Metastore)
-	}
-	if l.Pipeline != "" {
-		return fmt.Sprintf("pipeline: %s", l.Pipeline)
-	}
-	if l.Query != "" {
-		return fmt.Sprintf("query: %s", l.Query)
-	}
-	if l.ServicePrincipal != "" {
-		return fmt.Sprintf("service-principal: %s", l.ServicePrincipal)
-	}
-	if l.Warehouse != "" {
-		return fmt.Sprintf("warehouse: %s", l.Warehouse)
-	}
-
+	r, _ := l.constructResolver()
+	if r == nil {
 		return ""
-}
-
-func (l *Lookup) validate() error {
-	// Validate that only one field is set
-	count := 0
-	if l.Alert != "" {
-		count++
-	}
-	if l.ClusterPolicy != "" {
-		count++
-	}
-	if l.Cluster != "" {
-		count++
-	}
-	if l.Dashboard != "" {
-		count++
-	}
-	if l.InstancePool != "" {
-		count++
-	}
-	if l.Job != "" {
-		count++
-	}
-	if l.Metastore != "" {
-		count++
-	}
-	if l.Pipeline != "" {
-		count++
-	}
-	if l.Query != "" {
-		count++
-	}
-	if l.ServicePrincipal != "" {
-		count++
-	}
-	if l.Warehouse != "" {
-		count++
-	}
-
-	if count != 1 {
-		return fmt.Errorf("exactly one lookup field must be provided")
-	}
-
-	if strings.Contains(l.String(), "${var") {
-		return fmt.Errorf("lookup fields cannot contain variable references")
-	}
-
-	return nil
-}
-
-type resolverFunc func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error)
-type resolvers struct {
-	Alert            resolverFunc
-	ClusterPolicy    resolverFunc
-	Cluster          resolverFunc
-	Dashboard        resolverFunc
-	InstancePool     resolverFunc
-	Job              resolverFunc
-	Metastore        resolverFunc
-	Pipeline         resolverFunc
-	Query            resolverFunc
-	ServicePrincipal resolverFunc
-	Warehouse        resolverFunc
-}
-
-func allResolvers() *resolvers {
-	r := &resolvers{}
-	r.Alert = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["Alert"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.Alerts.GetByDisplayName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.Id), nil
-	}
-	r.ClusterPolicy = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["ClusterPolicy"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.ClusterPolicies.GetByName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.PolicyId), nil
-	}
-	r.Cluster = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["Cluster"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.Clusters.GetByClusterName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.ClusterId), nil
-	}
-	r.Dashboard = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["Dashboard"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.Dashboards.GetByName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.Id), nil
-	}
-	r.InstancePool = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["InstancePool"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.InstancePools.GetByInstancePoolName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.InstancePoolId), nil
-	}
-	r.Job = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["Job"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.Jobs.GetBySettingsName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.JobId), nil
-	}
-	r.Metastore = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["Metastore"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.Metastores.GetByName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.MetastoreId), nil
-	}
-	r.Pipeline = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["Pipeline"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.Pipelines.GetByName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.PipelineId), nil
-	}
-	r.Query = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["Query"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.Queries.GetByDisplayName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.Id), nil
-	}
-	r.ServicePrincipal = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["ServicePrincipal"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.ServicePrincipals.GetByDisplayName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.ApplicationId), nil
-	}
-	r.Warehouse = func(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
-		fn, ok := lookupOverrides["Warehouse"]
-		if ok {
-			return fn(ctx, w, name)
-		}
-		entity, err := w.Warehouses.GetByName(ctx, name)
-		if err != nil {
-			return "", err
-		}
-
-		return fmt.Sprint(entity.Id), nil
-	}
-
-	return r
+	}
+
+	return r.String()
 }

bundle/config/variable/lookup_test.go

@@ -0,0 +1,60 @@
+package variable
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLookup_Coverage(t *testing.T) {
+	var lookup Lookup
+	val := reflect.ValueOf(lookup)
+	typ := val.Type()
+
+	for i := 0; i < val.NumField(); i++ {
+		field := val.Field(i)
+		if field.Kind() != reflect.String {
+			t.Fatalf("Field %s is not a string", typ.Field(i).Name)
+		}
+
+		fieldType := typ.Field(i)
+		t.Run(fieldType.Name, func(t *testing.T) {
+			// Use a fresh instance of the struct in each test
+			var lookup Lookup
+
+			// Set the field to a non-empty string
+			reflect.ValueOf(&lookup).Elem().Field(i).SetString("value")
+
+			// Test the [String] function
+			assert.NotEmpty(t, lookup.String())
+		})
+	}
+}
+
+func TestLookup_Empty(t *testing.T) {
+	var lookup Lookup
+
+	// Resolve returns an error when no fields are provided
+	_, err := lookup.Resolve(context.Background(), nil)
+	assert.ErrorContains(t, err, "no valid lookup fields provided")
+
+	// No string representation for an invalid lookup
+	assert.Empty(t, lookup.String())
+
+}
+
+func TestLookup_Multiple(t *testing.T) {
+	lookup := Lookup{
+		Alert: "alert",
+		Query: "query",
+	}
+
+	// Resolve returns an error when multiple fields are provided
+	_, err := lookup.Resolve(context.Background(), nil)
+	assert.ErrorContains(t, err, "exactly one lookup field must be provided")
+
+	// No string representation for an invalid lookup
+	assert.Empty(t, lookup.String())
+}

bundle/config/variable/resolve_alert.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveAlert struct {
+	name string
+}
+
+func (l resolveAlert) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.Alerts.GetByDisplayName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.Id), nil
+}
+
+func (l resolveAlert) String() string {
+	return fmt.Sprintf("alert: %s", l.name)
+}

bundle/config/variable/resolve_alert_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/sql"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveAlert_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockAlertsAPI()
+	api.EXPECT().
+		GetByDisplayName(mock.Anything, "alert").
+		Return(&sql.ListAlertsResponseAlert{
+			Id: "1234",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveAlert{name: "alert"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "1234", result)
+}
+
+func TestResolveAlert_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockAlertsAPI()
+	api.EXPECT().
+		GetByDisplayName(mock.Anything, "alert").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveAlert{name: "alert"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveAlert_String(t *testing.T) {
+	l := resolveAlert{name: "name"}
+	assert.Equal(t, "alert: name", l.String())
+}

bundle/config/variable/resolve_cluster.go

@@ -8,13 +8,13 @@ import (
 	"github.com/databricks/databricks-sdk-go/service/compute"
 )
 
-var lookupOverrides = map[string]resolverFunc{
-	"Cluster": resolveCluster,
+type resolveCluster struct {
+	name string
 }
 
 // We added a custom resolver for the cluster to add filtering for the cluster source when we list all clusters.
 // Without the filtering listing could take a very long time (5-10 mins) which leads to lookup timeouts.
-func resolveCluster(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) {
+func (l resolveCluster) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
 	result, err := w.Clusters.ListAll(ctx, compute.ListClustersRequest{
 		FilterBy: &compute.ListClustersFilterBy{
 			ClusterSources: []compute.ClusterSource{compute.ClusterSourceApi, compute.ClusterSourceUi},
@@ -30,6 +30,8 @@ func resolveCluster(ctx context.Context, w *databricks.WorkspaceClient, name str
 		key := v.ClusterName
 		tmp[key] = append(tmp[key], v)
 	}
+
+	name := l.name
 	alternatives, ok := tmp[name]
 	if !ok || len(alternatives) == 0 {
 		return "", fmt.Errorf("cluster named '%s' does not exist", name)
@@ -39,3 +41,7 @@ func resolveCluster(ctx context.Context, w *databricks.WorkspaceClient, name str
 	}
 	return alternatives[0].ClusterId, nil
 }
+
+func (l resolveCluster) String() string {
+	return fmt.Sprintf("cluster: %s", l.name)
+}

bundle/config/variable/resolve_cluster_policy.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveClusterPolicy struct {
+	name string
+}
+
+func (l resolveClusterPolicy) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.ClusterPolicies.GetByName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.PolicyId), nil
+}
+
+func (l resolveClusterPolicy) String() string {
+	return fmt.Sprintf("cluster-policy: %s", l.name)
+}

bundle/config/variable/resolve_cluster_policy_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveClusterPolicy_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockClusterPoliciesAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "policy").
+		Return(&compute.Policy{
+			PolicyId: "1234",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveClusterPolicy{name: "policy"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "1234", result)
+}
+
+func TestResolveClusterPolicy_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockClusterPoliciesAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "policy").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveClusterPolicy{name: "policy"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveClusterPolicy_String(t *testing.T) {
+	l := resolveClusterPolicy{name: "name"}
+	assert.Equal(t, "cluster-policy: name", l.String())
+}

bundle/config/variable/resolve_cluster_test.go

@@ -0,0 +1,50 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveCluster_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockClustersAPI()
+	api.EXPECT().
+		ListAll(mock.Anything, mock.Anything).
+		Return([]compute.ClusterDetails{
+			{ClusterId: "1234", ClusterName: "cluster1"},
+			{ClusterId: "2345", ClusterName: "cluster2"},
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveCluster{name: "cluster2"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "2345", result)
+}
+
+func TestResolveCluster_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockClustersAPI()
+	api.EXPECT().
+		ListAll(mock.Anything, mock.Anything).
+		Return([]compute.ClusterDetails{}, nil)
+
+	ctx := context.Background()
+	l := resolveCluster{name: "cluster"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "cluster named 'cluster' does not exist")
+}
+
+func TestResolveCluster_String(t *testing.T) {
+	l := resolveCluster{name: "name"}
+	assert.Equal(t, "cluster: name", l.String())
+}

bundle/config/variable/resolve_dashboard.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveDashboard struct {
+	name string
+}
+
+func (l resolveDashboard) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.Dashboards.GetByName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.Id), nil
+}
+
+func (l resolveDashboard) String() string {
+	return fmt.Sprintf("dashboard: %s", l.name)
+}

bundle/config/variable/resolve_dashboard_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/sql"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveDashboard_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockDashboardsAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "dashboard").
+		Return(&sql.Dashboard{
+			Id: "1234",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveDashboard{name: "dashboard"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "1234", result)
+}
+
+func TestResolveDashboard_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockDashboardsAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "dashboard").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveDashboard{name: "dashboard"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveDashboard_String(t *testing.T) {
+	l := resolveDashboard{name: "name"}
+	assert.Equal(t, "dashboard: name", l.String())
+}

bundle/config/variable/resolve_instance_pool.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveInstancePool struct {
+	name string
+}
+
+func (l resolveInstancePool) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.InstancePools.GetByInstancePoolName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.InstancePoolId), nil
+}
+
+func (l resolveInstancePool) String() string {
+	return fmt.Sprintf("instance-pool: %s", l.name)
+}

bundle/config/variable/resolve_instance_pool_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveInstancePool_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockInstancePoolsAPI()
+	api.EXPECT().
+		GetByInstancePoolName(mock.Anything, "instance_pool").
+		Return(&compute.InstancePoolAndStats{
+			InstancePoolId: "5678",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveInstancePool{name: "instance_pool"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "5678", result)
+}
+
+func TestResolveInstancePool_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockInstancePoolsAPI()
+	api.EXPECT().
+		GetByInstancePoolName(mock.Anything, "instance_pool").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveInstancePool{name: "instance_pool"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveInstancePool_String(t *testing.T) {
+	l := resolveInstancePool{name: "name"}
+	assert.Equal(t, "instance-pool: name", l.String())
+}

bundle/config/variable/resolve_job.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveJob struct {
+	name string
+}
+
+func (l resolveJob) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.Jobs.GetBySettingsName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.JobId), nil
+}
+
+func (l resolveJob) String() string {
+	return fmt.Sprintf("job: %s", l.name)
+}

bundle/config/variable/resolve_job_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveJob_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockJobsAPI()
+	api.EXPECT().
+		GetBySettingsName(mock.Anything, "job").
+		Return(&jobs.BaseJob{
+			JobId: 5678,
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveJob{name: "job"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "5678", result)
+}
+
+func TestResolveJob_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockJobsAPI()
+	api.EXPECT().
+		GetBySettingsName(mock.Anything, "job").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveJob{name: "job"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveJob_String(t *testing.T) {
+	l := resolveJob{name: "name"}
+	assert.Equal(t, "job: name", l.String())
+}

bundle/config/variable/resolve_metastore.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveMetastore struct {
+	name string
+}
+
+func (l resolveMetastore) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.Metastores.GetByName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.MetastoreId), nil
+}
+
+func (l resolveMetastore) String() string {
+	return fmt.Sprintf("metastore: %s", l.name)
+}

bundle/config/variable/resolve_metastore_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/catalog"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveMetastore_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockMetastoresAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "metastore").
+		Return(&catalog.MetastoreInfo{
+			MetastoreId: "abcd",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveMetastore{name: "metastore"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "abcd", result)
+}
+
+func TestResolveMetastore_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockMetastoresAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "metastore").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveMetastore{name: "metastore"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveMetastore_String(t *testing.T) {
+	l := resolveMetastore{name: "name"}
+	assert.Equal(t, "metastore: name", l.String())
+}

bundle/config/variable/resolve_pipeline.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolvePipeline struct {
+	name string
+}
+
+func (l resolvePipeline) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.Pipelines.GetByName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.PipelineId), nil
+}
+
+func (l resolvePipeline) String() string {
+	return fmt.Sprintf("pipeline: %s", l.name)
+}

bundle/config/variable/resolve_pipeline_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/pipelines"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolvePipeline_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockPipelinesAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "pipeline").
+		Return(&pipelines.PipelineStateInfo{
+			PipelineId: "abcd",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolvePipeline{name: "pipeline"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "abcd", result)
+}
+
+func TestResolvePipeline_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockPipelinesAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "pipeline").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolvePipeline{name: "pipeline"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolvePipeline_String(t *testing.T) {
+	l := resolvePipeline{name: "name"}
+	assert.Equal(t, "pipeline: name", l.String())
+}

bundle/config/variable/resolve_query.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveQuery struct {
+	name string
+}
+
+func (l resolveQuery) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.Queries.GetByDisplayName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.Id), nil
+}
+
+func (l resolveQuery) String() string {
+	return fmt.Sprintf("query: %s", l.name)
+}

bundle/config/variable/resolve_query_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/sql"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveQuery_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockQueriesAPI()
+	api.EXPECT().
+		GetByDisplayName(mock.Anything, "query").
+		Return(&sql.ListQueryObjectsResponseQuery{
+			Id: "1234",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveQuery{name: "query"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "1234", result)
+}
+
+func TestResolveQuery_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockQueriesAPI()
+	api.EXPECT().
+		GetByDisplayName(mock.Anything, "query").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveQuery{name: "query"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveQuery_String(t *testing.T) {
+	l := resolveQuery{name: "name"}
+	assert.Equal(t, "query: name", l.String())
+}

bundle/config/variable/resolve_service_principal.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveServicePrincipal struct {
+	name string
+}
+
+func (l resolveServicePrincipal) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.ServicePrincipals.GetByDisplayName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.ApplicationId), nil
+}
+
+func (l resolveServicePrincipal) String() string {
+	return fmt.Sprintf("service-principal: %s", l.name)
+}

bundle/config/variable/resolve_service_principal_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/iam"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveServicePrincipal_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockServicePrincipalsAPI()
+	api.EXPECT().
+		GetByDisplayName(mock.Anything, "service-principal").
+		Return(&iam.ServicePrincipal{
+			ApplicationId: "5678",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveServicePrincipal{name: "service-principal"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "5678", result)
+}
+
+func TestResolveServicePrincipal_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockServicePrincipalsAPI()
+	api.EXPECT().
+		GetByDisplayName(mock.Anything, "service-principal").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveServicePrincipal{name: "service-principal"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveServicePrincipal_String(t *testing.T) {
+	l := resolveServicePrincipal{name: "name"}
+	assert.Equal(t, "service-principal: name", l.String())
+}

bundle/config/variable/resolve_warehouse.go

@@ -0,0 +1,24 @@
+package variable
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+type resolveWarehouse struct {
+	name string
+}
+
+func (l resolveWarehouse) Resolve(ctx context.Context, w *databricks.WorkspaceClient) (string, error) {
+	entity, err := w.Warehouses.GetByName(ctx, l.name)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprint(entity.Id), nil
+}
+
+func (l resolveWarehouse) String() string {
+	return fmt.Sprintf("warehouse: %s", l.name)
+}

bundle/config/variable/resolve_warehouse_test.go

@@ -0,0 +1,49 @@
+package variable
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/sql"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResolveWarehouse_ResolveSuccess(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockWarehousesAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "warehouse").
+		Return(&sql.EndpointInfo{
+			Id: "abcd",
+		}, nil)
+
+	ctx := context.Background()
+	l := resolveWarehouse{name: "warehouse"}
+	result, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.NoError(t, err)
+	assert.Equal(t, "abcd", result)
+}
+
+func TestResolveWarehouse_ResolveNotFound(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+
+	api := m.GetMockWarehousesAPI()
+	api.EXPECT().
+		GetByName(mock.Anything, "warehouse").
+		Return(nil, &apierr.APIError{StatusCode: 404})
+
+	ctx := context.Background()
+	l := resolveWarehouse{name: "warehouse"}
+	_, err := l.Resolve(ctx, m.WorkspaceClient)
+	require.ErrorIs(t, err, apierr.ErrNotFound)
+}
+
+func TestResolveWarehouse_String(t *testing.T) {
+	l := resolveWarehouse{name: "name"}
+	assert.Equal(t, "warehouse: name", l.String())
+}

internal/init_test.go

@@ -97,7 +97,7 @@ func TestAccBundleInitOnMlopsStacks(t *testing.T) {
 	require.NoError(t, err)
 	job, err := w.Jobs.GetByJobId(context.Background(), batchJobId)
 	assert.NoError(t, err)
-	assert.Equal(t, fmt.Sprintf("dev-%s-batch-inference-job", projectName), job.Settings.Name)
+	assert.Contains(t, job.Settings.Name, fmt.Sprintf("dev-%s-batch-inference-job", projectName))
 }
 
 func TestAccBundleInitHelpers(t *testing.T) {