fix: extensions in YAML format [#2795] (#2797)

* fix: Support extensions for YAML [#2795] * Move extensionsToMap to generator * bazel fix * define map size
grpc-ecosystem · Jul 18, 2022 · cec112b · cec112b
1 parent c75cfce
commit cec112b
Show file tree

Hide file tree

Showing 4 changed files with 217 additions and 13 deletions.
diff --git a/protoc-gen-openapiv2/internal/genopenapi/BUILD.bazel b/protoc-gen-openapiv2/internal/genopenapi/BUILD.bazel
@@ -61,6 +61,7 @@ go_test(
         "@in_gopkg_yaml_v3//:yaml_v3",
         "@io_bazel_rules_go//proto/wkt:field_mask_go_proto",
         "@org_golang_google_protobuf//encoding/protojson",
+        "@org_golang_google_protobuf//encoding/prototext",
         "@org_golang_google_protobuf//proto",
         "@org_golang_google_protobuf//reflect/protodesc",
         "@org_golang_google_protobuf//types/descriptorpb",

diff --git a/protoc-gen-openapiv2/internal/genopenapi/generator.go b/protoc-gen-openapiv2/internal/genopenapi/generator.go
@@ -89,36 +89,133 @@ func (so openapiSwaggerObject) MarshalJSON() ([]byte, error) {
 	return extensionMarshalJSON(alias(so), so.extensions)
 }
 
+// MarshalYAML implements yaml.Marshaler interface.
+//
+// It is required in order to pass extensions inline.
+//
+// Example:
+//   extensions: {x-key: x-value}
+//   type: string
+//
+// It will be rendered as:
+//   x-key: x-value
+//   type: string
+//
+// Use generics when the project will be upgraded to go 1.18+.
+func (so openapiSwaggerObject) MarshalYAML() (interface{}, error) {
+	type Alias openapiSwaggerObject
+
+	return struct {
+		Extension map[string]interface{} `yaml:",inline"`
+		Alias     `yaml:",inline"`
+	}{
+		Extension: extensionsToMap(so.extensions),
+		Alias:     Alias(so),
+	}, nil
+}
+
 func (so openapiInfoObject) MarshalJSON() ([]byte, error) {
 	type alias openapiInfoObject
 	return extensionMarshalJSON(alias(so), so.extensions)
 }
 
+func (so openapiInfoObject) MarshalYAML() (interface{}, error) {
+	type Alias openapiInfoObject
+
+	return struct {
+		Extension map[string]interface{} `yaml:",inline"`
+		Alias     `yaml:",inline"`
+	}{
+		Extension: extensionsToMap(so.extensions),
+		Alias:     Alias(so),
+	}, nil
+}
+
 func (so openapiSecuritySchemeObject) MarshalJSON() ([]byte, error) {
 	type alias openapiSecuritySchemeObject
 	return extensionMarshalJSON(alias(so), so.extensions)
 }
 
+func (so openapiSecuritySchemeObject) MarshalYAML() (interface{}, error) {
+	type Alias openapiSecuritySchemeObject
+
+	return struct {
+		Extension map[string]interface{} `yaml:",inline"`
+		Alias     `yaml:",inline"`
+	}{
+		Extension: extensionsToMap(so.extensions),
+		Alias:     Alias(so),
+	}, nil
+}
+
 func (so openapiOperationObject) MarshalJSON() ([]byte, error) {
 	type alias openapiOperationObject
 	return extensionMarshalJSON(alias(so), so.extensions)
 }
 
+func (so openapiOperationObject) MarshalYAML() (interface{}, error) {
+	type Alias openapiOperationObject
+
+	return struct {
+		Extension map[string]interface{} `yaml:",inline"`
+		Alias     `yaml:",inline"`
+	}{
+		Extension: extensionsToMap(so.extensions),
+		Alias:     Alias(so),
+	}, nil
+}
+
 func (so openapiResponseObject) MarshalJSON() ([]byte, error) {
 	type alias openapiResponseObject
 	return extensionMarshalJSON(alias(so), so.extensions)
 }
 
+func (so openapiResponseObject) MarshalYAML() (interface{}, error) {
+	type Alias openapiResponseObject
+
+	return struct {
+		Extension map[string]interface{} `yaml:",inline"`
+		Alias     `yaml:",inline"`
+	}{
+		Extension: extensionsToMap(so.extensions),
+		Alias:     Alias(so),
+	}, nil
+}
+
 func (so openapiSchemaObject) MarshalJSON() ([]byte, error) {
 	type alias openapiSchemaObject
 	return extensionMarshalJSON(alias(so), so.extensions)
 }
 
+func (so openapiSchemaObject) MarshalYAML() (interface{}, error) {
+	type Alias openapiSchemaObject
+
+	return struct {
+		Extension map[string]interface{} `yaml:",inline"`
+		Alias     `yaml:",inline"`
+	}{
+		Extension: extensionsToMap(so.extensions),
+		Alias:     Alias(so),
+	}, nil
+}
+
 func (so openapiParameterObject) MarshalJSON() ([]byte, error) {
 	type alias openapiParameterObject
 	return extensionMarshalJSON(alias(so), so.extensions)
 }
 
+func (so openapiParameterObject) MarshalYAML() (interface{}, error) {
+	type Alias openapiParameterObject
+
+	return struct {
+		Extension map[string]interface{} `yaml:",inline"`
+		Alias     `yaml:",inline"`
+	}{
+		Extension: extensionsToMap(so.extensions),
+		Alias:     Alias(so),
+	}, nil
+}
+
 func extensionMarshalJSON(so interface{}, extensions []extension) ([]byte, error) {
 	// To append arbitrary keys to the struct we'll render into json,
 	// we're creating another struct that embeds the original one, and
@@ -262,3 +359,13 @@ func AddErrorDefs(reg *descriptor.Registry) error {
 		},
 	})
 }
+
+func extensionsToMap(extensions []extension) map[string]interface{} {
+	m := make(map[string]interface{}, len(extensions))
+
+	for _, v := range extensions {
+		m[v.key] = RawExample(v.value)
+	}
+
+	return m
+}
diff --git a/protoc-gen-openapiv2/internal/genopenapi/generator_test.go b/protoc-gen-openapiv2/internal/genopenapi/generator_test.go
@@ -1,12 +1,14 @@
 package genopenapi_test
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/grpc-ecosystem/grpc-gateway/v2/internal/descriptor"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-openapiv2/internal/genopenapi"
 	"gopkg.in/yaml.v3"
 
+	"google.golang.org/protobuf/encoding/prototext"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/descriptorpb"
 	"google.golang.org/protobuf/types/pluginpb"
@@ -15,7 +17,6 @@ import (
 func TestGenerate_YAML(t *testing.T) {
 	t.Parallel()
 
-	reg := descriptor.NewRegistry()
 	req := &pluginpb.CodeGeneratorRequest{
 		ProtoFile: []*descriptorpb.FileDescriptorProto{{
 			Name:    proto.String("file.proto"),
@@ -29,31 +30,126 @@ func TestGenerate_YAML(t *testing.T) {
 		},
 	}
 
+	resp := requireGenerate(t, req, genopenapi.FormatYAML)
+	if len(resp) != 1 {
+		t.Fatalf("invalid count, expected: 1, actual: %d", len(resp))
+	}
+
+	var p map[string]interface{}
+	err := yaml.Unmarshal([]byte(resp[0].GetContent()), &p)
+	if err != nil {
+		t.Fatalf("failed to unmarshall yaml: %s", err)
+	}
+}
+
+func TestGenerateExtension(t *testing.T) {
+	t.Parallel()
+
+	const in = `
+	file_to_generate: "exampleproto/v1/example.proto"
+	parameter: "output_format=yaml,allow_delete_body=true"
+	proto_file: {
+		name: "exampleproto/v1/example.proto"
+		package: "example.v1"
+		message_type: {
+			name: "Foo"
+			field: {
+				name: "bar"
+				number: 1
+				label: LABEL_OPTIONAL
+				type: TYPE_STRING
+				json_name: "bar"
+				options: {
+					[grpc.gateway.protoc_gen_openapiv2.options.openapiv2_field]: {
+						description: "This is bar"
+						extensions: {
+							key: "x-go-default"
+							value: {
+								string_value: "0.5s"
+							}
+						}
+					}
+				}
+			}
+		}
+		service: {
+			name: "TestService"
+			method: {
+				name: "Test"
+				input_type: ".example.v1.Foo"
+				output_type: ".example.v1.Foo"
+				options: {}
+			}
+		}
+		options: {
+			go_package: "exampleproto/v1;exampleproto"
+		}
+	}`
+
+	var req pluginpb.CodeGeneratorRequest
+	if err := prototext.Unmarshal([]byte(in), &req); err != nil {
+		t.Fatalf("failed to marshall yaml: %s", err)
+	}
+
+	formats := [...]genopenapi.Format{
+		genopenapi.FormatJSON,
+		genopenapi.FormatYAML,
+	}
+
+	for _, format := range formats {
+		format := format
+
+		t.Run(string(format), func(t *testing.T) {
+			t.Parallel()
+
+			resp := requireGenerate(t, &req, format)
+			if len(resp) != 1 {
+				t.Fatalf("invalid count, expected: 1, actual: %d", len(resp))
+			}
+
+			content := resp[0].GetContent()
+
+			t.Log(content)
+
+			if !strings.Contains(content, "x-go-default") {
+				t.Fatal("x-go-default not found in content message")
+			}
+		})
+	}
+}
+
+func requireGenerate(
+	tb testing.TB,
+	req *pluginpb.CodeGeneratorRequest,
+	format genopenapi.Format,
+) []*descriptor.ResponseFile {
+	tb.Helper()
+
+	reg := descriptor.NewRegistry()
+
 	if err := reg.Load(req); err != nil {
-		t.Fatalf("failed to load request: %s", err)
+		tb.Fatalf("failed to load request: %s", err)
 	}
 
 	var targets []*descriptor.File
 	for _, target := range req.FileToGenerate {
 		f, err := reg.LookupFile(target)
 		if err != nil {
-			t.Fatalf("failed to lookup file: %s", err)
+			tb.Fatalf("failed to lookup file: %s", err)
 		}
+
 		targets = append(targets, f)
 	}
 
-	g := genopenapi.New(reg, genopenapi.FormatYAML)
+	g := genopenapi.New(reg, format)
+
 	resp, err := g.Generate(targets)
 	switch {
 	case err != nil:
-		t.Fatalf("failed to generate targets: %s", err)
-	case len(resp) != 1:
-		t.Fatalf("invalid count, expected: 1, actual: %d", len(resp))
+		tb.Fatalf("failed to generate targets: %s", err)
+	case len(resp) != len(targets):
+		tb.Fatalf("invalid count, expected: %d, actual: %d", len(targets), len(resp))
 	}
 
-	var p map[string]interface{}
-	err = yaml.Unmarshal([]byte(resp[0].GetContent()), &p)
-	if err != nil {
-		t.Fatalf("failed to unmarshall yaml: %s", err)
-	}
+	return resp
 }
diff --git a/protoc-gen-openapiv2/internal/genopenapi/types.go b/protoc-gen-openapiv2/internal/genopenapi/types.go
@@ -256,7 +256,7 @@ type keyVal struct {
 type openapiSchemaObjectProperties []keyVal
 
 func (p openapiSchemaObjectProperties) MarshalYAML() (interface{}, error) {
-	m := make(map[string]interface{})
+	m := make(map[string]interface{}, len(p))
 
 	for _, v := range p {
 		m[v.Key] = v.Value