Add support for variables in bundle config (#359)

## Changes
This PR now allows you to define variables in the bundle config and set
them in three ways
1. command line args
2. process environment variable
3. in the bundle config itself

## Tests
manually, unit, and black box tests

---------

Co-authored-by: Miles Yucht <miles@databricks.com>
This commit is contained in:
shreyas-goenka 2023-05-15 11:34:05 +02:00 committed by GitHub
parent c98b8dd583
commit c5e940f664
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
18 changed files with 474 additions and 12 deletions

View File

@ -10,6 +10,7 @@ import (
"strings" "strings"
"github.com/databricks/bricks/bundle" "github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/config/variable"
"golang.org/x/exp/maps" "golang.org/x/exp/maps"
"golang.org/x/exp/slices" "golang.org/x/exp/slices"
) )
@ -127,6 +128,13 @@ func (a *accumulator) walk(scope []string, rv reflect.Value, s setter) {
case reflect.String: case reflect.String:
path := strings.Join(scope, Delimiter) path := strings.Join(scope, Delimiter)
a.strings[path] = newStringField(path, anyGetter{rv}, s) a.strings[path] = newStringField(path, anyGetter{rv}, s)
// register alias for variable value. `var.foo` would be the alias for
// `variables.foo.value`
if len(scope) == 3 && scope[0] == "variables" && scope[2] == "value" {
aliasPath := strings.Join([]string{variable.VariableReferencePrefix, scope[1]}, Delimiter)
a.strings[aliasPath] = a.strings[path]
}
case reflect.Struct: case reflect.Struct:
a.walkStruct(scope, rv) a.walkStruct(scope, rv)
case reflect.Map: case reflect.Map:
@ -174,7 +182,7 @@ func (a *accumulator) Resolve(path string, seenPaths []string, fns ...LookupFunc
// fetch the string node to resolve // fetch the string node to resolve
field, ok := a.strings[path] field, ok := a.strings[path]
if !ok { 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 // return early if the string field has no variables to interpolate

View File

@ -3,6 +3,8 @@ package interpolation
import ( import (
"testing" "testing"
"github.com/databricks/bricks/bundle/config"
"github.com/databricks/bricks/bundle/config/variable"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
@ -125,3 +127,70 @@ func TestInterpolationVariableLoopError(t *testing.T) {
err := expand(&f) err := expand(&f)
assert.ErrorContains(t, err, "cycle detected in field resolution: b -> c -> d -> b") 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]*variable.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]*variable.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 -> var.foo -> var.bar -> var.foo")
}
func TestInterpolationInvalidVariableReference(t *testing.T) {
foo := "abc"
config := config.Root{
Variables: map[string]*variable.Variable{
"foo": {
Value: &foo,
},
},
Bundle: config.Bundle{
Name: "${vars.foo}",
},
}
err := expand(&config)
assert.ErrorContains(t, err, "could not resolve reference vars.foo")
}

View File

@ -0,0 +1,63 @@
package mutator
import (
"context"
"fmt"
"os"
"github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/config/variable"
)
const bundleVarPrefix = "BUNDLE_VAR_"
type setVariables struct{}
func SetVariables() bundle.Mutator {
return &setVariables{}
}
func (m *setVariables) Name() string {
return "SetVariables"
}
func setVariable(v *variable.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
envVarName := bundleVarPrefix + name
if val, ok := os.LookupEnv(envVarName); ok {
err := v.Set(val)
if err != nil {
return fmt.Errorf(`failed to assign value "%s" to variable %s from environment variable %s with error: %w`, val, name, envVarName, err)
}
return nil
}
// case: Set the variable to its default value
if v.HasDefault() {
err := v.Set(*v.Default)
if err != nil {
return fmt.Errorf(`failed to assign default value from config "%s" to variable %s with error: %w`, *v.Default, name, err)
}
return nil
}
// We should have had a value to set for the variable at this point.
// TODO: use cmdio to request values for unassigned variables if current
// terminal is a tty. Tracked in https://github.com/databricks/bricks/issues/379
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
}

View File

@ -0,0 +1,116 @@
package mutator
import (
"context"
"testing"
"github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/config"
"github.com/databricks/bricks/bundle/config/variable"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestSetVariableFromProcessEnvVar(t *testing.T) {
defaultVal := "default"
variable := variable.Variable{
Description: "a test variable",
Default: &defaultVal,
}
// set value for variable as an environment variable
t.Setenv("BUNDLE_VAR_foo", "process-env")
err := setVariable(&variable, "foo")
require.NoError(t, err)
assert.Equal(t, *variable.Value, "process-env")
}
func TestSetVariableUsingDefaultValue(t *testing.T) {
defaultVal := "default"
variable := variable.Variable{
Description: "a test variable",
Default: &defaultVal,
}
err := setVariable(&variable, "foo")
require.NoError(t, err)
assert.Equal(t, *variable.Value, "default")
}
func TestSetVariableWhenAlreadyAValueIsAssigned(t *testing.T) {
defaultVal := "default"
val := "assigned-value"
variable := variable.Variable{
Description: "a test variable",
Default: &defaultVal,
Value: &val,
}
// since a value is already assigned to the variable, it would not be overridden
// by the default value
err := setVariable(&variable, "foo")
require.NoError(t, err)
assert.Equal(t, *variable.Value, "assigned-value")
}
func TestSetVariableEnvVarValueDoesNotOverridePresetValue(t *testing.T) {
defaultVal := "default"
val := "assigned-value"
variable := variable.Variable{
Description: "a test variable",
Default: &defaultVal,
Value: &val,
}
// set value for variable as an environment variable
t.Setenv("BUNDLE_VAR_foo", "process-env")
// since a value is already assigned to the variable, it would not be overridden
// by the value from environment
err := setVariable(&variable, "foo")
require.NoError(t, err)
assert.Equal(t, *variable.Value, "assigned-value")
}
func TestSetVariablesErrorsIfAValueCouldNotBeResolved(t *testing.T) {
variable := variable.Variable{
Description: "a test variable with no default",
}
// fails because we could not resolve a value for the variable
err := setVariable(&variable, "foo")
assert.ErrorContains(t, err, "no value assigned to required variable foo. Assignment can be done through the \"--var\" flag or by setting the BUNDLE_VAR_foo environment variable")
}
func TestSetVariablesMutator(t *testing.T) {
defaultValForA := "default-a"
defaultValForB := "default-b"
valForC := "assigned-val-c"
bundle := &bundle.Bundle{
Config: config.Root{
Variables: map[string]*variable.Variable{
"a": {
Description: "resolved to default value",
Default: &defaultValForA,
},
"b": {
Description: "resolved from environment vairables",
Default: &defaultValForB,
},
"c": {
Description: "has already been assigned a value",
Value: &valForC,
},
},
},
}
t.Setenv("BUNDLE_VAR_b", "env-var-b")
_, err := SetVariables().Apply(context.Background(), bundle)
require.NoError(t, err)
assert.Equal(t, "default-a", *bundle.Config.Variables["a"].Value)
assert.Equal(t, "env-var-b", *bundle.Config.Variables["b"].Value)
assert.Equal(t, "assigned-val-c", *bundle.Config.Variables["c"].Value)
}

View File

@ -1,9 +1,12 @@
package config package config
import ( import (
"fmt"
"os" "os"
"path/filepath" "path/filepath"
"strings"
"github.com/databricks/bricks/bundle/config/variable"
"github.com/ghodss/yaml" "github.com/ghodss/yaml"
"github.com/imdario/mergo" "github.com/imdario/mergo"
) )
@ -16,6 +19,9 @@ type Root struct {
// It is set when loading `bundle.yml`. // It is set when loading `bundle.yml`.
Path string `json:"-" bundle:"readonly"` Path string `json:"-" bundle:"readonly"`
// Contains user defined variables
Variables map[string]*variable.Variable `json:"variables,omitempty"`
// Bundle contains details about this bundle, such as its name, // Bundle contains details about this bundle, such as its name,
// version of the spec (TODO), default cluster, default warehouse, etc. // version of the spec (TODO), default cluster, default warehouse, etc.
Bundle Bundle `json:"bundle"` 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 {
parsedVariable := strings.SplitN(variable, "=", 2)
if len(parsedVariable) != 2 {
return fmt.Errorf("unexpected flag value for variable assignment: %s", variable)
}
name := parsedVariable[0]
val := parsedVariable[1]
if _, ok := r.Variables[name]; !ok {
return fmt.Errorf("variable %s has not been defined", name)
}
err := r.Variables[name].Set(val)
if err != nil {
return fmt.Errorf("failed to assign %s to %s: %s", val, name, err)
}
}
return nil
}
func (r *Root) Load(path string) error { func (r *Root) Load(path string) error {
raw, err := os.ReadFile(path) raw, err := os.ReadFile(path)
if err != nil { if err != nil {

View File

@ -5,6 +5,7 @@ import (
"reflect" "reflect"
"testing" "testing"
"github.com/databricks/bricks/bundle/config/variable"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
@ -93,3 +94,63 @@ func TestDuplicateIdOnMergeReturnsError(t *testing.T) {
err = root.Merge(other) 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)") 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.Variable{
"foo": {
Default: &fooDefault,
Description: "an optional variable since default is defined",
},
"bar": {
Description: "a required variable",
},
},
}
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 TestInitializeVariablesWithAnEqualSignInValue(t *testing.T) {
root := &Root{
Variables: map[string]*variable.Variable{
"foo": {
Description: "a variable called foo",
},
},
}
err := root.InitializeVariables([]string{"foo=123=567"})
assert.NoError(t, err)
assert.Equal(t, "123=567", *(root.Variables["foo"].Value))
}
func TestInitializeVariablesInvalidFormat(t *testing.T) {
root := &Root{
Variables: map[string]*variable.Variable{
"foo": {
Description: "a variable called foo",
},
},
}
err := root.InitializeVariables([]string{"foo"})
assert.ErrorContains(t, err, "unexpected flag value for variable assignment: foo")
}
func TestInitializeVariablesUndefinedVariables(t *testing.T) {
root := &Root{
Variables: map[string]*variable.Variable{
"foo": {
Description: "A required variable",
},
},
}
err := root.InitializeVariables([]string{"bar=567"})
assert.ErrorContains(t, err, "variable bar has not been defined")
}

View File

@ -0,0 +1,45 @@
package variable
import (
"fmt"
)
const VariableReferencePrefix = "var"
// An input variable for the bundle config
type Variable struct {
// A default value which then makes the variable optional
Default *string `json:"default,omitempty"`
// Documentation for this input variable
Description string `json:"description,omitempty"`
// This field stores the resolved value for the variable. The variable are
// resolved in the following priority order (from highest to lowest)
//
// 1. Command line flag. For example: `--var="foo=bar"`
// 2. Environment variable. eg: BUNDLE_VAR_foo=bar
// 3. default value defined in bundle config
// 4. Throw error, since if no default value is defined, then the variable
// is required
Value *string `json:"value,omitempty" bundle:"readonly"`
}
// True if the variable has been assigned a default value. Variables without a
// a default value are by defination required
func (v *Variable) HasDefault() bool {
return v.Default != nil
}
// True if variable has already been assigned a value
func (v *Variable) HasValue() bool {
return v.Value != nil
}
func (v *Variable) Set(val string) error {
if v.HasValue() {
return fmt.Errorf("variable has already been assigned value: %s", *v.Value)
}
v.Value = &val
return nil
}

View File

@ -4,6 +4,7 @@ import (
"github.com/databricks/bricks/bundle" "github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/config/interpolation" "github.com/databricks/bricks/bundle/config/interpolation"
"github.com/databricks/bricks/bundle/config/mutator" "github.com/databricks/bricks/bundle/config/mutator"
"github.com/databricks/bricks/bundle/config/variable"
"github.com/databricks/bricks/bundle/deploy/terraform" "github.com/databricks/bricks/bundle/deploy/terraform"
) )
@ -18,9 +19,11 @@ func Initialize() bundle.Mutator {
mutator.DefineDefaultWorkspaceRoot(), mutator.DefineDefaultWorkspaceRoot(),
mutator.ExpandWorkspaceRoot(), mutator.ExpandWorkspaceRoot(),
mutator.DefineDefaultWorkspacePaths(), mutator.DefineDefaultWorkspacePaths(),
mutator.SetVariables(),
interpolation.Interpolate( interpolation.Interpolate(
interpolation.IncludeLookupsInPath("bundle"), interpolation.IncludeLookupsInPath("bundle"),
interpolation.IncludeLookupsInPath("workspace"), interpolation.IncludeLookupsInPath("workspace"),
interpolation.IncludeLookupsInPath(variable.VariableReferencePrefix),
), ),
mutator.TranslatePaths(), mutator.TranslatePaths(),
terraform.Initialize(), terraform.Initialize(),

View File

@ -0,0 +1,10 @@
variables:
a:
description: optional variable
default: abc
b:
description: required variable
bundle:
name: ${var.a} ${var.b}

View File

@ -0,0 +1,35 @@
package config_tests
import (
"context"
"testing"
"github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/config/interpolation"
"github.com/databricks/bricks/bundle/config/mutator"
"github.com/databricks/bricks/bundle/config/variable"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestVariables(t *testing.T) {
t.Setenv("BUNDLE_VAR_b", "def")
b := load(t, "./variables")
err := bundle.Apply(context.Background(), b, []bundle.Mutator{
mutator.SetVariables(),
interpolation.Interpolate(
interpolation.IncludeLookupsInPath(variable.VariableReferencePrefix),
)})
require.NoError(t, err)
assert.Equal(t, "abc def", b.Config.Bundle.Name)
}
func TestVariablesLoadingFailsWhenRequiredVariableIsNotSpecified(t *testing.T) {
b := load(t, "./variables")
err := bundle.Apply(context.Background(), b, []bundle.Mutator{
mutator.SetVariables(),
interpolation.Interpolate(
interpolation.IncludeLookupsInPath(variable.VariableReferencePrefix),
)})
assert.ErrorContains(t, err, "no value assigned to required variable b. Assignment can be done through the \"--var\" flag or by setting the BUNDLE_VAR_b environment variable")
}

View File

@ -4,14 +4,14 @@ import (
"fmt" "fmt"
"github.com/databricks/bricks/bundle" "github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/cmd/root" bundleCmd "github.com/databricks/bricks/cmd/bundle"
"github.com/spf13/cobra" "github.com/spf13/cobra"
) )
var whoamiCmd = &cobra.Command{ var whoamiCmd = &cobra.Command{
Use: "whoami", Use: "whoami",
PreRunE: root.MustConfigureBundle, PreRunE: bundleCmd.ConfigureBundleWithVariables,
RunE: func(cmd *cobra.Command, args []string) error { RunE: func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context() ctx := cmd.Context()
w := bundle.Get(ctx).WorkspaceClient() w := bundle.Get(ctx).WorkspaceClient()

View File

@ -3,7 +3,6 @@ package bundle
import ( import (
"github.com/databricks/bricks/bundle" "github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/phases" "github.com/databricks/bricks/bundle/phases"
"github.com/databricks/bricks/cmd/root"
"github.com/spf13/cobra" "github.com/spf13/cobra"
) )
@ -11,7 +10,7 @@ var deployCmd = &cobra.Command{
Use: "deploy", Use: "deploy",
Short: "Deploy bundle", Short: "Deploy bundle",
PreRunE: root.MustConfigureBundle, PreRunE: ConfigureBundleWithVariables,
RunE: func(cmd *cobra.Command, args []string) error { RunE: func(cmd *cobra.Command, args []string) error {
b := bundle.Get(cmd.Context()) b := bundle.Get(cmd.Context())

View File

@ -6,7 +6,6 @@ import (
"github.com/databricks/bricks/bundle" "github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/phases" "github.com/databricks/bricks/bundle/phases"
"github.com/databricks/bricks/cmd/root"
"github.com/databricks/bricks/libs/cmdio" "github.com/databricks/bricks/libs/cmdio"
"github.com/databricks/bricks/libs/flags" "github.com/databricks/bricks/libs/flags"
"github.com/spf13/cobra" "github.com/spf13/cobra"
@ -17,7 +16,7 @@ var destroyCmd = &cobra.Command{
Use: "destroy", Use: "destroy",
Short: "Destroy deployed bundle resources", Short: "Destroy deployed bundle resources",
PreRunE: root.MustConfigureBundle, PreRunE: ConfigureBundleWithVariables,
RunE: func(cmd *cobra.Command, args []string) error { RunE: func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context() ctx := cmd.Context()
b := bundle.Get(ctx) b := bundle.Get(ctx)

View File

@ -15,6 +15,9 @@ func AddCommand(cmd *cobra.Command) {
rootCmd.AddCommand(cmd) rootCmd.AddCommand(cmd)
} }
var variables []string
func init() { func init() {
root.RootCmd.AddCommand(rootCmd) root.RootCmd.AddCommand(rootCmd)
AddVariableFlag(rootCmd)
} }

View File

@ -20,7 +20,7 @@ var runCmd = &cobra.Command{
Short: "Run a workload (e.g. a job or a pipeline)", Short: "Run a workload (e.g. a job or a pipeline)",
Args: cobra.ExactArgs(1), Args: cobra.ExactArgs(1),
PreRunE: root.MustConfigureBundle, PreRunE: ConfigureBundleWithVariables,
RunE: func(cmd *cobra.Command, args []string) error { RunE: func(cmd *cobra.Command, args []string) error {
b := bundle.Get(cmd.Context()) b := bundle.Get(cmd.Context())
err := bundle.Apply(cmd.Context(), b, []bundle.Mutator{ err := bundle.Apply(cmd.Context(), b, []bundle.Mutator{

View File

@ -6,7 +6,6 @@ import (
"github.com/databricks/bricks/bundle" "github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/phases" "github.com/databricks/bricks/bundle/phases"
"github.com/databricks/bricks/cmd/root"
"github.com/databricks/bricks/libs/log" "github.com/databricks/bricks/libs/log"
"github.com/databricks/bricks/libs/sync" "github.com/databricks/bricks/libs/sync"
"github.com/spf13/cobra" "github.com/spf13/cobra"
@ -35,7 +34,7 @@ var syncCmd = &cobra.Command{
Short: "Synchronize bundle tree to the workspace", Short: "Synchronize bundle tree to the workspace",
Args: cobra.NoArgs, Args: cobra.NoArgs,
PreRunE: root.MustConfigureBundle, PreRunE: ConfigureBundleWithVariables,
RunE: func(cmd *cobra.Command, args []string) error { RunE: func(cmd *cobra.Command, args []string) error {
b := bundle.Get(cmd.Context()) b := bundle.Get(cmd.Context())

View File

@ -5,7 +5,6 @@ import (
"github.com/databricks/bricks/bundle" "github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/phases" "github.com/databricks/bricks/bundle/phases"
"github.com/databricks/bricks/cmd/root"
"github.com/spf13/cobra" "github.com/spf13/cobra"
) )
@ -13,9 +12,10 @@ var validateCmd = &cobra.Command{
Use: "validate", Use: "validate",
Short: "Validate configuration", Short: "Validate configuration",
PreRunE: root.MustConfigureBundle, PreRunE: ConfigureBundleWithVariables,
RunE: func(cmd *cobra.Command, args []string) error { RunE: func(cmd *cobra.Command, args []string) error {
b := bundle.Get(cmd.Context()) b := bundle.Get(cmd.Context())
err := bundle.Apply(cmd.Context(), b, []bundle.Mutator{ err := bundle.Apply(cmd.Context(), b, []bundle.Mutator{
phases.Initialize(), phases.Initialize(),
}) })

23
cmd/bundle/variables.go Normal file
View File

@ -0,0 +1,23 @@
package bundle
import (
"github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/cmd/root"
"github.com/spf13/cobra"
)
func ConfigureBundleWithVariables(cmd *cobra.Command, args []string) error {
// Load bundle config and apply environment
err := root.MustConfigureBundle(cmd, args)
if err != nil {
return err
}
// Initialize variables by assigning them values passed as command line flags
b := bundle.Get(cmd.Context())
return b.Config.InitializeVariables(variables)
}
func AddVariableFlag(cmd *cobra.Command) {
cmd.PersistentFlags().StringSliceVar(&variables, "var", []string{}, `set values for variables defined in bundle config. Example: --var="foo=bar"`)
}