chore: Extract annotation package

This commit is contained in:
Ilya Kuznetsov 2025-01-02 16:40:06 +01:00
parent 820dd5f34c
commit 460a4558e0
No known key found for this signature in database
GPG Key ID: 91F3DDCF5D21CDDF
6 changed files with 88 additions and 158 deletions

View File

@ -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"

View File

@ -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/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/bundle/internal/annotation"
"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) {
if a.Description != "" && a.Description != Placeholder {
func assignAnnotation(s *jsonschema.Schema, a annotation.Descriptor) {
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
}

View File

@ -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
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)
data, err := annotation.LoadAndMerge(sources)
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) {
if a.Description != Placeholder {
func assignAnnotation(s *jsonschema.Schema, a annotation.Descriptor) {
if a.Description != annotation.Placeholder {
s.Description = a.Description
}

View File

@ -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
: "private_key":
github.com/databricks/databricks-sdk-go/service/serving.GoogleCloudVertexAiConfig:
"private_key":
"description": |-
PLACEHOLDER
"private_key_plaintext":

View File

@ -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
}

View File

@ -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{}
overrides := annotationFile{}
annotations := annotation.File{}
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
}