chore: Extract annotation package

bundle/internal/annotation/main.go

@@ -0,0 +1,56 @@
+package annotation
+
+import (
+	"bytes"
+	"os"
+
+	"github.com/databricks/cli/libs/dyn"
+	"github.com/databricks/cli/libs/dyn/convert"
+	"github.com/databricks/cli/libs/dyn/merge"
+	"github.com/databricks/cli/libs/dyn/yamlloader"
+)
+
+type Descriptor struct {
+	Description         string `json:"description,omitempty"`
+	MarkdownDescription string `json:"markdown_description,omitempty"`
+	Title               string `json:"title,omitempty"`
+	Default             any    `json:"default,omitempty"`
+	Enum                []any  `json:"enum,omitempty"`
+	MarkdownExamples    string `json:"markdown_examples,omitempty"`
+}
+
+/**
+ * Parsed file with annotations, expected format:
+ * github.com/databricks/cli/bundle/config.Bundle:
+ *  	cluster_id:
+ *      description: "Description"
+ */
+type File map[string]map[string]Descriptor
+
+func LoadAndMerge(sources []string) (File, error) {
+	prev := dyn.NilValue
+	for _, path := range sources {
+		b, err := os.ReadFile(path)
+		if err != nil {
+			return nil, err
+		}
+		generated, err := yamlloader.LoadYAML(path, bytes.NewBuffer(b))
+		if err != nil {
+			return nil, err
+		}
+		prev, err = merge.Merge(prev, generated)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	var data File
+
+	err := convert.ToTyped(&data, prev)
+	if err != nil {
+		return nil, err
+	}
+	return data, nil
+}
+
+const Placeholder = "PLACEHOLDER"

bundle/internal/docs/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"bytes"
 	"fmt"
 	"log"
 	"os"
@@ -10,60 +9,10 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/bundle/config"
-	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/cli/bundle/internal/annotation"
-	"github.com/databricks/cli/libs/dyn"
-	"github.com/databricks/cli/libs/dyn/convert"
-	"github.com/databricks/cli/libs/dyn/merge"
-	"github.com/databricks/cli/libs/dyn/yamlloader"
 	"github.com/databricks/cli/libs/jsonschema"
-	"github.com/databricks/databricks-sdk-go/service/jobs"
 )
 
