Add openapi descriptions to bundle resources (#229)

This PR:
1. Adds autogeneration of descriptions for `resources` field
2. Autogenerates empty descriptions for any properties in DABs
3. Defines SOPs for how to refresh these descriptions
4. Adds command to generate this documentation
5. Adds Automatically copy any descriptions over to `environments`
property

Basically it provides a framework for adding descriptions to the
generated JSON schema

Tested manually and using unit tests

bundle/schema/README.md

@@ -0,0 +1,26 @@
+### Overview
+
+`docs/bundle_descriptions.json` contains both autogenerated as well as manually written
+descriptions for the json schema. Specifically
+1. `resources` : almost all descriptions are autogenerated from the OpenAPI spec
+2. `environments` : almost all descriptions are copied over from root level entities (eg: `bundle`, `artifacts`)
+3. `bundle` : manually editted
+4. `include` : manually editted
+5. `workspace` : manually editted
+6. `artifacts` : manually editted
+
+These descriptions are rendered in the inline documentation in an IDE
+
+### SOP: Add schema descriptions for new fields in bundle config
+
+1. You can autogenerate empty descriptions for the new fields by running
+`bricks bundle schema --only-docs > ~/bricks/bundle/schema/docs/bundle_descriptions.json`
+2. Manually edit bundle_descriptions.json to add your descriptions
+3. Again run `bricks bundle schema --only-docs > ~/bricks/bundle/schema/docs/bundle_descriptions.json` to copy over any applicable descriptions to `environments`
+4. push to repo
+
+
+### SOP: Update descriptions in resources from a newer openapi spec
+
+1. Run `bricks bundle schema --only-docs --openapi PATH_TO_SPEC > ~/bricks/bundle/schema/docs/bundle_descriptions.json`
+2. push to repo

bundle/schema/bundle_config_docs.yml

@@ -1,12 +0,0 @@
-documentation: Root of the bundle config
-children:
-  bundle:
-    documentation: |
-      Bundle contains details about this bundle, such as its name,
-      version of the spec (TODO), default cluster, default warehouse, etc.
-    children:
-      environment:
-        documentation: Environment is set by the mutator that selects the environment.
-
-  artifacts:
-    documentation: Artifacts contains a description of all code artifacts in this bundle.

bundle/schema/docs.go

@@ -2,37 +2,109 @@ package schema
 
 import (
 	_ "embed"
+	"encoding/json"
+	"fmt"
 	"os"
+	"reflect"
 
-	"gopkg.in/yaml.v3"
+	"github.com/databricks/bricks/bundle/config"
+	"github.com/databricks/databricks-sdk-go/openapi"
 )
 
+// A subset of Schema struct
 type Docs struct {
-	Documentation string          `json:"documentation"`
+	Description          string           `json:"description"`
-	Children      map[string]Docs `json:"children"`
+	Properties           map[string]*Docs `json:"properties,omitempty"`
+	Items                *Docs            `json:"items,omitempty"`
+	AdditionalProperties *Docs            `json:"additionalproperties,omitempty"`
 }
 
-func LoadDocs(path string) (*Docs, error) {
+//go:embed docs/bundle_descriptions.json
-	bytes, err := os.ReadFile(path)
-	if err != nil {
-		return nil, err
-	}
-	docs := Docs{}
-	err = yaml.Unmarshal(bytes, &docs)
-	if err != nil {
-		return nil, err
-	}
-	return &docs, nil
-}
-
-//go:embed bundle_config_docs.yml
 var bundleDocs []byte
 
-func GetBundleDocs() (*Docs, error) {
+func BundleDocs(openapiSpecPath string) (*Docs, error) {
-	docs := Docs{}
+	docs, err := initializeBundleDocs()
-	err := yaml.Unmarshal(bundleDocs, &docs)
 	if err != nil {
 		return nil, err
 	}
-	return &docs, nil
+	if openapiSpecPath != "" {
+		openapiSpec, err := os.ReadFile(openapiSpecPath)
+		if err != nil {
+			return nil, err
+		}
+		spec := &openapi.Specification{}
+		err = json.Unmarshal(openapiSpec, spec)
+		if err != nil {
+			return nil, err
+		}
+		openapiReader := &OpenapiReader{
+			OpenapiSpec: spec,
+			Memo:        make(map[string]*Schema),
+		}
+		resourcesDocs, err := openapiReader.ResourcesDocs()
+		if err != nil {
+			return nil, err
+		}
+		resourceSchema, err := New(reflect.TypeOf(config.Resources{}), resourcesDocs)
+		if err != nil {
+			return nil, err
+		}
+		docs.Properties["resources"] = schemaToDocs(resourceSchema)
+	}
+	docs.refreshEnvironmentsDocs()
+	return docs, nil
+}
+
+func (docs *Docs) refreshEnvironmentsDocs() error {
+	environmentsDocs, ok := docs.Properties["environments"]
+	if !ok || environmentsDocs.AdditionalProperties == nil ||
+		environmentsDocs.AdditionalProperties.Properties == nil {
+		return fmt.Errorf("invalid environments descriptions")
+	}
+	environmentProperties := environmentsDocs.AdditionalProperties.Properties
+	propertiesToCopy := []string{"artifacts", "bundle", "resources", "workspace"}
+	for _, p := range propertiesToCopy {
+		environmentProperties[p] = docs.Properties[p]
+	}
+	return nil
+}
+
+func initializeBundleDocs() (*Docs, error) {
+	// load embedded descriptions
+	embedded := Docs{}
+	err := json.Unmarshal(bundleDocs, &embedded)
+	if err != nil {
+		return nil, err
+	}
+	// generate schema with the embedded descriptions
+	schema, err := New(reflect.TypeOf(config.Root{}), &embedded)
+	if err != nil {
+		return nil, err
+	}
+	// converting the schema back to docs. This creates empty descriptions
+	// for any properties that were missing in the embedded descriptions
+	docs := schemaToDocs(schema)
+	return docs, nil
+}
+
+// *Docs are a subset of *Schema, this function selects that subset
+func schemaToDocs(schema *Schema) *Docs {
+	// terminate recursion if schema is nil
+	if schema == nil {
+		return nil
+	}
+	docs := &Docs{
+		Description: schema.Description,
+	}
+	if len(schema.Properties) > 0 {
+		docs.Properties = make(map[string]*Docs)
+	}
+	for k, v := range schema.Properties {
+		docs.Properties[k] = schemaToDocs(v)
+	}
+	docs.Items = schemaToDocs(schema.Items)
+	if additionalProperties, ok := schema.AdditionalProperties.(*Schema); ok {
+		docs.AdditionalProperties = schemaToDocs(additionalProperties)
+	}
+	return docs
 }

bundle/schema/docs_test.go

@@ -0,0 +1,61 @@
+package schema
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSchemaToDocs(t *testing.T) {
+	schema := &Schema{
+		Type:        "object",
+		Description: "root doc",
+		Properties: map[string]*Schema{
+			"foo": {Type: "number", Description: "foo doc"},
+			"bar": {Type: "string"},
+			"octave": {
+				Type:                 "object",
+				AdditionalProperties: &Schema{Type: "number"},
+				Description:          "octave docs",
+			},
+			"scales": {
+				Type:        "object",
+				Description: "scale docs",
+				Items:       &Schema{Type: "string"},
+			},
+		},
+	}
+	docs := schemaToDocs(schema)
+	docsJson, err := json.MarshalIndent(docs, "		", "	")
+	require.NoError(t, err)
+
+	expected :=
+		`{
+			"description": "root doc",
+			"properties": {
+				"bar": {
+					"description": ""
+				},
+				"foo": {
+					"description": "foo doc"
+				},
+				"octave": {
+					"description": "octave docs",
+					"additionalproperties": {
+						"description": ""
+					}
+				},
+				"scales": {
+					"description": "scale docs",
+					"items": {
+						"description": ""
+					}
+				}
+			}
+		}`
+	t.Log("[DEBUG] actual: ", string(docsJson))
+	t.Log("[DEBUG] expected: ", expected)
+	assert.Equal(t, expected, string(docsJson))
+}

bundle/schema/openapi.go

@@ -0,0 +1,214 @@
+package schema
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"github.com/databricks/databricks-sdk-go/openapi"
+)
+
+type OpenapiReader struct {
+	OpenapiSpec *openapi.Specification
+	Memo        map[string]*Schema
+}
+
+const SchemaPathPrefix = "#/components/schemas/"
+
+func (reader *OpenapiReader) readOpenapiSchema(path string) (*Schema, error) {
+	schemaKey := strings.TrimPrefix(path, SchemaPathPrefix)
+
+	// return early if we already have a computed schema
+	memoSchema, ok := reader.Memo[schemaKey]
+	if ok {
+		return memoSchema, nil
+	}
+
+	// check path is present in openapi spec
+	openapiSchema, ok := reader.OpenapiSpec.Components.Schemas[schemaKey]
+	if !ok {
+		return nil, fmt.Errorf("schema with path %s not found in openapi spec", path)
+	}
+
+	// convert openapi schema to the native schema struct
+	bytes, err := json.Marshal(*openapiSchema)
+	if err != nil {
+		return nil, err
+	}
+	jsonSchema := &Schema{}
+	err = json.Unmarshal(bytes, jsonSchema)
+	if err != nil {
+		return nil, err
+	}
+
+	// A hack to convert a map[string]interface{} to *Schema
+	// We rely on the type of a AdditionalProperties in downstream functions
+	// to do reference interpolation
+	_, ok = jsonSchema.AdditionalProperties.(map[string]interface{})
+	if ok {
+		b, err := json.Marshal(jsonSchema.AdditionalProperties)
+		if err != nil {
+			return nil, err
+		}
+		additionalProperties := &Schema{}
+		err = json.Unmarshal(b, additionalProperties)
+		if err != nil {
+			return nil, err
+		}
+		jsonSchema.AdditionalProperties = additionalProperties
+	}
+
+	// store read schema into memo
+	reader.Memo[schemaKey] = jsonSchema
+
+	return jsonSchema, nil
+}
+
+// safe againt loops in refs
+func (reader *OpenapiReader) safeResolveRefs(root *Schema, seenRefs map[string]struct{}) (*Schema, error) {
+	if root.Reference == nil {
+		return reader.traverseSchema(root, seenRefs)
+	}
+	key := *root.Reference
+	_, ok := seenRefs[key]
+	if ok {
+		// self reference loops can be supported however the logic is non-trivial because
+		// cross refernce loops are not allowed (see: http://json-schema.org/understanding-json-schema/structuring.html#recursion)
+		return nil, fmt.Errorf("references loop detected")
+	}
+	ref := *root.Reference
+	description := root.Description
+	seenRefs[ref] = struct{}{}
+
+	// Mark reference nil, so we do not traverse this again. This is tracked
+	// in the memo
+	root.Reference = nil
+
+	// unroll one level of reference
+	selfRef, err := reader.readOpenapiSchema(ref)
+	if err != nil {
+		return nil, err
+	}
+	root = selfRef
+	root.Description = description
+
+	// traverse again to find new references
+	root, err = reader.traverseSchema(root, seenRefs)
+	if err != nil {
+		return nil, err
+	}
+	delete(seenRefs, ref)
+	return root, err
+}
+
+func (reader *OpenapiReader) traverseSchema(root *Schema, seenRefs map[string]struct{}) (*Schema, error) {
+	// case primitive (or invalid)
+	if root.Type != Object && root.Type != Array {
+		return root, nil
+	}
+	// only root references are resolved
+	if root.Reference != nil {
+		return reader.safeResolveRefs(root, seenRefs)
+	}
+	// case struct
+	if len(root.Properties) > 0 {
+		for k, v := range root.Properties {
+			childSchema, err := reader.safeResolveRefs(v, seenRefs)
+			if err != nil {
+				return nil, err
+			}
+			root.Properties[k] = childSchema
+		}
+	}
+	// case array
+	if root.Items != nil {
+		itemsSchema, err := reader.safeResolveRefs(root.Items, seenRefs)
+		if err != nil {
+			return nil, err
+		}
+		root.Items = itemsSchema
+	}
+	// case map
+	additionionalProperties, ok := root.AdditionalProperties.(*Schema)
+	if ok && additionionalProperties != nil {
+		valueSchema, err := reader.safeResolveRefs(additionionalProperties, seenRefs)
+		if err != nil {
+			return nil, err
+		}
+		root.AdditionalProperties = valueSchema
+	}
+	return root, nil
+}
+
+func (reader *OpenapiReader) readResolvedSchema(path string) (*Schema, error) {
+	root, err := reader.readOpenapiSchema(path)
+	if err != nil {
+		return nil, err
+	}
+	seenRefs := make(map[string]struct{})
+	seenRefs[path] = struct{}{}
+	root, err = reader.safeResolveRefs(root, seenRefs)
+	if err != nil {
+		trace := ""
+		count := 0
+		for k := range seenRefs {
+			if count == len(seenRefs)-1 {
+				trace += k
+				break
+			}
+			trace += k + " -> "
+			count++
+		}
+		return nil, fmt.Errorf("%s. schema ref trace: %s", err, trace)
+	}
+	return root, nil
+}
+
+func (reader *OpenapiReader) jobsDocs() (*Docs, error) {
+	jobSettingsSchema, err := reader.readResolvedSchema(SchemaPathPrefix + "jobs.JobSettings")
+	if err != nil {
+		return nil, err
+	}
+	jobDocs := schemaToDocs(jobSettingsSchema)
+	// TODO: add description for id if needed.
+	// Tracked in https://github.com/databricks/bricks/issues/242
+	jobsDocs := &Docs{
+		Description:          "List of job definations",
+		AdditionalProperties: jobDocs,
+	}
+	return jobsDocs, nil
+}
+
+func (reader *OpenapiReader) pipelinesDocs() (*Docs, error) {
+	pipelineSpecSchema, err := reader.readResolvedSchema(SchemaPathPrefix + "pipelines.PipelineSpec")
+	if err != nil {
+		return nil, err
+	}
+	pipelineDocs := schemaToDocs(pipelineSpecSchema)
+	// TODO: Two fields in resources.Pipeline have the json tag id. Clarify the
+	// semantics and then add a description if needed. (https://github.com/databricks/bricks/issues/242)
+	pipelinesDocs := &Docs{
+		Description:          "List of pipeline definations",
+		AdditionalProperties: pipelineDocs,
+	}
+	return pipelinesDocs, nil
+}
+
+func (reader *OpenapiReader) ResourcesDocs() (*Docs, error) {
+	jobsDocs, err := reader.jobsDocs()
+	if err != nil {
+		return nil, err
+	}
+	pipelinesDocs, err := reader.pipelinesDocs()
+	if err != nil {
+		return nil, err
+	}
+
+	return &Docs{
+		Description: "Specification of databricks resources to instantiate",
+		Properties: map[string]*Docs{
+			"jobs":      jobsDocs,
+			"pipelines": pipelinesDocs,
+		},
+	}, nil
+}

bundle/schema/openapi_test.go

@@ -0,0 +1,435 @@
+package schema
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/openapi"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestReadSchemaForObject(t *testing.T) {
+	specString := `
+	{
+		"components": {
+			"schemas": {
+				"foo": {
+					"type": "number"
+				},
+				"fruits": {
+					"type": "object",
+					"description": "fruits that are cool",
+					"properties": {
+						"guava": {
+							"type": "string",
+							"description": "a guava for my schema"
+						},
+						"mango": {
+							"type": "object",
+							"description": "a mango for my schema",
+							"$ref": "#/components/schemas/mango"
+						}
+					}
+				},
+				"mango": {
+					"type": "object",
+					"properties": {
+						"foo": {
+							"$ref": "#/components/schemas/foo"
+						}
+					}
+				}
+			}
+		}
+	}
+	`
+	spec := &openapi.Specification{}
+	reader := &OpenapiReader{
+		OpenapiSpec: spec,
+		Memo:        make(map[string]*Schema),
+	}
+	err := json.Unmarshal([]byte(specString), spec)
+	require.NoError(t, err)
+
+	fruitsSchema, err := reader.readResolvedSchema("#/components/schemas/fruits")
+	require.NoError(t, err)
+
+	fruitsSchemaJson, err := json.MarshalIndent(fruitsSchema, "		", "	")
+	require.NoError(t, err)
+
+	expected := `{
+			"type": "object",
+			"description": "fruits that are cool",
+			"properties": {
+				"guava": {
+					"type": "string",
+					"description": "a guava for my schema"
+				},
+				"mango": {
+					"type": "object",
+					"description": "a mango for my schema",
+					"properties": {
+						"foo": {
+							"type": "number"
+						}
+					}
+				}
+			}
+		}`
+
+	t.Log("[DEBUG] actual: ", string(fruitsSchemaJson))
+	t.Log("[DEBUG] expected: ", expected)
+	assert.Equal(t, expected, string(fruitsSchemaJson))
+}
+
+func TestReadSchemaForArray(t *testing.T) {
+	specString := `
+	{
+		"components": {
+			"schemas": {
+				"fruits": {
+					"type": "object",
+					"description": "fruits that are cool",
+					"items": {
+						"description": "some papayas, because papayas are fruits too",
+						"$ref": "#/components/schemas/papaya"
+					}
+				},
+				"papaya": {
+					"type": "number"
+				}
+			}
+		}
+	}`
+	spec := &openapi.Specification{}
+	reader := &OpenapiReader{
+		OpenapiSpec: spec,
+		Memo:        make(map[string]*Schema),
+	}
+	err := json.Unmarshal([]byte(specString), spec)
+	require.NoError(t, err)
+
+	fruitsSchema, err := reader.readResolvedSchema("#/components/schemas/fruits")
+	require.NoError(t, err)
+
+	fruitsSchemaJson, err := json.MarshalIndent(fruitsSchema, "		", "	")
+	require.NoError(t, err)
+
+	expected := `{
+			"type": "object",
+			"description": "fruits that are cool",
+			"items": {
+				"type": "number",
+				"description": "some papayas, because papayas are fruits too"
+			}
+		}`
+
+	t.Log("[DEBUG] actual: ", string(fruitsSchemaJson))
+	t.Log("[DEBUG] expected: ", expected)
+	assert.Equal(t, expected, string(fruitsSchemaJson))
+}
+
+func TestReadSchemaForMap(t *testing.T) {
+	specString := `{
+		"components": {
+			"schemas": {
+				"fruits": {
+					"type": "object",
+					"description": "fruits that are meh",
+					"additionalProperties": {
+						"description": "watermelons. watermelons.",
+						"$ref": "#/components/schemas/watermelon"
+					}
+				},
+				"watermelon": {
+					"type": "number"
+				}
+			}
+		}
+	}`
+	spec := &openapi.Specification{}
+	reader := &OpenapiReader{
+		OpenapiSpec: spec,
+		Memo:        make(map[string]*Schema),
+	}
+	err := json.Unmarshal([]byte(specString), spec)
+	require.NoError(t, err)
+
+	fruitsSchema, err := reader.readResolvedSchema("#/components/schemas/fruits")
+	require.NoError(t, err)
+
+	fruitsSchemaJson, err := json.MarshalIndent(fruitsSchema, "		", "	")
+	require.NoError(t, err)
+
+	expected := `{
+			"type": "object",
+			"description": "fruits that are meh",
+			"additionalProperties": {
+				"type": "number",
+				"description": "watermelons. watermelons."
+			}
+		}`
+
+	t.Log("[DEBUG] actual: ", string(fruitsSchemaJson))
+	t.Log("[DEBUG] expected: ", expected)
+	assert.Equal(t, expected, string(fruitsSchemaJson))
+}
+
+func TestRootReferenceIsResolved(t *testing.T) {
+	specString := `{
+		"components": {
+			"schemas": {
+				"foo": {
+					"type": "object",
+					"description": "this description is ignored",
+					"properties": {
+						"abc": {
+							"type": "string"
+						}
+					}
+				},
+				"fruits": {
+					"type": "object",
+					"description": "foo fighters fighting fruits",
+					"$ref": "#/components/schemas/foo"
+				}
+			}
+		}
+	}`
+	spec := &openapi.Specification{}
+	reader := &OpenapiReader{
+		OpenapiSpec: spec,
+		Memo:        make(map[string]*Schema),
+	}
+	err := json.Unmarshal([]byte(specString), spec)
+	require.NoError(t, err)
+
+	schema, err := reader.readResolvedSchema("#/components/schemas/fruits")
+	require.NoError(t, err)
+	fruitsSchemaJson, err := json.MarshalIndent(schema, "		", "	")
+	require.NoError(t, err)
+
+	expected := `{
+			"type": "object",
+			"description": "foo fighters fighting fruits",
+			"properties": {
+				"abc": {
+					"type": "string"
+				}
+			}
+		}`
+
+	t.Log("[DEBUG] actual: ", string(fruitsSchemaJson))
+	t.Log("[DEBUG] expected: ", expected)
+	assert.Equal(t, expected, string(fruitsSchemaJson))
+}
+
+func TestSelfReferenceLoopErrors(t *testing.T) {
+	specString := `{
+		"components": {
+			"schemas": {
+				"foo": {
+					"type": "object",
+					"description": "this description is ignored",
+					"properties": {
+						"bar": {
+							"type": "object",
+							"$ref": "#/components/schemas/foo"
+						}
+					}
+				},
+				"fruits": {
+					"type": "object",
+					"description": "foo fighters fighting fruits",
+					"$ref": "#/components/schemas/foo"
+				}
+			}
+		}
+	}`
+	spec := &openapi.Specification{}
+	reader := &OpenapiReader{
+		OpenapiSpec: spec,
+		Memo:        make(map[string]*Schema),
+	}
+	err := json.Unmarshal([]byte(specString), spec)
+	require.NoError(t, err)
+
+	_, err = reader.readResolvedSchema("#/components/schemas/fruits")
+	assert.ErrorContains(t, err, "references loop detected. schema ref trace: #/components/schemas/fruits -> #/components/schemas/foo")
+}
+
+func TestCrossReferenceLoopErrors(t *testing.T) {
+	specString := `{
+		"components": {
+			"schemas": {
+				"foo": {
+					"type": "object",
+					"description": "this description is ignored",
+					"properties": {
+						"bar": {
+							"type": "object",
+							"$ref": "#/components/schemas/fruits"
+						}
+					}
+				},
+				"fruits": {
+					"type": "object",
+					"description": "foo fighters fighting fruits",
+					"$ref": "#/components/schemas/foo"
+				}
+			}
+		}
+	}`
+	spec := &openapi.Specification{}
+	reader := &OpenapiReader{
+		OpenapiSpec: spec,
+		Memo:        make(map[string]*Schema),
+	}
+	err := json.Unmarshal([]byte(specString), spec)
+	require.NoError(t, err)
+
+	_, err = reader.readResolvedSchema("#/components/schemas/fruits")
+	assert.ErrorContains(t, err, "references loop detected. schema ref trace: #/components/schemas/fruits -> #/components/schemas/foo")
+}
+
+func TestReferenceResolutionForMapInObject(t *testing.T) {
+	specString := `
+	{
+		"components": {
+			"schemas": {
+				"foo": {
+					"type": "number"
+				},
+				"fruits": {
+					"type": "object",
+					"description": "fruits that are cool",
+					"properties": {
+						"guava": {
+							"type": "string",
+							"description": "a guava for my schema"
+						},
+						"mangos": {
+							"type": "object",
+							"description": "multiple mangos",
+							"$ref": "#/components/schemas/mango"
+						}
+					}
+				},
+				"mango": {
+					"type": "object",
+					"additionalProperties": {
+						"description": "a single mango",
+						"$ref": "#/components/schemas/foo"
+					}
+				}
+			}
+		}
+	}`
+	spec := &openapi.Specification{}
+	reader := &OpenapiReader{
+		OpenapiSpec: spec,
+		Memo:        make(map[string]*Schema),
+	}
+	err := json.Unmarshal([]byte(specString), spec)
+	require.NoError(t, err)
+
+	fruitsSchema, err := reader.readResolvedSchema("#/components/schemas/fruits")
+	require.NoError(t, err)
+
+	fruitsSchemaJson, err := json.MarshalIndent(fruitsSchema, "		", "	")
+	require.NoError(t, err)
+
+	expected := `{
+			"type": "object",
+			"description": "fruits that are cool",
+			"properties": {
+				"guava": {
+					"type": "string",
+					"description": "a guava for my schema"
+				},
+				"mangos": {
+					"type": "object",
+					"description": "multiple mangos",
+					"additionalProperties": {
+						"type": "number",
+						"description": "a single mango"
+					}
+				}
+			}
+		}`
+
+	t.Log("[DEBUG] actual: ", string(fruitsSchemaJson))
+	t.Log("[DEBUG] expected: ", expected)
+	assert.Equal(t, expected, string(fruitsSchemaJson))
+}
+
+func TestReferenceResolutionForArrayInObject(t *testing.T) {
+	specString := `{
+		"components": {
+			"schemas": {
+				"foo": {
+					"type": "number"
+				},
+				"fruits": {
+					"type": "object",
+					"description": "fruits that are cool",
+					"properties": {
+						"guava": {
+							"type": "string",
+							"description": "a guava for my schema"
+						},
+						"mangos": {
+							"type": "object",
+							"description": "multiple mangos",
+							"$ref": "#/components/schemas/mango"
+						}
+					}
+				},
+				"mango": {
+					"type": "object",
+					"items": {
+						"description": "a single mango",
+						"$ref": "#/components/schemas/foo"
+					}
+				}
+			}
+		}
+	}`
+	spec := &openapi.Specification{}
+	reader := &OpenapiReader{
+		OpenapiSpec: spec,
+		Memo:        make(map[string]*Schema),
+	}
+	err := json.Unmarshal([]byte(specString), spec)
+	require.NoError(t, err)
+
+	fruitsSchema, err := reader.readResolvedSchema("#/components/schemas/fruits")
+	require.NoError(t, err)
+
+	fruitsSchemaJson, err := json.MarshalIndent(fruitsSchema, "		", "	")
+	require.NoError(t, err)
+
+	expected := `{
+			"type": "object",
+			"description": "fruits that are cool",
+			"properties": {
+				"guava": {
+					"type": "string",
+					"description": "a guava for my schema"
+				},
+				"mangos": {
+					"type": "object",
+					"description": "multiple mangos",
+					"items": {
+						"type": "number",
+						"description": "a single mango"
+					}
+				}
+			}
+		}`
+
+	t.Log("[DEBUG] actual: ", string(fruitsSchemaJson))
+	t.Log("[DEBUG] expected: ", expected)
+	assert.Equal(t, expected, string(fruitsSchemaJson))
+}

