From 5aee5fc87522a772df4e10e6b90ddc37fef079b4 Mon Sep 17 00:00:00 2001
From: Taufik Rama <taufik-rama+github.com@pm.me>
Date: Sat, 6 Apr 2024 16:06:15 +0700
Subject: [PATCH] openapi3gen: allow overriding how a Schema is generated
 (#920)

* openapi3gen: adds custom schema definition for client implementing our interface

* Add Unitest and trigger docs.sh (#1)

* add unitest, generate docs

* test: change test setschema to example

---------

Co-authored-by: Taufik Rama <ruff.dev@gmail.com>
Co-authored-by: Ferianto Surya Wijaya <33758092+blackhun11@users.noreply.github.com>
Co-authored-by: blackhun11 <feriantosw77@gmail.com>
---
 .github/docs/openapi3gen.txt    |  7 ++++++
 openapi3gen/openapi3gen.go      | 15 +++++++++++++
 openapi3gen/openapi3gen_test.go | 40 +++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/.github/docs/openapi3gen.txt b/.github/docs/openapi3gen.txt
index 9cf293cd2..e692b8594 100644
--- a/.github/docs/openapi3gen.txt
+++ b/.github/docs/openapi3gen.txt
@@ -70,3 +70,10 @@ type SchemaCustomizerFn func(name string, t reflect.Type, tag reflect.StructTag,
     A SchemaCustomizerFn can return an ExcludeSchemaSentinel error to indicate
     that the schema for this field should not be included in the final output
 
+type SetSchemar interface {
+	SetSchema(*openapi3.Schema)
+}
+    SetSchemar allows client to set their own schema definition according to
+    their specification. Useful when some custom datatype is needed and/or some
+    custom logic is needed on how the schema values would be generated
+
diff --git a/openapi3gen/openapi3gen.go b/openapi3gen/openapi3gen.go
index 95805ee25..a301337e1 100644
--- a/openapi3gen/openapi3gen.go
+++ b/openapi3gen/openapi3gen.go
@@ -35,6 +35,13 @@ type Option func(*generatorOpt)
 // the final output
 type SchemaCustomizerFn func(name string, t reflect.Type, tag reflect.StructTag, schema *openapi3.Schema) error
 
+// SetSchemar allows client to set their own schema definition according to
+// their specification. Useful when some custom datatype is needed and/or some custom logic
+// is needed on how the schema values would be generated
+type SetSchemar interface {
+	SetSchema(*openapi3.Schema)
+}
+
 type generatorOpt struct {
 	useAllExportedFields bool
 	throwErrorOnCycle    bool
@@ -347,6 +354,14 @@ func (g *Generator) generateWithoutSaving(parents []*theTypeInfo, t reflect.Type
 				schema.Type = &openapi3.Types{"object"}
 			}
 		}
+
+	default:
+		// Object has their own schema's implementation, so we'll use those
+		if v := reflect.New(t); v.CanInterface() {
+			if v, ok := v.Interface().(SetSchemar); ok {
+				v.SetSchema(schema)
+			}
+		}
 	}
 
 	if g.opts.schemaCustomizer != nil {
diff --git a/openapi3gen/openapi3gen_test.go b/openapi3gen/openapi3gen_test.go
index daea9ffd7..bb9ea4daf 100644
--- a/openapi3gen/openapi3gen_test.go
+++ b/openapi3gen/openapi3gen_test.go
@@ -590,3 +590,43 @@ func ExampleNewSchemaRefForValue_recursive() {
 	//   "type": "object"
 	// }
 }
+
+type ID [16]byte
+
+// T implements SetSchemar, allowing it to set an OpenAPI schema.
+type T struct {
+	ID ID `json:"id"`
+}
+
+func (_ *ID) SetSchema(schema *openapi3.Schema) {
+	schema.Type = &openapi3.Types{"string"} // Assuming this matches your custom implementation
+	schema.Format = "uuid"
+}
+
+func ExampleID_SetSchema() {
+	schemas := make(openapi3.Schemas)
+	instance := &T{
+		ID: ID{},
+	}
+
+	// Generate the schema for the instance
+	schemaRef, err := openapi3gen.NewSchemaRefForValue(instance, schemas)
+	if err != nil {
+		panic(err)
+	}
+	data, err := json.MarshalIndent(schemaRef, "", "  ")
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("schemaRef: %s\n", data)
+	// Output:
+	// schemaRef: {
+	//   "properties": {
+	//     "id": {
+	//       "format": "uuid",
+	//       "type": "string"
+	//     }
+	//   },
+	//   "type": "object"
+	// }
+}