Add enum support for bundle templates (#668)

## Changes
This PR includes:
1. Adding enum field to the json schema struct
2. Adding prompting logic for enum values. See demo for how it looks
3. Validation rules, validating the default value and config values when
an enum list is specified

This will now enable template authors to use enums for input parameters.

## Tests
Manually and new unit tests

libs/cmdio/logger.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/libs/flags"
+	"github.com/manifoldco/promptui"
 )
 
 // This is the interface for all io interactions with a user
@@ -104,6 +105,36 @@ func AskYesOrNo(ctx context.Context, question string) (bool, error) {
 	return false, nil
 }
 
+func AskSelect(ctx context.Context, question string, choices []string) (string, error) {
+	logger, ok := FromContext(ctx)
+	if !ok {
+		logger = Default()
+	}
+	return logger.AskSelect(question, choices)
+}
+
+func (l *Logger) AskSelect(question string, choices []string) (string, error) {
+	if l.Mode == flags.ModeJson {
+		return "", fmt.Errorf("question prompts are not supported in json mode")
+	}
+
+	prompt := promptui.Select{
+		Label:    question,
+		Items:    choices,
+		HideHelp: true,
+		Templates: &promptui.SelectTemplates{
+			Label:    "{{.}}: ",
+			Selected: fmt.Sprintf("%s: {{.}}", question),
+		},
+	}
+
+	_, ans, err := prompt.Run()
+	if err != nil {
+		return "", err
+	}
+	return ans, nil
+}
+
 func (l *Logger) Ask(question string, defaultVal string) (string, error) {
 	if l.Mode == flags.ModeJson {
 		return "", fmt.Errorf("question prompts are not supported in json mode")

libs/cmdio/logger_test.go

@@ -1,6 +1,7 @@
 package cmdio
 
 import (
+	"context"
 	"testing"
 
 	"github.com/databricks/cli/libs/flags"
@@ -12,3 +13,11 @@ func TestAskFailedInJsonMode(t *testing.T) {
 	_, err := l.Ask("What is your spirit animal?", "")
 	assert.ErrorContains(t, err, "question prompts are not supported in json mode")
 }
+
+func TestAskChoiceFailsInJsonMode(t *testing.T) {
+	l := NewLogger(flags.ModeJson)
+	ctx := NewContext(context.Background(), l)
+
+	_, err := AskSelect(ctx, "what is a question?", []string{"b", "c", "a"})
+	assert.EqualError(t, err, "question prompts are not supported in json mode")
+}

libs/jsonschema/instance.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"slices"
 )
 
 // Load a JSON document and validate it against the JSON schema. Instance here
@@ -39,13 +40,18 @@ func (s *Schema) LoadInstance(path string) (map[string]any, error) {
 }
 
 func (s *Schema) ValidateInstance(instance map[string]any) error {
-	if err := s.validateAdditionalProperties(instance); err != nil {
+	for _, fn := range []func(map[string]any) error{
+		s.validateAdditionalProperties,
+		s.validateEnum,
+		s.validateRequired,
+		s.validateTypes,
+	} {
+		err := fn(instance)
+		if err != nil {
 			return err
 		}
-	if err := s.validateRequired(instance); err != nil {
-		return err
 	}
-	return s.validateTypes(instance)
+	return nil
 }
 
 // If additional properties is set to false, this function validates instance only
@@ -89,3 +95,19 @@ func (s *Schema) validateTypes(instance map[string]any) error {
 	}
 	return nil
 }
+
+func (s *Schema) validateEnum(instance map[string]any) error {
+	for k, v := range instance {
+		fieldInfo, ok := s.Properties[k]
+		if !ok {
+			continue
+		}
+		if fieldInfo.Enum == nil {
+			continue
+		}
+		if !slices.Contains(fieldInfo.Enum, v) {
+			return fmt.Errorf("expected value of property %s to be one of %v. Found: %v", k, fieldInfo.Enum, v)
+		}
+	}
+	return nil
+}

libs/jsonschema/instance_test.go

@@ -127,3 +127,29 @@ func TestLoadInstance(t *testing.T) {
 	_, err = schema.LoadInstance("./testdata/instance-load/invalid-type-instance.json")
 	assert.EqualError(t, err, "incorrect type for property string_val: expected type string, but value is 123")
 }
+
+func TestValidateInstanceEnum(t *testing.T) {
+	schema, err := Load("./testdata/instance-validate/test-schema-enum.json")
+	require.NoError(t, err)
+
+	validInstance := map[string]any{
+		"foo": "b",
+		"bar": int64(6),
+	}
+	assert.NoError(t, schema.validateEnum(validInstance))
+	assert.NoError(t, schema.ValidateInstance(validInstance))
+
+	invalidStringInstance := map[string]any{
+		"foo": "d",
+		"bar": int64(2),
+	}
+	assert.EqualError(t, schema.validateEnum(invalidStringInstance), "expected value of property foo to be one of [a b c]. Found: d")
+	assert.EqualError(t, schema.ValidateInstance(invalidStringInstance), "expected value of property foo to be one of [a b c]. Found: d")
+
+	invalidIntInstance := map[string]any{
+		"foo": "a",
+		"bar": int64(1),
+	}
+	assert.EqualError(t, schema.validateEnum(invalidIntInstance), "expected value of property bar to be one of [2 4 6]. Found: 1")
+	assert.EqualError(t, schema.ValidateInstance(invalidIntInstance), "expected value of property bar to be one of [2 4 6]. Found: 1")
+}

libs/jsonschema/schema.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"slices"
 )
 
 // defines schema for a json object
@@ -41,6 +42,9 @@ type Schema struct {
 	// Default value for the property / object
 	Default any `json:"default,omitempty"`
 
+	// List of valid values for a JSON instance for this schema.
+	Enum []any `json:"enum,omitempty"`
+
 	// Extension embeds our custom JSON schema extensions.
 	Extension
 }
@@ -84,6 +88,30 @@ func (schema *Schema) validate() error {
 		}
 	}
 
+	// Validate enum field values for properties are consistent with types.
+	for name, property := range schema.Properties {
+		if property.Enum == nil {
+			continue
+		}
+		for i, enum := range property.Enum {
+			err := validateType(enum, property.Type)
+			if err != nil {
+				return fmt.Errorf("type validation for enum at index %v failed for property %s: %w", i, name, err)
+			}
+		}
+	}
+
+	// Validate default value is contained in the list of enums if both are defined.
+	for name, property := range schema.Properties {
+		if property.Default == nil || property.Enum == nil {
+			continue
+		}
+		// We expect the default value to be consistent with the list of enum
+		// values.
+		if !slices.Contains(property.Enum, property.Default) {
+			return fmt.Errorf("list of enum values for property %s does not contain default value %v: %v", name, property.Default, property.Enum)
+		}
+	}
 	return nil
 }
 
