Custom annotations for bundle-specific JSON schema fields

.codegen.json

@@ -11,7 +11,7 @@
         "required": ["go"],
         "post_generate": [
             "go test -timeout 240s -run TestConsistentDatabricksSdkVersion github.com/databricks/cli/internal/build",
-            "go run ./bundle/internal/schema/*.go ./bundle/schema/jsonschema.json",
+            "go run ./bundle/internal/schema/*.go ./bundle/internal/schema/annotations.yml ./bundle/schema/jsonschema.json",
             "echo 'bundle/internal/tf/schema/\\*.go linguist-generated=true' >> ./.gitattributes",
             "echo 'go.sum linguist-generated=true' >> ./.gitattributes",
             "echo 'bundle/schema/jsonschema.json linguist-generated=true' >> ./.gitattributes"

bundle/internal/schema/annotations.go

@@ -0,0 +1,109 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"reflect"
+	"strings"
+
+	"github.com/ghodss/yaml"
+
+	"github.com/databricks/cli/libs/jsonschema"
+)
+
+type annotation struct {
+	Description         string `json:"description,omitempty"`
+	MarkdownDescription string `json:"markdown_description,omitempty"`
+	Title               string `json:"title,omitempty"`
+	Default             any    `json:"default,omitempty"`
+}
+
+type annotationHandler struct {
+	filePath string
+	op       *openapiParser
+	ref      map[string]annotation
+}
+
+const Placeholder = "PLACEHOLDER"
+
+// Adds annotations to the JSON schema reading from the annotations file.
+// More details https://json-schema.org/understanding-json-schema/reference/annotations
+func newAnnotationHandler(path string, op *openapiParser) (*annotationHandler, error) {
+	b, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+
+	data := map[string]annotation{}
+	err = yaml.Unmarshal(b, &data)
+	if err != nil {
+		return nil, err
+	}
+
+	if data == nil {
+		data = map[string]annotation{}
+	}
+
+	d := &annotationHandler{}
+	d.ref = data
+	d.op = op
+	d.filePath = path
+	return d, nil
+}
+
+func (d *annotationHandler) addAnnotations(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema {
+	// Skips the type if it is a part of OpenAPI spec, as it is already annotated.
+	_, ok := d.op.findRef(typ)
+	if ok {
+		return s
+	}
+
+	refPath := jsonschema.TypePath(typ)
+	items := map[string]*jsonschema.Schema{}
+	items[refPath] = &s
+
+	for k, v := range s.Properties {
+		itemRefPath := fmt.Sprintf("%s.%s", refPath, k)
+		items[itemRefPath] = v
+	}
+
+	for k, v := range items {
+		// Skipping all default types like int, string, etc.
+		shouldHandle := strings.HasPrefix(refPath, "github.com")
+		if !shouldHandle {
+			continue
+		}
+
+		item, ok := d.ref[k]
+		if !ok {
+			d.ref[k] = annotation{
+				Description: Placeholder,
+			}
+			continue
+		}
+		if item.Description == Placeholder {
+			continue
+		}
+
+		v.Description = item.Description
+		if item.Default != nil {
+			v.Default = item.Default
+		}
+		v.MarkdownDescription = item.MarkdownDescription
+		v.Title = item.Title
+	}
+	return s
+}
+
+func (d *annotationHandler) overwrite() error {
+	data, err := yaml.Marshal(d.ref)
+	if err != nil {
+		return err
+	}
+
+	err = os.WriteFile(d.filePath, data, 0644)
+	if err != nil {
+		return err
+	}
+	return nil
+}