From 7c96270db8c95a54e8d59893c15e370dff1f2f76 Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Fri, 8 Sep 2023 14:07:22 +0200
Subject: [PATCH] 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                          | 31 ++++++++++
 libs/cmdio/logger_test.go                     |  9 +++
 libs/jsonschema/instance.go                   | 34 +++++++++--
 libs/jsonschema/instance_test.go              | 26 ++++++++
 libs/jsonschema/schema.go                     | 34 +++++++++++
 libs/jsonschema/schema_test.go                | 60 +++++++++++++++++++
 .../instance-validate/test-schema-enum.json   | 12 ++++
 .../schema-load-int/schema-invalid-enum.json  | 10 ++++
 .../schema-load-int/schema-valid.json         |  3 +-
 libs/jsonschema/utils.go                      | 12 ++++
 libs/jsonschema/utils_test.go                 | 10 ++++
 libs/template/config.go                       | 20 ++++++-
 libs/template/config_test.go                  | 27 +++++++++
 13 files changed, 278 insertions(+), 10 deletions(-)
 create mode 100644 libs/jsonschema/testdata/instance-validate/test-schema-enum.json
 create mode 100644 libs/jsonschema/testdata/schema-load-int/schema-invalid-enum.json

diff --git a/libs/cmdio/logger.go b/libs/cmdio/logger.go
index 0663306e1..7d760b998 100644
--- a/libs/cmdio/logger.go
+++ b/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")
diff --git a/libs/cmdio/logger_test.go b/libs/cmdio/logger_test.go
index da6190462..c5c00d022 100644
--- a/libs/cmdio/logger_test.go
+++ b/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")
+}
diff --git a/libs/jsonschema/instance.go b/libs/jsonschema/instance.go
index 02ab9f281..229a45b53 100644
--- a/libs/jsonschema/instance.go
+++ b/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 {
-		return err
+	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
+}
diff --git a/libs/jsonschema/instance_test.go b/libs/jsonschema/instance_test.go
index d5e0766dd..ffd10ca43 100644
--- a/libs/jsonschema/instance_test.go
+++ b/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")
+}
diff --git a/libs/jsonschema/schema.go b/libs/jsonschema/schema.go
index 44c65ecc6..108102a64 100644
--- a/libs/jsonschema/schema.go
+++ b/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()
diff --git a/libs/jsonschema/schema_test.go b/libs/jsonschema/schema_test.go
index 5b92d8466..db559ea88 100644
--- a/libs/jsonschema/schema_test.go
+++ b/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)
+}
diff --git a/libs/jsonschema/testdata/instance-validate/test-schema-enum.json b/libs/jsonschema/testdata/instance-validate/test-schema-enum.json
new file mode 100644
index 000000000..75ffd6eb8
--- /dev/null
+++ b/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]
+        }
+    }
+}
diff --git a/libs/jsonschema/testdata/schema-load-int/schema-invalid-enum.json b/libs/jsonschema/testdata/schema-load-int/schema-invalid-enum.json
new file mode 100644
index 000000000..5bd2b3f2b
--- /dev/null
+++ b/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]
+        }
+    }
+}
diff --git a/libs/jsonschema/testdata/schema-load-int/schema-valid.json b/libs/jsonschema/testdata/schema-load-int/schema-valid.json
index 599ac04d0..a1167a6c9 100644
--- a/libs/jsonschema/testdata/schema-load-int/schema-valid.json
+++ b/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]
         }
     }
 }
diff --git a/libs/jsonschema/utils.go b/libs/jsonschema/utils.go
index 21866965e..66db9603e 100644
--- a/libs/jsonschema/utils.go
+++ b/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
diff --git a/libs/jsonschema/utils_test.go b/libs/jsonschema/utils_test.go
index 9686cf39b..29529aaa9 100644
--- a/libs/jsonschema/utils_test.go
+++ b/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)
+}
diff --git a/libs/template/config.go b/libs/template/config.go
index 6f980f613..21618ac9a 100644
--- a/libs/template/config.go
+++ b/libs/template/config.go
@@ -102,9 +102,23 @@ func (c *config) promptForValues() error {
 		}
 
 		// Get user input by running the prompt
-		userInput, err := cmdio.Ask(c.ctx, property.Description, defaultVal)
-		if err != nil {
-			return err
+		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
diff --git a/libs/template/config_test.go b/libs/template/config_test.go
index bba22c758..1b1fc3383 100644
--- a/libs/template/config_test.go
+++ b/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())
+}