From 455a11b470eb937843dca1fb3276128e698480ce Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Wed, 4 Dec 2024 10:46:32 +0100
Subject: [PATCH] update JSON schema

---
 bundle/internal/schema/main.go                | 19 +++++
 .../testdata/fail/incorrect_volume_type.yml   | 10 +++
 .../internal/schema/testdata/pass/volume.yml  |  9 +++
 bundle/phases/initialize.go                   |  1 +
 bundle/schema/jsonschema.json                 | 72 +++++++++++++++++++
 5 files changed, 111 insertions(+)
 create mode 100644 bundle/internal/schema/testdata/fail/incorrect_volume_type.yml
 create mode 100644 bundle/internal/schema/testdata/pass/volume.yml

diff --git a/bundle/internal/schema/main.go b/bundle/internal/schema/main.go
index ddeffe2fda..881ce3496c 100644
--- a/bundle/internal/schema/main.go
+++ b/bundle/internal/schema/main.go
@@ -93,6 +93,24 @@ func removeJobsFields(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema {
 	return s
 }
 
+// While volume_type is required in the volume create API, DABs automatically sets
+// it's value to "MANAGED" if it's not provided. Thus, we make it optional
+// in the bundle schema.
+func makeVolumeTypeOptional(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema {
+	if typ != reflect.TypeOf(resources.Volume{}) {
+		return s
+	}
+
+	res := []string{}
+	for _, r := range s.Required {
+		if r != "volume_type" {
+			res = append(res, r)
+		}
+	}
+	s.Required = res
+	return s
+}
+
 func main() {
 	if len(os.Args) != 2 {
 		fmt.Println("Usage: go run main.go <output-file>")
@@ -118,6 +136,7 @@ func main() {
 		p.addDescriptions,
 		p.addEnums,
 		removeJobsFields,
+		makeVolumeTypeOptional,
 		addInterpolationPatterns,
 	})
 	if err != nil {
diff --git a/bundle/internal/schema/testdata/fail/incorrect_volume_type.yml b/bundle/internal/schema/testdata/fail/incorrect_volume_type.yml
new file mode 100644
index 0000000000..d540056d4a
--- /dev/null
+++ b/bundle/internal/schema/testdata/fail/incorrect_volume_type.yml
@@ -0,0 +1,10 @@
+bundle:
+  name: volume with incorrect type
+
+resources:
+  volumes:
+    foo:
+      catalog_name: main
+      name: my_volume
+      schema_name: myschema
+      volume_type: incorrect_type
diff --git a/bundle/internal/schema/testdata/pass/volume.yml b/bundle/internal/schema/testdata/pass/volume.yml
new file mode 100644
index 0000000000..14684f36e5
--- /dev/null
+++ b/bundle/internal/schema/testdata/pass/volume.yml
@@ -0,0 +1,9 @@
+bundle:
+  name: a volume
+
+resources:
+  volumes:
+    foo:
+      catalog_name: main
+      name: my_volume
+      schema_name: myschema
diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go
index 3345872022..6fa0e5fede 100644
--- a/bundle/phases/initialize.go
+++ b/bundle/phases/initialize.go
@@ -68,6 +68,7 @@ func Initialize() bundle.Mutator {
 			mutator.SetRunAs(),
 			mutator.OverrideCompute(),
 			mutator.ConfigureDashboardDefaults(),
+			mutator.ConfigureVolumeDefaults(),
 			mutator.ProcessTargetMode(),
 			mutator.ApplyPresets(),
 			mutator.DefaultQueueing(),
diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json
index cf003f4235..f791b8440e 100644
--- a/bundle/schema/jsonschema.json
+++ b/bundle/schema/jsonschema.json
@@ -791,6 +791,51 @@
                   }
                 ]
               },
+              "resources.Volume": {
+                "anyOf": [
+                  {
+                    "type": "object",
+                    "properties": {
+                      "catalog_name": {
+                        "description": "The name of the catalog where the schema and the volume are",
+                        "$ref": "#/$defs/string"
+                      },
+                      "comment": {
+                        "description": "The comment attached to the volume",
+                        "$ref": "#/$defs/string"
+                      },
+                      "grants": {
+                        "$ref": "#/$defs/slice/github.com/databricks/cli/bundle/config/resources.Grant"
+                      },
+                      "name": {
+                        "description": "The name of the volume",
+                        "$ref": "#/$defs/string"
+                      },
+                      "schema_name": {
+                        "description": "The name of the schema where the volume is",
+                        "$ref": "#/$defs/string"
+                      },
+                      "storage_location": {
+                        "description": "The storage location on the cloud",
+                        "$ref": "#/$defs/string"
+                      },
+                      "volume_type": {
+                        "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.VolumeType"
+                      }
+                    },
+                    "additionalProperties": false,
+                    "required": [
+                      "catalog_name",
+                      "name",
+                      "schema_name"
+                    ]
+                  },
+                  {
+                    "type": "string",
+                    "pattern": "\\$\\{(var(\\.[a-zA-Z]+([-_]?[a-zA-Z0-9]+)*(\\[[0-9]+\\])*)+)\\}"
+                  }
+                ]
+              },
               "variable.Lookup": {
                 "anyOf": [
                   {
@@ -963,6 +1008,9 @@
                     },
                     "name": {
                       "$ref": "#/$defs/string"
+                    },
+                    "uuid": {
+                      "$ref": "#/$defs/string"
                     }
                   },
                   "additionalProperties": false,
@@ -1157,6 +1205,9 @@
                     },
                     "schemas": {
                       "$ref": "#/$defs/map/github.com/databricks/cli/bundle/config/resources.Schema"
+                    },
+                    "volumes": {
+                      "$ref": "#/$defs/map/github.com/databricks/cli/bundle/config/resources.Volume"
                     }
                   },
                   "additionalProperties": false
@@ -1558,6 +1609,13 @@
                 }
               ]
             },
+            "catalog.VolumeType": {
+              "type": "string",
+              "enum": [
+                "EXTERNAL",
+                "MANAGED"
+              ]
+            },
             "compute.Adlsgen2Info": {
               "anyOf": [
                 {
@@ -5565,6 +5623,20 @@
                     }
                   ]
                 },
+                "resources.Volume": {
+                  "anyOf": [
+                    {
+                      "type": "object",
+                      "additionalProperties": {
+                        "$ref": "#/$defs/github.com/databricks/cli/bundle/config/resources.Volume"
+                      }
+                    },
+                    {
+                      "type": "string",
+                      "pattern": "\\$\\{(var(\\.[a-zA-Z]+([-_]?[a-zA-Z0-9]+)*(\\[[0-9]+\\])*)+)\\}"
+                    }
+                  ]
+                },
                 "variable.TargetVariable": {
                   "anyOf": [
                     {