-const Placeholder = "PLACEHOLDER"
-
-func removeJobsFields(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema {
-	switch typ {
-	case reflect.TypeOf(resources.Job{}):
-		// This field has been deprecated in jobs API v2.1 and is always set to
-		// "MULTI_TASK" in the backend. We should not expose it to the user.
-		delete(s.Properties, "format")
-
-		// These fields are only meant to be set by the DABs client (ie the CLI)
-		// and thus should not be exposed to the user. These are used to annotate
-		// jobs that were created by DABs.
-		delete(s.Properties, "deployment")
-		delete(s.Properties, "edit_mode")
-
-	case reflect.TypeOf(jobs.GitSource{}):
-		// These fields are readonly and are not meant to be set by the user.
-		delete(s.Properties, "job_source")
-		delete(s.Properties, "git_snapshot")
-
-	default:
-		// Do nothing
-	}
-
-	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) != 3 {
 		fmt.Println("Usage: go run main.go <annotation-file> <output-file>")
@@ -79,21 +28,10 @@ func main() {
 	}
 }
 
-type annotationFile map[string]map[string]annotation
-
-type annotation struct {
-	Description         string `json:"description,omitempty"`
-	MarkdownDescription string `json:"markdown_description,omitempty"`
-	Title               string `json:"title,omitempty"`
-	Default             any    `json:"default,omitempty"`
-	Enum                []any  `json:"enum,omitempty"`
-	MarkdownExamples    string `json:"markdown_examples,omitempty"`
-}
-
 func generateDocs(workdir, outputPath string) error {
 	annotationsPath := filepath.Join(workdir, "annotations.yml")
 
-	annotations, err := LoadAndMergeAnnotations([]string{annotationsPath})
+	annotations, err := annotation.LoadAndMerge([]string{annotationsPath})
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -102,9 +40,6 @@ func generateDocs(workdir, outputPath string) error {
 	customFields := map[string]bool{}
 
 	s, err := jsonschema.FromType(reflect.TypeOf(config.Root{}), []func(reflect.Type, jsonschema.Schema) jsonschema.Schema{
-		removeJobsFields,
-		makeVolumeTypeOptional,
-
 		func(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema {
 			_, isCustomField := annotations[jsonschema.TypePath(typ)]
 			if isCustomField {
@@ -120,7 +55,7 @@ func generateDocs(workdir, outputPath string) error {
 
 			a := annotations[refPath]
 			if a == nil {
-				a = map[string]annotation{}
+				a = map[string]annotation.Descriptor{}
 			}
 
 			rootTypeAnnotation, ok := a["_"]
@@ -151,8 +86,8 @@ func getPath(typ reflect.Type) string {
 	return typ.PkgPath() + "." + typ.Name()
 }
 
-func assignAnnotation(s *jsonschema.Schema, a annotation) {
+func assignAnnotation(s *jsonschema.Schema, a annotation.Descriptor) {
-	if a.Description != "" && a.Description != Placeholder {
+	if a.Description != "" && a.Description != annotation.Placeholder {
 		s.Description = a.Description
 	}
 	if a.MarkdownDescription != "" {
@@ -162,29 +97,3 @@ func assignAnnotation(s *jsonschema.Schema, a annotation) {
 		s.Examples = []any{a.MarkdownExamples}
 	}
 }
-
-func LoadAndMergeAnnotations(sources []string) (annotationFile, error) {
-	prev := dyn.NilValue
-	for _, path := range sources {
-		b, err := os.ReadFile(path)
-		if err != nil {
-			return nil, err
-		}
-		generated, err := yamlloader.LoadYAML(path, bytes.NewBuffer(b))
-		if err != nil {
-			return nil, err
-		}
-		prev, err = merge.Merge(prev, generated)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	var data annotationFile
-
-	err := convert.ToTyped(&data, prev)
-	if err != nil {
-		return nil, err
-	}
-	return data, nil
-}

bundle/internal/schema/annotations.go

@@ -10,6 +10,7 @@ import (
 
 	yaml3 "gopkg.in/yaml.v3"
 
+	"github.com/databricks/cli/bundle/internal/annotation"
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/dyn/convert"
 	"github.com/databricks/cli/libs/dyn/merge"
@@ -18,61 +19,23 @@ import (
 	"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"`
-	Enum                []any  `json:"enum,omitempty"`
-	MarkdownExamples    string `json:"markdown_examples,omitempty"`
-}
-
 type annotationHandler struct {
 	// Annotations read from all annotation files including all overrides
-	parsedAnnotations annotationFile
+	parsedAnnotations annotation.File
 	// Missing annotations for fields that are found in config that need to be added to the annotation file
-	missingAnnotations annotationFile
+	missingAnnotations annotation.File
 }
 
-/**
- * Parsed file with annotations, expected format:
- * github.com/databricks/cli/bundle/config.Bundle:
- *  	cluster_id:
- *      description: "Description"
- */
-type annotationFile map[string]map[string]annotation
-
-const Placeholder = "PLACEHOLDER"
-
 // Adds annotations to the JSON schema reading from the annotation files.
 // More details https://json-schema.org/understanding-json-schema/reference/annotations
 func newAnnotationHandler(sources []string) (*annotationHandler, error) {
-	prev := dyn.NilValue
+	data, err := annotation.LoadAndMerge(sources)
-	for _, path := range sources {
-		b, err := os.ReadFile(path)
 	if err != nil {
 		return nil, err
 	}
-		generated, err := yamlloader.LoadYAML(path, bytes.NewBuffer(b))
-		if err != nil {
-			return nil, err
-		}
-		prev, err = merge.Merge(prev, generated)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	var data annotationFile
-
-	err := convert.ToTyped(&data, prev)
-	if err != nil {
-		return nil, err
-	}
-
 	d := &annotationHandler{}
 	d.parsedAnnotations = data
-	d.missingAnnotations = annotationFile{}
+	d.missingAnnotations = annotation.File{}
 	return d, nil
 }
 
@@ -85,7 +48,7 @@ func (d *annotationHandler) addAnnotations(typ reflect.Type, s jsonschema.Schema
 
 	annotations := d.parsedAnnotations[refPath]
 	if annotations == nil {
-		annotations = map[string]annotation{}
+		annotations = map[string]annotation.Descriptor{}
 	}
 
 	rootTypeAnnotation, ok := annotations[RootTypeKey]
@@ -96,11 +59,11 @@ func (d *annotationHandler) addAnnotations(typ reflect.Type, s jsonschema.Schema
 	for k, v := range s.Properties {
 		item := annotations[k]
 		if item.Description == "" {
-			item.Description = Placeholder
+			item.Description = annotation.Placeholder
 
 			emptyAnnotations := d.missingAnnotations[refPath]
 			if emptyAnnotations == nil {
-				emptyAnnotations = map[string]annotation{}
+				emptyAnnotations = map[string]annotation.Descriptor{}
 				d.missingAnnotations[refPath] = emptyAnnotations
 			}
 			emptyAnnotations[k] = item
@@ -141,8 +104,8 @@ func getPath(typ reflect.Type) string {
 	return typ.PkgPath() + "." + typ.Name()
 }
 
-func assignAnnotation(s *jsonschema.Schema, a annotation) {
+func assignAnnotation(s *jsonschema.Schema, a annotation.Descriptor) {
-	if a.Description != Placeholder {
+	if a.Description != annotation.Placeholder {
 		s.Description = a.Description
 	}
 

bundle/internal/schema/annotations.yml

@@ -463,8 +463,8 @@ github.com/databricks/databricks-sdk-go/service/serving.Ai21LabsConfig:
   "ai21labs_api_key_plaintext":
     "description": |-
       PLACEHOLDER
-? github.com/databricks/databricks-sdk-go/service/serving.GoogleCloudVertexAiConfig
+github.com/databricks/databricks-sdk-go/service/serving.GoogleCloudVertexAiConfig:
-: "private_key":
+  "private_key":
     "description": |-
       PLACEHOLDER
   "private_key_plaintext":

bundle/internal/schema/main_test.go

@@ -11,6 +11,7 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/bundle/config"
+	"github.com/databricks/cli/bundle/internal/annotation"
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/dyn/merge"
 	"github.com/databricks/cli/libs/dyn/yamlloader"
@@ -113,13 +114,13 @@ func TestNoDetachedAnnotations(t *testing.T) {
 	assert.Empty(t, types, "Detached annotations found, regenerate schema and check for package path changes")
 }
 
-func getAnnotations(path string) (annotationFile, error) {
+func getAnnotations(path string) (annotation.File, error) {
 	b, err := os.ReadFile(path)
 	if err != nil {
 		return nil, err
 	}
 
-	var data annotationFile
+	var data annotation.File
 	err = yaml.Unmarshal(b, &data)
 	return data, err
 }

bundle/internal/schema/parser.go

@@ -9,6 +9,7 @@ import (
 	"reflect"
 	"strings"
 
+	"github.com/databricks/cli/bundle/internal/annotation"
 	"github.com/databricks/cli/libs/dyn/yamlloader"
 	"github.com/databricks/cli/libs/jsonschema"
 	"gopkg.in/yaml.v3"
@@ -95,8 +96,8 @@ func (p *openapiParser) findRef(typ reflect.Type) (jsonschema.Schema, bool) {
 
 // Use the OpenAPI spec to load descriptions for the given type.
 func (p *openapiParser) extractAnnotations(typ reflect.Type, outputPath, overridesPath string) error {
-	annotations := annotationFile{}
+	annotations := annotation.File{}
-	overrides := annotationFile{}
+	overrides := annotation.File{}
 
 	b, err := os.ReadFile(overridesPath)
 	if err != nil {
@@ -107,7 +108,7 @@ func (p *openapiParser) extractAnnotations(typ reflect.Type, outputPath, overrid
 		return err
 	}
 	if overrides == nil {
-		overrides = annotationFile{}
+		overrides = annotation.File{}
 	}
 
 	_, err = jsonschema.FromType(typ, []func(reflect.Type, jsonschema.Schema) jsonschema.Schema{
@@ -118,16 +119,16 @@ func (p *openapiParser) extractAnnotations(typ reflect.Type, outputPath, overrid
 			}
 
 			basePath := getPath(typ)
-			pkg := map[string]annotation{}
+			pkg := map[string]annotation.Descriptor{}
 			annotations[basePath] = pkg
 
 			if ref.Description != "" || ref.Enum != nil {
-				pkg[RootTypeKey] = annotation{Description: ref.Description, Enum: ref.Enum}
+				pkg[RootTypeKey] = annotation.Descriptor{Description: ref.Description, Enum: ref.Enum}
 			}
 
 			for k := range s.Properties {
 				if refProp, ok := ref.Properties[k]; ok {
-					pkg[k] = annotation{Description: refProp.Description, Enum: refProp.Enum}
+					pkg[k] = annotation.Descriptor{Description: refProp.Description, Enum: refProp.Enum}
 					if refProp.Description == "" {
 						addEmptyOverride(k, basePath, overrides)
 					}
@@ -167,22 +168,22 @@ func (p *openapiParser) extractAnnotations(typ reflect.Type, outputPath, overrid
 	return nil
 }
 
-func addEmptyOverride(key, pkg string, overridesFile annotationFile) {
+func addEmptyOverride(key, pkg string, overridesFile annotation.File) {
 	if overridesFile[pkg] == nil {
-		overridesFile[pkg] = map[string]annotation{}
+		overridesFile[pkg] = map[string]annotation.Descriptor{}
 	}
 
 	overrides := overridesFile[pkg]
 	if overrides[key].Description == "" {
-		overrides[key] = annotation{Description: Placeholder}
+		overrides[key] = annotation.Descriptor{Description: annotation.Placeholder}
 	}
 
 	a, ok := overrides[key]
 	if !ok {
-		a = annotation{}
+		a = annotation.Descriptor{}
 	}
 	if a.Description == "" {
-		a.Description = Placeholder
+		a.Description = annotation.Placeholder
 	}
 	overrides[key] = a
 }