From 0c5a5c07bb1686122bfb68c3677c61d6738c3d2e Mon Sep 17 00:00:00 2001
From: Jefftree <jeffrey.ying86@live.com>
Date: Tue, 17 Jan 2023 18:26:36 +0000
Subject: [PATCH] Add AdditionalDefinitions

---
 pkg/builder/openapi.go  | 41 ++++++++++++++++++++++++++++++++---------
 pkg/builder3/openapi.go | 26 ++++++++++++++++++++++----
 pkg/common/common.go    | 10 ++++++++++
 3 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/pkg/builder/openapi.go b/pkg/builder/openapi.go
index 1c4cb5bf8..2187b9856 100644
--- a/pkg/builder/openapi.go
+++ b/pkg/builder/openapi.go
@@ -35,10 +35,24 @@ const (
 )
 
 type openAPI struct {
-	config       *common.Config
-	swagger      *spec.Swagger
-	protocolList []string
-	definitions  map[string]common.OpenAPIDefinition
+	config                *common.Config
+	swagger               *spec.Swagger
+	protocolList          []string
+	definitions           map[string]common.OpenAPIDefinition
+	additionalDefinitions map[string]common.OpenAPIDefinition
+}
+
+// getDefinition is a getter that checks two maps for the existence of a key
+func (o *openAPI) getDefinition(key string) (common.OpenAPIDefinition, bool) {
+	val, ok := o.definitions[key]
+	if ok {
+		return val, ok
+	}
+	if o.additionalDefinitions != nil {
+		val, ok = o.additionalDefinitions[key]
+		return val, ok
+	}
+	return common.OpenAPIDefinition{}, false
 }
 
 // BuildOpenAPISpec builds OpenAPI spec given a list of route containers and common.Config to customize it.
@@ -127,10 +141,19 @@ func newOpenAPI(config *common.Config) openAPI {
 			return name[strings.LastIndex(name, "/")+1:], nil
 		}
 	}
-	o.definitions = o.config.GetDefinitions(func(name string) spec.Ref {
-		defName, _ := o.config.GetDefinitionName(name)
-		return spec.MustCreateRef("#/definitions/" + common.EscapeJsonPointer(defName))
-	})
+	if o.config.Definitions != nil {
+		o.definitions = o.config.Definitions
+	} else {
+		o.definitions = o.config.GetDefinitions(func(name string) spec.Ref {
+			defName, _ := o.config.GetDefinitionName(name)
+			return spec.MustCreateRef("#/definitions/" + common.EscapeJsonPointer(defName))
+		})
+	}
+
+	if o.config.AdditionalDefinitions != nil {
+		o.additionalDefinitions = o.config.AdditionalDefinitions
+	}
+
 	if o.config.CommonResponses == nil {
 		o.config.CommonResponses = map[int]spec.Response{}
 	}
@@ -160,7 +183,7 @@ func (o *openAPI) buildDefinitionRecursively(name string) error {
 	if _, ok := o.swagger.Definitions[uniqueName]; ok {
 		return nil
 	}
-	if item, ok := o.definitions[name]; ok {
+	if item, ok := o.getDefinition(name); ok {
 		schema := spec.Schema{
 			VendorExtensible:   item.Schema.VendorExtensible,
 			SchemaProps:        item.Schema.SchemaProps,
diff --git a/pkg/builder3/openapi.go b/pkg/builder3/openapi.go
index e59844786..8b7024704 100644
--- a/pkg/builder3/openapi.go
+++ b/pkg/builder3/openapi.go
@@ -37,9 +37,23 @@ const (
 )
 
 type openAPI struct {
-	config      *common.OpenAPIV3Config
-	spec        *spec3.OpenAPI
-	definitions map[string]common.OpenAPIDefinition
+	config                *common.OpenAPIV3Config
+	spec                  *spec3.OpenAPI
+	definitions           map[string]common.OpenAPIDefinition
+	additionalDefinitions map[string]common.OpenAPIDefinition
+}
+
+// getDefinition is a getter that checks two maps for the existence of a key
+func (o *openAPI) getDefinition(key string) (common.OpenAPIDefinition, bool) {
+	val, ok := o.definitions[key]
+	if ok {
+		return val, ok
+	}
+	if o.additionalDefinitions != nil {
+    val, ok = o.additionalDefinitions[key]
+    return val, ok
+  }
+  return nil, false
 }
 
 func groupRoutesByPath(routes []common.Route) map[string][]common.Route {
@@ -237,6 +251,10 @@ func newOpenAPI(config *common.OpenAPIV3Config) openAPI {
 		})
 	}
 
+	if o.config.AdditionalDefinitions != nil {
+		o.additionalDefinitions = o.config.AdditionalDefinitions
+	}
+
 	return o
 }
 
@@ -433,7 +451,7 @@ func (o *openAPI) buildDefinitionRecursively(name string) error {
 	if _, ok := o.spec.Components.Schemas[uniqueName]; ok {
 		return nil
 	}
-	if item, ok := o.definitions[name]; ok {
+	if item, ok := o.getDefinition(name); ok {
 		schema := &spec.Schema{
 			VendorExtensible:   item.Schema.VendorExtensible,
 			SchemaProps:        item.Schema.SchemaProps,
diff --git a/pkg/common/common.go b/pkg/common/common.go
index 2e15e163c..4c272adb7 100644
--- a/pkg/common/common.go
+++ b/pkg/common/common.go
@@ -97,6 +97,11 @@ type Config struct {
 	// This takes precedent over the GetDefinitions function
 	Definitions map[string]OpenAPIDefinition
 
+	// Provides an additional set of definitions for models used by routes.
+	// This is unioned with Definitions when building the spec, and allows
+	// the Definitions map to be reused across different builders.
+	AdditionalDefinitions map[string]OpenAPIDefinition
+
 	// GetOperationIDAndTags returns operation id and tags for a restful route. It is an optional function to customize operation IDs.
 	//
 	// Deprecated: GetOperationIDAndTagsFromRoute should be used instead. This cannot be specified if using the new Route
@@ -151,6 +156,11 @@ type OpenAPIV3Config struct {
 	// This takes precedent over the GetDefinitions function
 	Definitions map[string]OpenAPIDefinition
 
+	// Provides an additional set of definitions for models used by routes.
+	// This is unioned with Definitions when building the spec, and allows
+	// the Definitions map to be reused across different builders.
+	AdditionalDefinitions map[string]OpenAPIDefinition
+
 	// GetOperationIDAndTags returns operation id and tags for a restful route. It is an optional function to customize operation IDs.
 	//
 	// Deprecated: GetOperationIDAndTagsFromRoute should be used instead. This cannot be specified if using the new Route