From 04946d268a9268203ad4c619747a92a7d1c9ddda Mon Sep 17 00:00:00 2001
From: Arpit Jasapara <arpit.jasapara@databricks.com>
Date: Fri, 19 Jan 2024 01:21:24 -0800
Subject: [PATCH] Add AnyOf Support to SkipPromptIf

---
 libs/jsonschema/schema.go    |  3 ++
 libs/template/config.go      | 31 +++++++++++--
 libs/template/config_test.go | 86 ++++++++++++++++++++++++++++++++++++
 3 files changed, 116 insertions(+), 4 deletions(-)

diff --git a/libs/jsonschema/schema.go b/libs/jsonschema/schema.go
index 443e7af6e..967e2e9cd 100644
--- a/libs/jsonschema/schema.go
+++ b/libs/jsonschema/schema.go
@@ -60,6 +60,9 @@ type Schema struct {
 
 	// Extension embeds our custom JSON schema extensions.
 	Extension
+
+	// Schema that must match any of the schemas in the array
+	AnyOf []*Schema `json:"anyOf,omitempty"`
 }
 
 // Default value defined in a JSON Schema, represented as a string.
diff --git a/libs/template/config.go b/libs/template/config.go
index 14e09fe56..71c9cd52d 100644
--- a/libs/template/config.go
+++ b/libs/template/config.go
@@ -127,11 +127,34 @@ func (c *config) skipPrompt(p jsonschema.Property, r *renderer) (bool, error) {
 
 	// Check if conditions specified by template author for skipping the prompt
 	// are satisfied. If they are not, we have to prompt for a user input.
-	for name, property := range p.Schema.SkipPromptIf.Properties {
-		if v, ok := c.values[name]; ok && v == property.Const {
-			continue
+	if p.Schema.SkipPromptIf.Properties != nil {
+		for name, property := range p.Schema.SkipPromptIf.Properties {
+			if v, ok := c.values[name]; ok && v == property.Const {
+				continue
+			}
+			return false, nil
+		}
+	}
+
+	// Check if AnyOf is set first. If so, then iterate over all property sets to find a match.
+	// If no match is found (allFalse stays true), we have to prompt for a user input.
+	if p.Schema.SkipPromptIf.AnyOf != nil {
+		allFalse := true
+		for _, properties := range p.Schema.SkipPromptIf.AnyOf {
+			match := true
+			for name, property := range properties.Properties {
+				if v, ok := c.values[name]; ok && v == property.Const {
+					continue
+				}
+				match = false
+			}
+			if match {
+				allFalse = false
+			}
+		}
+		if allFalse {
+			return false, nil
 		}
-		return false, nil
 	}
 
 	if p.Schema.Default == nil {
diff --git a/libs/template/config_test.go b/libs/template/config_test.go
index c4968ee1a..a02c80b50 100644
--- a/libs/template/config_test.go
+++ b/libs/template/config_test.go
@@ -398,3 +398,89 @@ func TestPromptIsSkipped(t *testing.T) {
 	assert.True(t, skip)
 	assert.Equal(t, "hello-world", c.values["xyz"])
 }
+
+func TestPromptIsSkippedAnyOf(t *testing.T) {
+	c := config{
+		ctx:    context.Background(),
+		values: make(map[string]any),
+		schema: &jsonschema.Schema{
+			Properties: map[string]*jsonschema.Schema{
+				"abc": {
+					Type: "string",
+				},
+				"def": {
+					Type: "integer",
+				},
+				"xyz": {
+					Type:    "string",
+					Default: "hello-world",
+					Extension: jsonschema.Extension{
+						SkipPromptIf: &jsonschema.Schema{
+							AnyOf: []*jsonschema.Schema{
+								{
+									Properties: map[string]*jsonschema.Schema{
+										"abc": {
+											Const: "foobar",
+										},
+										"def": {
+											Const: 123,
+										},
+									},
+								},
+								{
+									Properties: map[string]*jsonschema.Schema{
+										"abc": {
+											Const: "barfoo",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	// No skip condition defined. Prompt should not be skipped.
+	skip, err := c.skipPrompt(jsonschema.Property{
+		Name:   "abc",
+		Schema: c.schema.Properties["abc"],
+	}, testRenderer())
+	assert.NoError(t, err)
+	assert.False(t, skip)
+
+	// Values do not match skip condition. Prompt should not be skipped.
+	c.values["abc"] = "foobar"
+	c.values["def"] = 1234
+	skip, err = c.skipPrompt(jsonschema.Property{
+		Name:   "xyz",
+		Schema: c.schema.Properties["xyz"],
+	}, testRenderer())
+	assert.NoError(t, err)
+	assert.False(t, skip)
+	assert.NotContains(t, c.values, "xyz")
+
+	// Values match skip condition. Prompt should be skipped. Default value should
+	// be assigned to "xyz".
+	c.values["abc"] = "foobar"
+	c.values["def"] = 123
+	skip, err = c.skipPrompt(jsonschema.Property{
+		Name:   "xyz",
+		Schema: c.schema.Properties["xyz"],
+	}, testRenderer())
+	assert.NoError(t, err)
+	assert.True(t, skip)
+	assert.Equal(t, "hello-world", c.values["xyz"])
+
+	// Values match skip condition. Prompt should be skipped. Default value should
+	// be assigned to "xyz".
+	c.values["abc"] = "barfoo"
+	skip, err = c.skipPrompt(jsonschema.Property{
+		Name:   "xyz",
+		Schema: c.schema.Properties["xyz"],
+	}, testRenderer())
+	assert.NoError(t, err)
+	assert.True(t, skip)
+	assert.Equal(t, "hello-world", c.values["xyz"])
+}