Add support for variables in bundle config

bundle/config/interpolation/expand_variables.go

@@ -0,0 +1,26 @@
+package interpolation
+
+import (
+	"fmt"
+	"strings"
+)
+
+const VariableReferencePrefix = "var."
+
+func isVariableReference(s string) bool {
+	if !strings.HasPrefix(s, VariableReferencePrefix) || strings.Count(s, ".") != 1 {
+		return false
+	}
+	name := strings.TrimPrefix(s, VariableReferencePrefix)
+	return len(name) > 0
+}
+
+// Expands variable references of the form `var.foo` to `variables.foo.value`
+// Errors out if input string is not in the correct format
+func expandVariable(s string) (string, error) {
+	if !isVariableReference(s) {
+		return "", fmt.Errorf("%s is not a valid variable reference", s)
+	}
+	name := strings.TrimPrefix(s, VariableReferencePrefix)
+	return strings.Join([]string{"variables", name, "value"}, "."), nil
+}

bundle/config/interpolation/expand_variables_test.go

@@ -0,0 +1,38 @@
+package interpolation
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestVariableInterpolationIsVariableReference(t *testing.T) {
+	assert.True(t, isVariableReference("var.foo"))
+	assert.True(t, isVariableReference("var.bar"))
+	assert.True(t, isVariableReference("var.var"))
+	assert.True(t, isVariableReference("var.a-b-c"))
+
+	assert.False(t, isVariableReference("var.foo.bar"))
+	assert.False(t, isVariableReference("vars.bar"))
+	assert.False(t, isVariableReference("var"))
+	assert.False(t, isVariableReference("var."))
+}
+
+func TestVariableInterpolationExpandVariable(t *testing.T) {
+	ans, err := expandVariable("var.foo")
+	assert.NoError(t, err)
+	assert.Equal(t, "variables.foo.value", ans)
+
+	ans, err = expandVariable("var.foo-BAR")
+	assert.NoError(t, err)
+	assert.Equal(t, "variables.foo-BAR.value", ans)
+
+	_, err = expandVariable("var.foo.bar")
+	assert.ErrorContains(t, err, "var.foo.bar is not a valid variable reference")
+
+	_, err = expandVariable("var.")
+	assert.ErrorContains(t, err, "var. is not a valid variable reference")
+
+	_, err = expandVariable("vars.foo")
+	assert.ErrorContains(t, err, "vars.foo is not a valid variable reference")
+}

bundle/config/interpolation/interpolation.go

