From 59731be704e4aebd2e7656410963a2f3fdd78b90 Mon Sep 17 00:00:00 2001
From: Zaid Ajaj <zaid.naom@gmail.com>
Date: Mon, 5 Dec 2022 18:19:37 +0100
Subject: [PATCH] Implement UnmarshalYAML for ReturnTypeSpec

---
 pkg/codegen/schema/schema.go | 30 +++++++++++++++++++++++++++---
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/pkg/codegen/schema/schema.go b/pkg/codegen/schema/schema.go
index b1bc56192ee9..4a09210d076a 100644
--- a/pkg/codegen/schema/schema.go
+++ b/pkg/codegen/schema/schema.go
@@ -1564,9 +1564,6 @@ func (returnTypeSpec *ReturnTypeSpec) UnmarshalJSON(inputJSON []byte) error {
 	}
 
 	if len(data) == 0 {
-		// then the schema supplied "outputs":{}
-		// treat as of outputs was not defined
-		returnTypeSpec = nil
 		return nil
 	}
 
@@ -1587,6 +1584,33 @@ func (returnTypeSpec *ReturnTypeSpec) UnmarshalJSON(inputJSON []byte) error {
 	return nil
 }
 
+func (returnTypeSpec *ReturnTypeSpec) UnmarshalYAML(node *yaml.Node) error {
+	var data map[string]interface{}
+	if err := node.Decode(&data); err != nil {
+		return err
+	}
+
+	if len(data) == 0 {
+		return nil
+	}
+
+	var objectSpec *ObjectTypeSpec
+	var typeSpec *TypeSpec
+	if _, hasProperties := data["properties"]; hasProperties {
+		if err := node.Decode(&objectSpec); err != nil {
+			return err
+		}
+	} else {
+		if err := node.Decode(&typeSpec); err != nil {
+			return err
+		}
+	}
+
+	returnTypeSpec.TypeSpec = typeSpec
+	returnTypeSpec.ObjectTypeSpec = objectSpec
+	return nil
+}
+
 // FunctionSpec is the serializable form of a function description.
 type FunctionSpec struct {
 	// Description is the description of the function, if any.