bundle/schema/schema.go

@@ -35,6 +35,9 @@ type Schema struct {
 	// Required properties for the object. Any fields missing the "omitempty"
 	// json tag will be included
 	Required []string `json:"required,omitempty"`
+
+	// URI to a json schema
+	Reference *string `json:"$ref,omitempty"`
 }
 
 // This function translates golang types into json schema. Here is the mapping
@@ -187,7 +190,7 @@ func toSchema(golangType reflect.Type, docs *Docs, tracker *tracker) (*Schema, e
 	schema := &Schema{Type: rootJavascriptType}
 
 	if docs != nil {
-		schema.Description = docs.Documentation
+		schema.Description = docs.Description
 	}
 
 	// case array/slice
@@ -197,7 +200,11 @@ func toSchema(golangType reflect.Type, docs *Docs, tracker *tracker) (*Schema, e
 		if err != nil {
 			return nil, err
 		}
-		elemProps, err := safeToSchema(elemGolangType, docs, "", tracker)
+		var childDocs *Docs
+		if docs != nil {
+			childDocs = docs.Items
+		}
+		elemProps, err := safeToSchema(elemGolangType, childDocs, "", tracker)
 		if err != nil {
 			return nil, err
 		}
@@ -215,7 +222,11 @@ func toSchema(golangType reflect.Type, docs *Docs, tracker *tracker) (*Schema, e
 		if golangType.Key().Kind() != reflect.String {
 			return nil, fmt.Errorf("only string keyed maps allowed")
 		}
-		schema.AdditionalProperties, err = safeToSchema(golangType.Elem(), docs, "", tracker)
+		var childDocs *Docs
+		if docs != nil {
+			childDocs = docs.AdditionalProperties
+		}
+		schema.AdditionalProperties, err = safeToSchema(golangType.Elem(), childDocs, "", tracker)
 		if err != nil {
 			return nil, err
 		}
@@ -240,8 +251,8 @@ func toSchema(golangType reflect.Type, docs *Docs, tracker *tracker) (*Schema, e
 			// get docs for the child if they exist
 			var childDocs *Docs
 			if docs != nil {
-				if val, ok := docs.Children[childName]; ok {
+				if val, ok := docs.Properties[childName]; ok {
-					childDocs = &val
+					childDocs = val
 				}
 			}
 

bundle/schema/schema_test.go

@@ -1062,60 +1062,31 @@ func TestFieldsWithoutOmitEmptyAreRequired(t *testing.T) {
 	assert.Equal(t, expectedSchema, string(jsonSchema))
 }
 
-func TestDocIngestionInSchema(t *testing.T) {
+func TestDocIngestionForObject(t *testing.T) {
 	docs := &Docs{
-		Documentation: "docs for root",
+		Description: "docs for root",
-		Children: map[string]Docs{
+		Properties: map[string]*Docs{
 			"my_struct": {
-				Documentation: "docs for my struct",
+				Description: "docs for my struct",
-			},
+				Properties: map[string]*Docs{
-			"my_val": {
+					"a": {
-				Documentation: "docs for my val",
+						Description: "docs for a",
-			},
-			"my_slice": {
-				Documentation: "docs for my slice",
-				Children: map[string]Docs{
-					"guava": {
-						Documentation: "docs for guava",
 					},
-					"pineapple": {
+					"c": {
-						Documentation: "docs for pineapple",
+						Description: "docs for c which does not exist on my_struct",
-					},
-				},
-			},
-			"my_map": {
-				Documentation: "docs for my map",
-				Children: map[string]Docs{
-					"apple": {
-						Documentation: "docs for apple",
-					},
-					"mango": {
-						Documentation: "docs for mango",
 					},
 				},
 			},
 		},
 	}
 
-	type Foo struct {
-		Apple int `json:"apple"`
-		Mango int `json:"mango"`
-	}
-
-	type Bar struct {
-		Guava     int `json:"guava"`
-		Pineapple int `json:"pineapple"`
-	}
-
 	type MyStruct struct {
 		A string `json:"a"`
+		B int    `json:"b"`
 	}
 
 	type Root struct {
-		MyStruct *MyStruct       `json:"my_struct"`
+		MyStruct *MyStruct `json:"my_struct"`
-		MyVal    int             `json:"my_val"`
-		MySlice  []Bar           `json:"my_slice"`
-		MyMap    map[string]*Foo `json:"my_map"`
 	}
 
 	elem := Root{}
@@ -1131,29 +1102,82 @@ func TestDocIngestionInSchema(t *testing.T) {
 			"type": "object",
 			"description": "docs for root",
 			"properties": {
-				"my_map": {
+				"my_struct": {
 					"type": "object",
-					"description": "docs for my map",
+					"description": "docs for my struct",
-					"additionalProperties": {
+					"properties": {
-						"type": "object",
+						"a": {
-						"description": "docs for my map",
+							"type": "string",
-						"properties": {
+							"description": "docs for a"
-							"apple": {
-								"type": "number",
-								"description": "docs for apple"
-							},
-							"mango": {
-								"type": "number",
-								"description": "docs for mango"
-							}
 						},
-						"additionalProperties": false,
+						"b": {
-						"required": [
+							"type": "number"
-							"apple",
+						}
-							"mango"
+					},
-						]
+					"additionalProperties": false,
-					}
+					"required": [
+						"a",
+						"b"
+					]
+				}
+			},
+			"additionalProperties": false,
+			"required": [
+				"my_struct"
+			]
+		}`
+
+	t.Log("[DEBUG] actual: ", string(jsonSchema))
+	t.Log("[DEBUG] expected: ", expectedSchema)
+
+	assert.Equal(t, expectedSchema, string(jsonSchema))
+}
+
+func TestDocIngestionForSlice(t *testing.T) {
+	docs := &Docs{
+		Description: "docs for root",
+		Properties: map[string]*Docs{
+			"my_slice": {
+				Description: "docs for my slice",
+				Items: &Docs{
+					Properties: map[string]*Docs{
+						"guava": {
+							Description: "docs for guava",
+						},
+						"pineapple": {
+							Description: "docs for pineapple",
+						},
+						"watermelon": {
+							Description: "docs for watermelon which does not exist in schema",
+						},
+					},
 				},
+			},
+		},
+	}
+
+	type Bar struct {
+		Guava     int `json:"guava"`
+		Pineapple int `json:"pineapple"`
+	}
+
+	type Root struct {
+		MySlice []Bar `json:"my_slice"`
+	}
+
+	elem := Root{}
+
+	schema, err := New(reflect.TypeOf(elem), docs)
+	require.NoError(t, err)
+
+	jsonSchema, err := json.MarshalIndent(schema, "		", "	")
+	assert.NoError(t, err)
+
+	expectedSchema :=
+		`{
+			"type": "object",
+			"description": "docs for root",
+			"properties": {
 				"my_slice": {
 					"type": "array",
 					"description": "docs for my slice",
@@ -1175,20 +1199,130 @@ func TestDocIngestionInSchema(t *testing.T) {
 							"pineapple"
 						]
 					}
-				},
+				}
-				"my_struct": {
+			},
-					"type": "object",
+			"additionalProperties": false,
-					"description": "docs for my struct",
+			"required": [
-					"properties": {
+				"my_slice"
-						"a": {
+			]
-							"type": "string"
+		}`
-						}
+
+	t.Log("[DEBUG] actual: ", string(jsonSchema))
+	t.Log("[DEBUG] expected: ", expectedSchema)
+
+	assert.Equal(t, expectedSchema, string(jsonSchema))
+}
+
+func TestDocIngestionForMap(t *testing.T) {
+	docs := &Docs{
+		Description: "docs for root",
+		Properties: map[string]*Docs{
+			"my_map": {
+				Description: "docs for my map",
+				AdditionalProperties: &Docs{
+					Properties: map[string]*Docs{
+						"apple": {
+							Description: "docs for apple",
+						},
+						"mango": {
+							Description: "docs for mango",
+						},
+						"watermelon": {
+							Description: "docs for watermelon which does not exist in schema",
+						},
+						"papaya": {
+							Description: "docs for papaya which does not exist in schema",
+						},
 					},
-					"additionalProperties": false,
-					"required": [
-						"a"
-					]
 				},
+			},
+		},
+	}
+
+	type Foo struct {
+		Apple int `json:"apple"`
+		Mango int `json:"mango"`
+	}
+
+	type Root struct {
+		MyMap map[string]*Foo `json:"my_map"`
+	}
+
+	elem := Root{}
+
+	schema, err := New(reflect.TypeOf(elem), docs)
+	require.NoError(t, err)
+
+	jsonSchema, err := json.MarshalIndent(schema, "		", "	")
+	assert.NoError(t, err)
+
+	expectedSchema :=
+		`{
+			"type": "object",
+			"description": "docs for root",
+			"properties": {
+				"my_map": {
+					"type": "object",
+					"description": "docs for my map",
+					"additionalProperties": {
+						"type": "object",
+						"properties": {
+							"apple": {
+								"type": "number",
+								"description": "docs for apple"
+							},
+							"mango": {
+								"type": "number",
+								"description": "docs for mango"
+							}
+						},
+						"additionalProperties": false,
+						"required": [
+							"apple",
+							"mango"
+						]
+					}
+				}
+			},
+			"additionalProperties": false,
+			"required": [
+				"my_map"
+			]
+		}`
+
+	t.Log("[DEBUG] actual: ", string(jsonSchema))
+	t.Log("[DEBUG] expected: ", expectedSchema)
+
+	assert.Equal(t, expectedSchema, string(jsonSchema))
+}
+
+func TestDocIngestionForTopLevelPrimitive(t *testing.T) {
+	docs := &Docs{
+		Description: "docs for root",
+		Properties: map[string]*Docs{
+			"my_val": {
+				Description: "docs for my val",
+			},
+		},
+	}
+
+	type Root struct {
+		MyVal int `json:"my_val"`
+	}
+
+	elem := Root{}
+
+	schema, err := New(reflect.TypeOf(elem), docs)
+	require.NoError(t, err)
+
+	jsonSchema, err := json.MarshalIndent(schema, "		", "	")
+	assert.NoError(t, err)
+
+	expectedSchema :=
+		`{
+			"type": "object",
+			"description": "docs for root",
+			"properties": {
 				"my_val": {
 					"type": "number",
 					"description": "docs for my val"
@@ -1196,10 +1330,7 @@ func TestDocIngestionInSchema(t *testing.T) {
 			},
 			"additionalProperties": false,
 			"required": [
-				"my_struct",
+				"my_val"
-				"my_val",
-				"my_slice",
-				"my_map"
 			]
 		}`
 

cmd/bundle/schema.go

@@ -14,7 +14,7 @@ var schemaCmd = &cobra.Command{
 	Short: "Generate JSON Schema for bundle configuration",
 
 	RunE: func(cmd *cobra.Command, args []string) error {
-		docs, err := schema.GetBundleDocs()
+		docs, err := schema.BundleDocs(openapi)
 		if err != nil {
 			return err
 		}
@@ -22,15 +22,26 @@ var schemaCmd = &cobra.Command{
 		if err != nil {
 			return err
 		}
-		jsonSchema, err := json.MarshalIndent(schema, "", "  ")
+		result, err := json.MarshalIndent(schema, "", "  ")
 		if err != nil {
 			return err
 		}
-		cmd.OutOrStdout().Write(jsonSchema)
+		if onlyDocs {
+			result, err = json.MarshalIndent(docs, "", "  ")
+			if err != nil {
+				return err
+			}
+		}
+		cmd.OutOrStdout().Write(result)
 		return nil
 	},
 }
 
+var openapi string
+var onlyDocs bool
+
 func init() {
 	AddCommand(schemaCmd)
+	schemaCmd.Flags().StringVar(&openapi, "openapi", "", "path to a databricks openapi spec")
+	schemaCmd.Flags().BoolVar(&onlyDocs, "only-docs", false, "only generate descriptions for the schema")
 }

go.mod

@@ -62,5 +62,5 @@ require (
 	google.golang.org/grpc v1.53.0 // indirect
 	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )