From 26466d23b435bd6a60833b247a53ce0bcda2105f Mon Sep 17 00:00:00 2001
From: Ilya Kuznetsov <ilya.kuznetsov@databricks.com>
Date: Wed, 15 Jan 2025 19:47:39 +0100
Subject: [PATCH] fix: Remove type inconsistency in maps

---
 bundle/internal/docs/main.go  |  6 +++---
 bundle/internal/docs/nodes.go |  6 +++---
 bundle/internal/docs/refs.go  | 15 ++++++++-------
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/bundle/internal/docs/main.go b/bundle/internal/docs/main.go
index 0f00565c6..ed2897687 100644
--- a/bundle/internal/docs/main.go
+++ b/bundle/internal/docs/main.go
@@ -69,7 +69,7 @@ func generateDocs(inputPaths []string, outputPath string, rootType reflect.Type,
 		log.Fatal(err)
 	}
 
-	schemas := map[string]jsonschema.Schema{}
+	schemas := map[string]*jsonschema.Schema{}
 	customFields := map[string]bool{}
 
 	s, err := jsonschema.FromType(rootType, []func(reflect.Type, jsonschema.Schema) jsonschema.Schema{
@@ -82,7 +82,7 @@ func generateDocs(inputPaths []string, outputPath string, rootType reflect.Type,
 			refPath := getPath(typ)
 			shouldHandle := strings.HasPrefix(refPath, "github.com")
 			if !shouldHandle {
-				schemas[jsonschema.TypePath(typ)] = s
+				schemas[jsonschema.TypePath(typ)] = &s
 				return s
 			}
 
@@ -100,7 +100,7 @@ func generateDocs(inputPaths []string, outputPath string, rootType reflect.Type,
 				assignAnnotation(v, a[k])
 			}
 
-			schemas[jsonschema.TypePath(typ)] = s
+			schemas[jsonschema.TypePath(typ)] = &s
 			return s
 		},
 	})
diff --git a/bundle/internal/docs/nodes.go b/bundle/internal/docs/nodes.go
index 9d9bbf8b5..4771dd0b2 100644
--- a/bundle/internal/docs/nodes.go
+++ b/bundle/internal/docs/nodes.go
@@ -42,7 +42,7 @@ type rootProp struct {
 
 const MapType = "Map"
 
-func getNodes(s jsonschema.Schema, refs map[string]jsonschema.Schema, customFields map[string]bool) []rootNode {
+func getNodes(s jsonschema.Schema, refs map[string]*jsonschema.Schema, customFields map[string]bool) []rootNode {
 	rootProps := []rootProp{}
 	for k, v := range s.Properties {
 		rootProps = append(rootProps, rootProp{k, v, true, false})
@@ -125,7 +125,7 @@ func getHumanReadableType(t jsonschema.Type) string {
 	return typesMapping[string(t)]
 }
 
-func getAttributes(props map[string]*jsonschema.Schema, refs map[string]jsonschema.Schema, customFields map[string]bool, prefix string, circular bool) []attributeNode {
+func getAttributes(props, refs map[string]*jsonschema.Schema, customFields map[string]bool, prefix string, circular bool) []attributeNode {
 	attributes := []attributeNode{}
 	for k, v := range props {
 		v = resolveRefs(v, refs)
@@ -165,7 +165,7 @@ func shouldExtract(ref string, customFields map[string]bool) bool {
 	return isCustomField
 }
 
-func extractNodes(prefix string, props map[string]*jsonschema.Schema, refs map[string]jsonschema.Schema, customFields map[string]bool) []rootProp {
+func extractNodes(prefix string, props, refs map[string]*jsonschema.Schema, customFields map[string]bool) []rootProp {
 	nodes := []rootProp{}
 	for k, v := range props {
 		if !shouldExtract(*v.Reference, customFields) {
diff --git a/bundle/internal/docs/refs.go b/bundle/internal/docs/refs.go
index 7d9f6ff03..2cc613fbb 100644
--- a/bundle/internal/docs/refs.go
+++ b/bundle/internal/docs/refs.go
@@ -7,7 +7,7 @@ import (
 	"github.com/databricks/cli/libs/jsonschema"
 )
 
-func isReferenceType(v *jsonschema.Schema, refs map[string]jsonschema.Schema, customFields map[string]bool) bool {
+func isReferenceType(v *jsonschema.Schema, refs map[string]*jsonschema.Schema, customFields map[string]bool) bool {
 	if v.Type != "object" && v.Type != "array" {
 		return false
 	}
@@ -50,7 +50,7 @@ func resolveAdditionalProperties(v *jsonschema.Schema) *jsonschema.Schema {
 	return additionalProps
 }
 
-func resolveRefs(s *jsonschema.Schema, schemas map[string]jsonschema.Schema) *jsonschema.Schema {
+func resolveRefs(s *jsonschema.Schema, schemas map[string]*jsonschema.Schema) *jsonschema.Schema {
 	node := s
 
 	description := s.Description
@@ -74,14 +74,15 @@ func resolveRefs(s *jsonschema.Schema, schemas map[string]jsonschema.Schema) *js
 			examples = newNode.Examples
 		}
 
-		node = &newNode
+		node = newNode
 	}
 
-	node.Description = description
-	node.MarkdownDescription = markdownDescription
-	node.Examples = examples
+	newNode := *node
+	newNode.Description = description
+	newNode.MarkdownDescription = markdownDescription
+	newNode.Examples = examples
 
-	return node
+	return &newNode
 }
 
 func getRefType(node *jsonschema.Schema) string {