3388 - Add feature to support passing values to go templates via cli

grpc-ecosystem · Dec 3, 2023 · 71c6fdc · 71c6fdc
1 parent 46e828a
commit 71c6fdc
Show file tree

Hide file tree

Showing 6 changed files with 208 additions and 0 deletions.
diff --git a/docs/docs/mapping/customizing_openapi_output.md b/docs/docs/mapping/customizing_openapi_output.md
@@ -231,6 +231,61 @@ This is how the OpenAPI file would be rendered in [Postman](https://www.getpostm
 
 For a more detailed example of a proto file that has Go, templates enabled, [see the examples](https://github.com/grpc-ecosystem/grpc-gateway/blob/main/examples/internal/proto/examplepb/use_go_template.proto).
 
+### Using custom values
+
+Custom values can be specified in the [Go templates](https://golang.org/pkg/text/template/) that generate your proto file comments.
+
+A use case might be to interpolate different external documentation URLs when rendering documentation for different environments.
+
+#### How to use it
+
+The `use_go_templates` option has to be enabled as a prerequisite.
+
+Provide customized values in the format of `go_template_args=my_key=my_value`. `{{arg "my_key"}}` will be replaced with `my_value` in the Go template.
+
+Specify the `go_template_args` option multiple times if needed.
+
+```sh
+--openapiv2_out . --openapiv2_opt use_go_templates=true --openapiv2_opt go_template_args=my_key1=my_value1 --openapiv2_opt go_template_args=my_key2=my_value2
+...
+```
+
+#### Example script
+
+Example of a bash script with the `use_go_templates` flag set to true and custom template values set:
+
+```sh
+$ protoc -I. \
+    --go_out . --go-grpc_out . \
+    --grpc-gateway_out . \
+    --openapiv2_out . \
+    --openapiv2_opt use_go_templates=true \
+    --openapiv2_opt go_template_args=environment=test1 \
+    --openapiv2_opt go_template_args=environment_label=Test1 \
+    path/to/my/proto/v1/myproto.proto
+```
+
+#### Example proto file
+
+Example of a proto file with Go templates and custom values:
+
+```protobuf
+service LoginService {
+    // Login (Environment: {{arg "environment_label"}})
+    //
+    // {{.MethodDescriptorProto.Name}} is a call with the method(s) {{$first := true}}{{range .Bindings}}{{if $first}}{{$first = false}}{{else}}, {{end}}{{.HTTPMethod}}{{end}} within the "{{.Service.Name}}" service.
+    // It takes in "{{.RequestType.Name}}" and returns a "{{.ResponseType.Name}}".
+    // This only works in the {{arg "environment"}} domain.
+    //
+    rpc Login (LoginRequest) returns (LoginReply) {
+        option (google.api.http) = {
+            post: "/v1/example/login"
+            body: "*"
+        };
+    }
+}
+```
+
 ## Other plugin options
 
 A comprehensive list of OpenAPI plugin options can be found [here](https://github.com/grpc-ecosystem/grpc-gateway/blob/main/protoc-gen-openapiv2/main.go). Options can be passed via `protoc` CLI:

diff --git a/internal/descriptor/registry.go b/internal/descriptor/registry.go
@@ -78,6 +78,9 @@ type Registry struct {
 	// in your protofile comments
 	useGoTemplate bool
 
+	// goTemplateArgs specifies a list of key value pair inputs to be displayed in Go templates
+	goTemplateArgs map[string]string
+
 	// ignoreComments determines whether all protofile comments should be excluded from output
 	ignoreComments bool
 
@@ -583,6 +586,19 @@ func (r *Registry) GetUseGoTemplate() bool {
 	return r.useGoTemplate
 }
 
+func (r *Registry) SetGoTemplateArgs(kvs []string) {
+	r.goTemplateArgs = make(map[string]string)
+	for _, kv := range kvs {
+		if key, value, found := strings.Cut(kv, "="); found {
+			r.goTemplateArgs[key] = value
+		}
+	}
+}
+
+func (r *Registry) GetGoTemplateArgs() map[string]string {
+	return r.goTemplateArgs
+}
+
 // SetIgnoreComments sets ignoreComments
 func (r *Registry) SetIgnoreComments(ignore bool) {
 	r.ignoreComments = ignore

diff --git a/protoc-gen-openapiv2/defs.bzl b/protoc-gen-openapiv2/defs.bzl
@@ -60,6 +60,7 @@ def _run_proto_gen_openapi(
         fqn_for_openapi_name,
         openapi_naming_strategy,
         use_go_templates,
+        go_template_args,
         ignore_comments,
         remove_internal_comments,
         disable_default_errors,
@@ -110,6 +111,9 @@ def _run_proto_gen_openapi(
     if use_go_templates:
         args.add("--openapiv2_opt", "use_go_templates=true")
 
+    for go_template_arg in go_template_args:
+        args.add("--openapiv2_opt", "go_template_args=%s" % go_template_arg)
+
     if ignore_comments:
         args.add("--openapiv2_opt", "ignore_comments=true")
 
@@ -232,6 +236,7 @@ def _proto_gen_openapi_impl(ctx):
                     fqn_for_openapi_name = ctx.attr.fqn_for_openapi_name,
                     openapi_naming_strategy = ctx.attr.openapi_naming_strategy,
                     use_go_templates = ctx.attr.use_go_templates,
+                    go_template_args = ctx.attr.go_template_args,
                     ignore_comments = ctx.attr.ignore_comments,
                     remove_internal_comments = ctx.attr.remove_internal_comments,
                     disable_default_errors = ctx.attr.disable_default_errors,
@@ -312,6 +317,12 @@ protoc_gen_openapiv2 = rule(
             mandatory = False,
             doc = "if set, you can use Go templates in protofile comments",
         ),
+        "go_template_args": attr.string_list(
+            mandatory = False,
+            doc = "specify a key value pair as inputs to the Go template of the protofile" +
+                  " comments. Repeat this option to specify multiple template arguments." +
+                  " Requires the `use_go_templates` option to be set.",
+        ),
         "ignore_comments": attr.bool(
             default = False,
             mandatory = False,

diff --git a/protoc-gen-openapiv2/internal/genopenapi/template.go b/protoc-gen-openapiv2/internal/genopenapi/template.go
@@ -2486,6 +2486,12 @@ func goTemplateComments(comment string, data interface{}, reg *descriptor.Regist
 		"fieldcomments": func(msg *descriptor.Message, field *descriptor.Field) string {
 			return strings.ReplaceAll(fieldProtoComments(reg, msg, field), "\n", "<br>")
 		},
+		"arg": func(name string) string {
+			if v, f := reg.GetGoTemplateArgs()[name]; f {
+				return v
+			}
+			return fmt.Sprintf("goTemplateArg %s not found", name)
+		},
 	}).Parse(comment)
 	if err != nil {
 		// If there is an error parsing the templating insert the error as string in the comment

diff --git a/protoc-gen-openapiv2/internal/genopenapi/template_test.go b/protoc-gen-openapiv2/internal/genopenapi/template_test.go
@@ -5841,6 +5841,7 @@ func TestUpdateOpenAPIDataFromComments(t *testing.T) {
 		expectedError         error
 		expectedOpenAPIObject interface{}
 		useGoTemplate         bool
+		goTemplateArgs        []string
 	}{
 		{
 			descr:                 "empty comments",
@@ -5968,6 +5969,33 @@ func TestUpdateOpenAPIDataFromComments(t *testing.T) {
 			expectedError: nil,
 			useGoTemplate: true,
 		},
+		{
+			descr:                "template with use_go_template and go_template_args",
+			openapiSwaggerObject: &openapiSchemaObject{},
+			expectedOpenAPIObject: &openapiSchemaObject{
+				Title:       "Template",
+				Description: `Description "which means nothing" for environment test with value my_value`,
+			},
+			comments: "Template\n\nDescription {{with \"which means nothing\"}}{{printf \"%q\" .}}{{end}} for " +
+				"environment {{arg \"environment\"}} with value {{arg \"my_key\"}}",
+			expectedError:  nil,
+			useGoTemplate:  true,
+			goTemplateArgs: []string{"my_key=my_value", "environment=test"},
+		},
+		{
+			descr:                "template with use_go_template and undefined go_template_args",
+			openapiSwaggerObject: &openapiSchemaObject{},
+			expectedOpenAPIObject: &openapiSchemaObject{
+				Title: "Template",
+				Description: `Description "which means nothing" for environment test with value ` +
+					`goTemplateArg something_undefined not found`,
+			},
+			comments: "Template\n\nDescription {{with \"which means nothing\"}}{{printf \"%q\" .}}{{end}} for " +
+				"environment {{arg \"environment\"}} with value {{arg \"something_undefined\"}}",
+			expectedError:  nil,
+			useGoTemplate:  true,
+			goTemplateArgs: []string{"environment=test"},
+		},
 	}
 
 	for _, test := range tests {
@@ -5976,6 +6004,9 @@ func TestUpdateOpenAPIDataFromComments(t *testing.T) {
 			if test.useGoTemplate {
 				reg.SetUseGoTemplate(true)
 			}
+			if len(test.goTemplateArgs) > 0 {
+				reg.SetGoTemplateArgs(test.goTemplateArgs)
+			}
 			err := updateOpenAPIDataFromComments(reg, test.openapiSwaggerObject, nil, test.comments, false)
 			if test.expectedError == nil {
 				if err != nil {
@@ -6007,6 +6038,7 @@ func TestMessageOptionsWithGoTemplate(t *testing.T) {
 		defs           openapiDefinitionsObject
 		openAPIOptions *openapiconfig.OpenAPIOptions
 		useGoTemplate  bool
+		goTemplateArgs []string
 	}{
 		{
 			descr: "external docs option",
@@ -6068,6 +6100,39 @@ func TestMessageOptionsWithGoTemplate(t *testing.T) {
 			},
 			useGoTemplate: false,
 		},
+		{
+			descr: "external docs option with go template args",
+			msgDescs: []*descriptorpb.DescriptorProto{
+				{Name: proto.String("Message")},
+			},
+			schema: map[string]*openapi_options.Schema{
+				"Message": {
+					JsonSchema: &openapi_options.JSONSchema{
+						Title: "{{.Name}}",
+						Description: "Description {{with \"which means nothing\"}}{{printf \"%q\" .}}{{end}} " +
+							"{{arg \"my_key\"}}",
+					},
+					ExternalDocs: &openapi_options.ExternalDocumentation{
+						Description: "Description {{with \"which means nothing\"}}{{printf \"%q\" .}}{{end}} " +
+							"{{arg \"my_key\"}}",
+					},
+				},
+			},
+			defs: map[string]openapiSchemaObject{
+				"Message": {
+					schemaCore: schemaCore{
+						Type: "object",
+					},
+					Title:       "Message",
+					Description: `Description "which means nothing" too`,
+					ExternalDocs: &openapiExternalDocumentationObject{
+						Description: `Description "which means nothing" too`,
+					},
+				},
+			},
+			useGoTemplate:  true,
+			goTemplateArgs: []string{"my_key=too"},
+		},
 		{
 			descr: "registered OpenAPIOption",
 			msgDescs: []*descriptorpb.DescriptorProto{
@@ -6103,6 +6168,44 @@ func TestMessageOptionsWithGoTemplate(t *testing.T) {
 			},
 			useGoTemplate: true,
 		},
+		{
+			descr: "registered OpenAPIOption with go template args",
+			msgDescs: []*descriptorpb.DescriptorProto{
+				{Name: proto.String("Message")},
+			},
+			openAPIOptions: &openapiconfig.OpenAPIOptions{
+				Message: []*openapiconfig.OpenAPIMessageOption{
+					{
+						Message: "example.Message",
+						Option: &openapi_options.Schema{
+							JsonSchema: &openapi_options.JSONSchema{
+								Title: "{{.Name}}",
+								Description: "Description {{with \"which means nothing\"}}{{printf \"%q\" .}}{{end}} " +
+									"{{arg \"my_key\"}}",
+							},
+							ExternalDocs: &openapi_options.ExternalDocumentation{
+								Description: "Description {{with \"which means nothing\"}}{{printf \"%q\" .}}{{end}} " +
+									"{{arg \"my_key\"}}",
+							},
+						},
+					},
+				},
+			},
+			defs: map[string]openapiSchemaObject{
+				"Message": {
+					schemaCore: schemaCore{
+						Type: "object",
+					},
+					Title:       "Message",
+					Description: `Description "which means nothing" too`,
+					ExternalDocs: &openapiExternalDocumentationObject{
+						Description: `Description "which means nothing" too`,
+					},
+				},
+			},
+			useGoTemplate:  true,
+			goTemplateArgs: []string{"my_key=too"},
+		},
 	}
 
 	for _, test := range tests {
@@ -6116,6 +6219,7 @@ func TestMessageOptionsWithGoTemplate(t *testing.T) {
 
 			reg := descriptor.NewRegistry()
 			reg.SetUseGoTemplate(test.useGoTemplate)
+			reg.SetGoTemplateArgs(test.goTemplateArgs)
 			file := descriptor.File{
 				FileDescriptorProto: &descriptorpb.FileDescriptorProto{
 					SourceCodeInfo: &descriptorpb.SourceCodeInfo{},
@@ -6174,6 +6278,7 @@ func TestMessageOptionsWithGoTemplate(t *testing.T) {
 func TestTagsWithGoTemplate(t *testing.T) {
 	reg := descriptor.NewRegistry()
 	reg.SetUseGoTemplate(true)
+	reg.SetGoTemplateArgs([]string{"my_key=my_value"})
 
 	svc := &descriptorpb.ServiceDescriptorProto{
 		Name:    proto.String("ExampleService"),
@@ -6220,6 +6325,10 @@ func TestTagsWithGoTemplate(t *testing.T) {
 				Name:        "not a service tag 2",
 				Description: "{{ import \"file\" }}",
 			},
+			{
+				Name:        "Service with my_key",
+				Description: "the {{arg \"my_key\"}}",
+			},
 		},
 	}
 	proto.SetExtension(proto.Message(file.FileDescriptorProto.Options), openapi_options.E_Openapiv2Swagger, &swagger)
@@ -6241,6 +6350,10 @@ func TestTagsWithGoTemplate(t *testing.T) {
 			Name:        "not a service tag 2",
 			Description: "open file: no such file or directory",
 		},
+		{
+			Name:        "Service with my_key",
+			Description: "the my_value",
+		},
 	}
 	if !reflect.DeepEqual(actual.Tags, expectedTags) {
 		t.Errorf("Expected tags %+v, not %+v", expectedTags, actual.Tags)

diff --git a/protoc-gen-openapiv2/main.go b/protoc-gen-openapiv2/main.go
@@ -30,6 +30,7 @@ var (
 	useFQNForOpenAPIName           = flag.Bool("fqn_for_openapi_name", false, "if set, the object's OpenAPI names will use the fully qualified names from the proto definition (ie my.package.MyMessage.MyInnerMessage). DEPRECATED: prefer `openapi_naming_strategy=fqn`")
 	openAPINamingStrategy          = flag.String("openapi_naming_strategy", "", "use the given OpenAPI naming strategy. Allowed values are `legacy`, `fqn`, `simple`. If unset, either `legacy` or `fqn` are selected, depending on the value of the `fqn_for_openapi_name` flag")
 	useGoTemplate                  = flag.Bool("use_go_templates", false, "if set, you can use Go templates in protofile comments")
+	goTemplateArgs                 = utilities.StringArrayFlag(flag.CommandLine, "go_template_args", "provide a custom value that can override a key in the Go template. Requires the `use_go_templates` option to be set")
 	ignoreComments                 = flag.Bool("ignore_comments", false, "if set, all protofile comments are excluded from output")
 	removeInternalComments         = flag.Bool("remove_internal_comments", false, "if set, removes all substrings in comments that start with `(--` and end with `--)` as specified in https://google.aip.dev/192#internal-comments")
 	disableDefaultErrors           = flag.Bool("disable_default_errors", false, "if set, disables generation of default errors. This is useful if you have defined custom error handling")
@@ -135,6 +136,12 @@ func main() {
 	reg.SetIgnoreComments(*ignoreComments)
 	reg.SetRemoveInternalComments(*removeInternalComments)
 
+	if len(*goTemplateArgs) > 0 && !*useGoTemplate {
+		emitError(fmt.Errorf("`go_template_args` requires `use_go_templates` to be enabled"))
+		return
+	}
+	reg.SetGoTemplateArgs(*goTemplateArgs)
+
 	reg.SetOpenAPINamingStrategy(namingStrategy)
 	reg.SetEnumsAsInts(*enumsAsInts)
 	reg.SetDisableDefaultErrors(*disableDefaultErrors)