@@ -174,7 +174,7 @@ func (a *accumulator) Resolve(path string, seenPaths []string, fns ...LookupFunc
 	// fetch the string node to resolve
 	field, ok := a.strings[path]
 	if !ok {
-		return fmt.Errorf("could not find string field with path %s", path)
+		return fmt.Errorf("could not resolve reference %s", path)
 	}
 
 	// return early if the string field has no variables to interpolate
@@ -185,6 +185,15 @@ func (a *accumulator) Resolve(path string, seenPaths []string, fns ...LookupFunc
 
 	// resolve all variables refered in the root string field
 	for _, childFieldPath := range field.dependsOn() {
+		// expand child path if it's an alias for a variable
+		if isVariableReference(childFieldPath) {
+			var err error
+			childFieldPath, err = expandVariable(childFieldPath)
+			if err != nil {
+				return err
+			}
+		}
+
 		// error if there is a loop in variable interpolation
 		if slices.Contains(seenPaths, childFieldPath) {
 			return fmt.Errorf("cycle detected in field resolution: %s", strings.Join(append(seenPaths, childFieldPath), " -> "))

bundle/config/interpolation/interpolation_test.go

@@ -3,6 +3,7 @@ package interpolation
 import (
 	"testing"
 
+	"github.com/databricks/bricks/bundle/config"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -31,7 +32,7 @@ type foo struct {
 func expand(v any) error {
 	a := accumulator{}
 	a.start(v)
-	return a.expand(DefaultLookup)
+	return a.expand(IncludeVariableLookups(), DefaultLookup)
 }
 
 func TestInterpolationVariables(t *testing.T) {
@@ -125,3 +126,70 @@ func TestInterpolationVariableLoopError(t *testing.T) {
 	err := expand(&f)
 	assert.ErrorContains(t, err, "cycle detected in field resolution: b -> c -> d -> b")
 }
+
+func TestInterpolationForVariables(t *testing.T) {
+	foo := "abc"
+	bar := "${var.foo} def"
+	apple := "${var.foo} ${var.bar}"
+	config := config.Root{
+		Variables: map[string]*config.Variable{
+			"foo": {
+				Value: &foo,
+			},
+			"bar": {
+				Value: &bar,
+			},
+			"apple": {
+				Value: &apple,
+			},
+		},
+		Bundle: config.Bundle{
+			Name: "${var.apple} ${var.foo}",
+		},
+	}
+
+	err := expand(&config)
+	assert.NoError(t, err)
+	assert.Equal(t, "abc", *(config.Variables["foo"].Value))
+	assert.Equal(t, "abc def", *(config.Variables["bar"].Value))
+	assert.Equal(t, "abc abc def", *(config.Variables["apple"].Value))
+	assert.Equal(t, "abc abc def abc", config.Bundle.Name)
+}
+
+func TestInterpolationLoopForVariables(t *testing.T) {
+	foo := "${var.bar}"
+	bar := "${var.foo}"
+	config := config.Root{
+		Variables: map[string]*config.Variable{
+			"foo": {
+				Value: &foo,
+			},
+			"bar": {
+				Value: &bar,
+			},
+		},
+		Bundle: config.Bundle{
+			Name: "${var.foo}",
+		},
+	}
+
+	err := expand(&config)
+	assert.ErrorContains(t, err, "cycle detected in field resolution: bundle.name -> variables.foo.value -> variables.bar.value -> variables.foo.value")
+}
+
+func TestInterpolationInvalidVariableReference(t *testing.T) {
+	foo := "abc"
+	config := config.Root{
+		Variables: map[string]*config.Variable{
+			"foo": {
+				Value: &foo,
+			},
+		},
+		Bundle: config.Bundle{
+			Name: "${vars.foo}",
+		},
+	}
+
+	err := expand(&config)
+	assert.ErrorContains(t, err, "could not resolve reference vars.foo")
+}

bundle/config/interpolation/lookup.go

@@ -50,3 +50,18 @@ func IncludeLookupsInPath(include ...string) LookupFunction {
 		return DefaultLookup(path, lookup)
 	}
 }
+
+// IncludeVariableLookups is a lookup function that resolves variable lookups to
+// their value. ${var.foo} is resolved to ${variables.foo.value}
+func IncludeVariableLookups() LookupFunction {
+	return func(path string, lookup map[string]string) (string, error) {
+		if !isVariableReference(path) {
+			return "", ErrSkipInterpolation
+		}
+		expandedPath, err := expandVariable(path)
+		if err != nil {
+			return "", err
+		}
+		return DefaultLookup(expandedPath, lookup)
+	}
+}

bundle/config/interpolation/lookup_test.go

@@ -79,3 +79,24 @@ func TestIncludePathMultiple(t *testing.T) {
 	assert.Equal(t, "1", tmp.C["ax"])
 	assert.Equal(t, "2", tmp.C["bx"])
 }
+
+func TestIncludeVariableLookups(t *testing.T) {
+	f := IncludeVariableLookups()
+
+	_, err := f("abc.def", nil)
+	assert.ErrorIs(t, err, ErrSkipInterpolation)
+
+	_, err = f("var.", nil)
+	assert.ErrorIs(t, err, ErrSkipInterpolation)
+
+	_, err = f("variables.foo.value", map[string]string{
+		"variables.foo.value": "abc",
+	})
+	assert.ErrorIs(t, err, ErrSkipInterpolation)
+
+	ans, err := f("var.foo", map[string]string{
+		"variables.foo.value": "abc",
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, "abc", ans)
+}

bundle/config/mutator/set_variables.go

@@ -0,0 +1,60 @@
+package mutator
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/databricks/bricks/bundle"
+	"github.com/databricks/bricks/bundle/config"
+)
+
+const bundleVarPrefix = "BUNDLE_VAR_"
+
+type setVariables struct{}
+
+func SetVariables() bundle.Mutator {
+	return &setVariables{}
+}
+
+func (m *setVariables) Name() string {
+	return "SetVariables"
+}
+
+func setVariable(v *config.Variable, name string) error {
+	// case: variable already has value initialized, so skip
+	if v.HasValue() {
+		return nil
+	}
+
+	// case: read and set variable value from process environment
+	if val, ok := os.LookupEnv(bundleVarPrefix + name); ok {
+		err := v.Set(val)
+		if err != nil {
+			return fmt.Errorf("failed to assign %s to %s with error: %w", val, name, err)
+		}
+		return nil
+	}
+
+	// case: Set the varliable to it's default value
+	if v.HasDefault() {
+		err := v.Set(*v.Default)
+		if err != nil {
+			return fmt.Errorf("failed to assign %s to %s with error: %w", *v.Default, name, err)
+		}
+		return nil
+	}
+
+	// We should have had a value to set for the variable at this point.
+	return fmt.Errorf(`no value assigned to required variable %s. Assignment can be done through the "--var" flag or by setting the %s environment variable`, name, bundleVarPrefix+name)
+}
+
+func (m *setVariables) Apply(ctx context.Context, b *bundle.Bundle) ([]bundle.Mutator, error) {
+	for name, variable := range b.Config.Variables {
+		err := setVariable(variable, name)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return nil, nil
+}

bundle/config/root.go

@@ -1,8 +1,10 @@
 package config
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/ghodss/yaml"
 	"github.com/imdario/mergo"
@@ -16,6 +18,10 @@ type Root struct {
 	// It is set when loading `bundle.yml`.
 	Path string `json:"-" bundle:"readonly"`
 
+	// Contains user defined variables, either defined in the bundle config or
+	// loaded through env vars
+	Variables map[string]*Variable `json:"variables"`
+
 	// Bundle contains details about this bundle, such as its name,
 	// version of the spec (TODO), default cluster, default warehouse, etc.
 	Bundle Bundle `json:"bundle"`
@@ -79,6 +85,29 @@ func (r *Root) SetConfigFilePath(path string) {
 	}
 }
 
+// Initializes variables using values passed from the command line flag
+// Input has to be a string of the form `foo=bar`. In this case the variable with
+// name `foo` is assigned the value `bar`
+func (r *Root) InitializeVariables(vars []string) error {
+	for _, variable := range vars {
+		varComponents := strings.Split(variable, "=")
+		if len(varComponents) != 2 {
+			return fmt.Errorf("variable assignment %s has unexpected format", variable)
+		}
+		varName := varComponents[0]
+		varValue := varComponents[1]
+
+		if _, ok := r.Variables[varName]; !ok {
+			return fmt.Errorf("variable %s has not been defined", varName)
+		}
+		err := r.Variables[varName].Set(varValue)
+		if err != nil {
+			return fmt.Errorf("failed to assign %s to %s: %s", varValue, varName, err)
+		}
+	}
+	return nil
+}
+
 func (r *Root) Load(path string) error {
 	raw, err := os.ReadFile(path)
 	if err != nil {

bundle/config/root_test.go

@@ -93,3 +93,49 @@ func TestDuplicateIdOnMergeReturnsError(t *testing.T) {
 	err = root.Merge(other)
 	assert.ErrorContains(t, err, "multiple resources named foo (job at ./testdata/duplicate_resource_name_in_subconfiguration/bundle.yml, pipeline at ./testdata/duplicate_resource_name_in_subconfiguration/resources.yml)")
 }
+
+func TestInitializeVariables(t *testing.T) {
+	fooDefault := "abc"
+	root := &Root{
+		Variables: map[string]*Variable{
+			"foo": {
+				Default: &fooDefault,
+				Type:    "string",
+			},
+			"bar": {
+				Type: "string",
+			},
+		},
+	}
+
+	err := root.InitializeVariables([]string{"foo=123", "bar=456"})
+	assert.NoError(t, err)
+	assert.Equal(t, "123", *(root.Variables["foo"].Value))
+	assert.Equal(t, "456", *(root.Variables["bar"].Value))
+}
+
+func TestInitializeVariablesInvalidFormat(t *testing.T) {
+	root := &Root{
+		Variables: map[string]*Variable{
+			"foo": {
+				Type: "string",
+			},
+		},
+	}
+
+	err := root.InitializeVariables([]string{"foo=123=567"})
+	assert.ErrorContains(t, err, "variable assignment foo=123=567 has unexpected format")
+}
+
+func TestInitializeVariablesUndefinedVariables(t *testing.T) {
+	root := &Root{
+		Variables: map[string]*Variable{
+			"foo": {
+				Type: "string",
+			},
+		},
+	}
+
+	err := root.InitializeVariables([]string{"bar=567"})
+	assert.ErrorContains(t, err, "variable bar has not been defined")
+}