@@ -115,6 +143,12 @@ func Load(path string) (*Schema, error) {
 				return nil, fmt.Errorf("failed to parse default value for property %s: %w", name, err)
 			}
 		}
+		for i, enum := range property.Enum {
+			property.Enum[i], err = toInteger(enum)
+			if err != nil {
+				return nil, fmt.Errorf("failed to parse enum value %v at index %v for property %s: %w", enum, i, name, err)
+			}
+		}
 	}
 
 	return schema, schema.validate()

libs/jsonschema/schema_test.go

@@ -47,6 +47,7 @@ func TestSchemaLoadIntegers(t *testing.T) {
 	schema, err := Load("./testdata/schema-load-int/schema-valid.json")
 	assert.NoError(t, err)
 	assert.Equal(t, int64(1), schema.Properties["abc"].Default)
+	assert.Equal(t, []any{int64(1), int64(2), int64(3)}, schema.Properties["abc"].Enum)
 }
 
 func TestSchemaLoadIntegersWithInvalidDefault(t *testing.T) {
@@ -54,6 +55,11 @@ func TestSchemaLoadIntegersWithInvalidDefault(t *testing.T) {
 	assert.EqualError(t, err, "failed to parse default value for property abc: expected integer value, got: 1.1")
 }
 
+func TestSchemaLoadIntegersWithInvalidEnums(t *testing.T) {
+	_, err := Load("./testdata/schema-load-int/schema-invalid-enum.json")
+	assert.EqualError(t, err, "failed to parse enum value 2.4 at index 1 for property abc: expected integer value, got: 2.4")
+}
+
 func TestSchemaValidateDefaultType(t *testing.T) {
 	invalidSchema := &Schema{
 		Properties: map[string]*Schema{
@@ -79,3 +85,57 @@ func TestSchemaValidateDefaultType(t *testing.T) {
 	err = validSchema.validate()
 	assert.NoError(t, err)
 }
+
+func TestSchemaValidateEnumType(t *testing.T) {
+	invalidSchema := &Schema{
+		Properties: map[string]*Schema{
+			"foo": {
+				Type: "boolean",
+				Enum: []any{true, "false"},
+			},
+		},
+	}
+
+	err := invalidSchema.validate()
+	assert.EqualError(t, err, "type validation for enum at index 1 failed for property foo: expected type boolean, but value is \"false\"")
+
+	validSchema := &Schema{
+		Properties: map[string]*Schema{
+			"foo": {
+				Type: "boolean",
+				Enum: []any{true, false},
+			},
+		},
+	}
+
+	err = validSchema.validate()
+	assert.NoError(t, err)
+}
+
+func TestSchemaValidateErrorWhenDefaultValueIsNotInEnums(t *testing.T) {
+	invalidSchema := &Schema{
+		Properties: map[string]*Schema{
+			"foo": {
+				Type:    "string",
+				Default: "abc",
+				Enum:    []any{"def", "ghi"},
+			},
+		},
+	}
+
+	err := invalidSchema.validate()
+	assert.EqualError(t, err, "list of enum values for property foo does not contain default value abc: [def ghi]")
+
+	validSchema := &Schema{
+		Properties: map[string]*Schema{
+			"foo": {
+				Type:    "string",
+				Default: "abc",
+				Enum:    []any{"def", "ghi", "abc"},
+			},
+		},
+	}
+
+	err = validSchema.validate()
+	assert.NoError(t, err)
+}

libs/jsonschema/testdata/instance-validate/test-schema-enum.json

@@ -0,0 +1,12 @@
+{
+    "properties": {
+        "foo": {
+            "type": "string",
+            "enum": ["a", "b", "c"]
+        },
+        "bar": {
+            "type": "integer",
+            "enum": [2,4,6]
+        }
+    }
+}

libs/jsonschema/testdata/schema-load-int/schema-invalid-enum.json

@@ -0,0 +1,10 @@
+{
+    "type": "object",
+    "properties": {
+        "abc": {
+            "type": "integer",
+            "default": 1,
+            "enum": [1,2.4,3]
+        }
+    }
+}

libs/jsonschema/testdata/schema-load-int/schema-valid.json

@@ -3,7 +3,8 @@
     "properties": {
         "abc": {
             "type": "integer",
-            "default": 1
+            "default": 1,
+            "enum": [1,2,3]
         }
     }
 }

libs/jsonschema/utils.go

@@ -71,6 +71,18 @@ func ToString(v any, T Type) (string, error) {
 	}
 }
 
+func ToStringSlice(arr []any, T Type) ([]string, error) {
+	res := []string{}
+	for _, v := range arr {
+		s, err := ToString(v, T)
+		if err != nil {
+			return nil, err
+		}
+		res = append(res, s)
+	}
+	return res, nil
+}
+
 func FromString(s string, T Type) (any, error) {
 	if T == StringType {
 		return s, nil

libs/jsonschema/utils_test.go

@@ -118,3 +118,13 @@ func TestTemplateFromString(t *testing.T) {
 	_, err = FromString("1.0", "foobar")
 	assert.EqualError(t, err, "unknown json schema type: \"foobar\"")
 }
+
+func TestTemplateToStringSlice(t *testing.T) {
+	s, err := ToStringSlice([]any{"a", "b", "c"}, StringType)
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"a", "b", "c"}, s)
+
+	s, err = ToStringSlice([]any{1.1, 2.2, 3.3}, NumberType)
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"1.1", "2.2", "3.3"}, s)
+}

libs/template/config.go

@@ -102,10 +102,24 @@ func (c *config) promptForValues() error {
 		}
 
 		// Get user input by running the prompt
-		userInput, err := cmdio.Ask(c.ctx, property.Description, defaultVal)
+		var userInput string
+		if property.Enum != nil {
+			// convert list of enums to string slice
+			enums, err := jsonschema.ToStringSlice(property.Enum, property.Type)
 			if err != nil {
 				return err
 			}
+			userInput, err = cmdio.AskSelect(c.ctx, property.Description, enums)
+			if err != nil {
+				return err
+			}
+		} else {
+			userInput, err = cmdio.Ask(c.ctx, property.Description, defaultVal)
+			if err != nil {
+				return err
+			}
+
+		}
 
 		// Convert user input string back to a value
 		c.values[name], err = jsonschema.FromString(userInput, property.Type)

libs/template/config_test.go

@@ -142,3 +142,30 @@ func TestTemplateValidateSchema(t *testing.T) {
 	err = validateSchema(toSchema("array"))
 	assert.EqualError(t, err, "property type array is not supported by bundle templates")
 }
+
+func TestTemplateEnumValidation(t *testing.T) {
+	schema := jsonschema.Schema{
+		Properties: map[string]*jsonschema.Schema{
+			"abc": {
+				Type: "integer",
+				Enum: []any{1, 2, 3, 4},
+			},
+		},
+	}
+
+	c := &config{
+		schema: &schema,
+		values: map[string]any{
+			"abc": 5,
+		},
+	}
+	assert.EqualError(t, c.validate(), "validation for template input parameters failed. expected value of property abc to be one of [1 2 3 4]. Found: 5")
+
+	c = &config{
+		schema: &schema,
+		values: map[string]any{
+			"abc": 4,
+		},
+	}
+	assert.NoError(t, c.validate())
+}