Fix bundle schema for variables (#1396)

## Changes
This PR fixes the variable schema to:
1. Allow non-string values in the "default" value of a variable.
2. Allow non-string overrides in a target for a variable. 

## Tests
Manually. There are no longer squiggly lines. 

Before:
<img width="329" alt="Screenshot 2024-04-24 at 3 26 43 PM"
src="https://github.com/databricks/cli/assets/88374338/43be02c2-80a4-4f80-bd79-0f3e1e93ee17">


After:
<img width="361" alt="Screenshot 2024-04-24 at 3 26 10 PM"
src="https://github.com/databricks/cli/assets/88374338/2c1fb892-a2a2-478b-8d2e-9bda6d844b54">

bundle/config/mutator/set_variables.go

@@ -53,8 +53,6 @@ func setVariable(ctx context.Context, v *variable.Variable, name string) diag.Di
 	}
 
 	// 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/cli/issues/379
 	return diag.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)
 }
 

cmd/bundle/schema.go

@@ -7,9 +7,58 @@ import (
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/schema"
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/jsonschema"
 	"github.com/spf13/cobra"
 )
 
+func overrideVariables(s *jsonschema.Schema) error {
+	// Override schema for default values to allow for multiple primitive types.
+	// These are normalized to strings when converted to the typed representation.
+	err := s.SetByPath("variables.*.default", jsonschema.Schema{
+		AnyOf: []*jsonschema.Schema{
+			{
+				Type: jsonschema.StringType,
+			},
+			{
+				Type: jsonschema.BooleanType,
+			},
+			{
+				Type: jsonschema.NumberType,
+			},
+			{
+				Type: jsonschema.IntegerType,
+			},
+		},
+	})
+	if err != nil {
+		return err
+	}
+
+	// Override schema for variables in targets to allow just specifying the value
+	// along side overriding the variable definition if needed.
+	ns, err := s.GetByPath("variables.*")
+	if err != nil {
+		return err
+	}
+	return s.SetByPath("targets.*.variables.*", jsonschema.Schema{
+		AnyOf: []*jsonschema.Schema{
+			{
+				Type: jsonschema.StringType,
+			},
+			{
+				Type: jsonschema.BooleanType,
+			},
+			{
+				Type: jsonschema.NumberType,
+			},
+			{
+				Type: jsonschema.IntegerType,
+			},
+			&ns,
+		},
+	})
+}
+
 func newSchemaCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "schema",
@@ -30,6 +79,13 @@ func newSchemaCommand() *cobra.Command {
 			return err
 		}
 
+		// Override schema for variables to take into account normalization of default
+		// variable values and variable overrides in a target.
+		err = overrideVariables(schema)
+		if err != nil {
+			return err
+		}
+
 		// Print the JSON schema to stdout.
 		result, err := json.MarshalIndent(schema, "", "  ")
 		if err != nil {

libs/jsonschema/schema.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"regexp"
 	"slices"
+	"strings"
 
 	"github.com/databricks/cli/internal/build"
 	"golang.org/x/mod/semver"
@@ -81,6 +82,41 @@ func (s *Schema) ParseString(v string) (any, error) {
 	return fromString(v, s.Type)
 }
 
+func (s *Schema) getByPath(path string) (*Schema, error) {
+	p := strings.Split(path, ".")
+
+	res := s
+	for _, node := range p {
+		if node == "*" {
+			res = res.AdditionalProperties.(*Schema)
+			continue
+		}
+		var ok bool
+		res, ok = res.Properties[node]
+		if !ok {
+			return nil, fmt.Errorf("property %q not found in schema. Query path: %s", node, path)
+		}
+	}
+	return res, nil
+}
+
+func (s *Schema) GetByPath(path string) (Schema, error) {
+	v, err := s.getByPath(path)
+	if err != nil {
+		return Schema{}, err
+	}
+	return *v, nil
+}
+
+func (s *Schema) SetByPath(path string, v Schema) error {
+	dst, err := s.getByPath(path)
+	if err != nil {
+		return err
+	}
+	*dst = v
+	return nil
+}
+
 type Type string
 
 const (
@@ -97,7 +133,7 @@ const (
 func (schema *Schema) validateSchemaPropertyTypes() error {
 	for _, v := range schema.Properties {
 		switch v.Type {
-		case NumberType, BooleanType, StringType, IntegerType:
+		case NumberType, BooleanType, StringType, IntegerType, ObjectType, ArrayType:
 			continue
 		case "int", "int32", "int64":
 			return fmt.Errorf("type %s is not a recognized json schema type. Please use \"integer\" instead", v.Type)

libs/jsonschema/schema_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestSchemaValidateTypeNames(t *testing.T) {
@@ -305,3 +306,92 @@ func TestValidateSchemaSkippedPropertiesHaveDefaults(t *testing.T) {
 	err = s.validate()
 	assert.NoError(t, err)
 }
+
+func testSchema() *Schema {
+	return &Schema{
+		Type: "object",
+		Properties: map[string]*Schema{
+			"int_val": {
+				Type:    "integer",
+				Default: int64(123),
+			},
+			"string_val": {
+				Type: "string",
+			},
+			"object_val": {
+				Type: "object",
+				Properties: map[string]*Schema{
+					"bar": {
+						Type:    "string",
+						Default: "baz",
+					},
+				},
+				AdditionalProperties: &Schema{
+					Type: "object",
+					Properties: map[string]*Schema{
+						"foo": {
+							Type:    "string",
+							Default: "zab",
+						},
+					},
+				},
+			},
+		},
+	}
+
+}
+
+func TestSchemaGetByPath(t *testing.T) {
+	s := testSchema()
+
+	ss, err := s.GetByPath("int_val")
+	require.NoError(t, err)
+	assert.Equal(t, Schema{
+		Type:    IntegerType,
+		Default: int64(123),
+	}, ss)
+
+	ss, err = s.GetByPath("string_val")
+	require.NoError(t, err)
+	assert.Equal(t, Schema{
+		Type: StringType,
+	}, ss)
+
+	ss, err = s.GetByPath("object_val.bar")
+	require.NoError(t, err)
+	assert.Equal(t, Schema{
+		Type:    StringType,
+		Default: "baz",
+	}, ss)
+
+	ss, err = s.GetByPath("object_val.*.foo")
+	require.NoError(t, err)
+	assert.Equal(t, Schema{
+		Type:    StringType,
+		Default: "zab",
+	}, ss)
+}
+
+func TestSchemaSetByPath(t *testing.T) {
+	s := testSchema()
+
+	err := s.SetByPath("int_val", Schema{
+		Type:    IntegerType,
+		Default: int64(456),
+	})
+	require.NoError(t, err)
+	assert.Equal(t, int64(456), s.Properties["int_val"].Default)
+
+	err = s.SetByPath("object_val.*.foo", Schema{
+		Type:    StringType,
+		Default: "zooby",
+	})
+	require.NoError(t, err)
+
+	ns, err := s.GetByPath("object_val.*.foo")
+	require.NoError(t, err)
+	assert.Equal(t, Schema{
+		Type:    StringType,
+		Default: "zooby",
+	}, ns)
+}