Add same_package option to protoc-gen-connect-go (#803)

This adds an option `package_suffix` to `protoc-gen-connect-go`. It allows to
customize the output directory for the generated code. The sub-package of the
package containing the base .pb.go files using the given suffix. An empty
suffix denotes to generate into the same package as the base pb.go files.
Default is "connect".

---------

Signed-off-by: bufdev <[email protected]>
Signed-off-by: Edward McFarlane <[email protected]>
Co-authored-by: Edward McFarlane <[email protected]>

Author: bufdev

Makefile

@@ -11,7 +11,7 @@ export PATH := $(BIN):$(PATH)
 export GOBIN := $(abspath $(BIN))
 COPYRIGHT_YEARS := 2021-2024
 LICENSE_IGNORE := --ignore /testdata/
-BUF_VERSION := 1.34.0
+BUF_VERSION := 1.47.2
 
 .PHONY: help
 help: ## Describe useful make targets
@@ -81,6 +81,7 @@ generate: $(BIN)/buf $(BIN)/protoc-gen-go $(BIN)/protoc-gen-connect-go $(BIN)/li
 	go mod tidy
 	rm -rf internal/gen
 	PATH="$(abspath $(BIN))" buf generate
+	( cd cmd/protoc-gen-connect-go; ./generate.sh )
 	license-header \
 		--license-type apache \
 		--copyright-holder "The Connect Authors" \

buf.gen.yaml

@@ -11,3 +11,4 @@ plugins:
   - local: protoc-gen-connect-go
     out: internal/gen
     opt: paths=source_relative
+clean: true

cmd/protoc-gen-connect-go/generate.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+find testdata -maxdepth 1 -type d \( ! -name testdata \) -exec bash -c "cd '{}' && buf generate" \;

cmd/protoc-gen-connect-go/main.go

@@ -23,26 +23,50 @@
 //
 // With [buf], your buf.gen.yaml will look like this:
 //
-//	version: v1
+//	version: v2
 //	plugins:
-//	  - name: go
+//	  - local: protoc-gen-go
 //	    out: gen
-//	  - name: connect-go
+//	  - local: protoc-gen-connect-go
 //	    out: gen
 //
 // This generates service definitions for the Protobuf types and services
 // defined by file.proto. If file.proto defines the foov1 Protobuf package, the
 // invocations above will write output to:
 //
 //	gen/path/to/file.pb.go
-//	gen/path/to/connectfoov1/file.connect.go
+//	gen/path/to/foov1connect/file.connect.go
+//
+// The generated code is configurable with the same parameters as the protoc-gen-go
+// plugin, with the following additional parameters:
+//
+//   - package_suffix: To generate into a sub-package of the package containing the
+//     base .pb.go files using the given suffix. An empty suffix denotes to
+//     generate into the same package as the base pb.go files. Default is "connect".
+//
+// For example, to generate into the same package as the base .pb.go files:
+//
+//	version: v2
+//	plugins:
+//	  - local: protoc-gen-go
+//	    out: gen
+//	  - local: protoc-gen-connect-go
+//	    out: gen
+//	    opts: package_suffix
+//
+// This will generate output to:
+//
+//	gen/path/to/file.pb.go
+//	gen/path/to/file.connect.go
 //
 // [buf]: https://buf.build
 package main
 
 import (
 	"bytes"
+	"flag"
 	"fmt"
+	"go/token"
 	"os"
 	"path"
 	"path/filepath"
@@ -64,7 +88,8 @@ const (
 	connectPackage = protogen.GoImportPath("connectrpc.com/connect")
 
 	generatedFilenameExtension = ".connect.go"
-	generatedPackageSuffix     = "connect"
+	defaultPackageSuffix       = "connect"
+	packageSuffixFlagName      = "package_suffix"
 
 	usage = "See https://connectrpc.com/docs/go/getting-started to learn how to use this plugin.\n\nFlags:\n  -h, --help\tPrint this help and exit.\n      --version\tPrint the version and exit."
 
@@ -89,46 +114,63 @@ func main() {
 		fmt.Fprintln(os.Stderr, usage)
 		os.Exit(1)
 	}
-	protogen.Options{}.Run(
+	var flagSet flag.FlagSet
+	packageSuffix := flagSet.String(
+		packageSuffixFlagName,
+		defaultPackageSuffix,
+		"Generate files into a sub-package of the package containing the base .pb.go files using the given suffix. An empty suffix denotes to generate into the same package as the base pb.go files.",
+	)
+	protogen.Options{
+		ParamFunc: flagSet.Set,
+	}.Run(
 		func(plugin *protogen.Plugin) error {
 			plugin.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL) | uint64(pluginpb.CodeGeneratorResponse_FEATURE_SUPPORTS_EDITIONS)
 			plugin.SupportedEditionsMinimum = descriptorpb.Edition_EDITION_PROTO2
 			plugin.SupportedEditionsMaximum = descriptorpb.Edition_EDITION_2023
 			for _, file := range plugin.Files {
 				if file.Generate {
-					generate(plugin, file)
+					generate(plugin, file, *packageSuffix)
 				}
 			}
 			return nil
 		},
 	)
 }
 
-func generate(plugin *protogen.Plugin, file *protogen.File) {
+func generate(plugin *protogen.Plugin, file *protogen.File, packageSuffix string) {
 	if len(file.Services) == 0 {
 		return
 	}
-	file.GoPackageName += generatedPackageSuffix
 
-	generatedFilenamePrefixToSlash := filepath.ToSlash(file.GeneratedFilenamePrefix)
-	file.GeneratedFilenamePrefix = path.Join(
-		path.Dir(generatedFilenamePrefixToSlash),
-		string(file.GoPackageName),
-		path.Base(generatedFilenamePrefixToSlash),
-	)
-	generatedFile := plugin.NewGeneratedFile(
-		file.GeneratedFilenamePrefix+generatedFilenameExtension,
-		protogen.GoImportPath(path.Join(
+	goImportPath := file.GoImportPath
+	if packageSuffix != "" {
+		if !token.IsIdentifier(packageSuffix) {
+			plugin.Error(fmt.Errorf("package_suffix %q is not a valid Go identifier", packageSuffix))
+			return
+		}
+		file.GoPackageName += protogen.GoPackageName(packageSuffix)
+		generatedFilenamePrefixToSlash := filepath.ToSlash(file.GeneratedFilenamePrefix)
+		file.GeneratedFilenamePrefix = path.Join(
+			path.Dir(generatedFilenamePrefixToSlash),
+			string(file.GoPackageName),
+			path.Base(generatedFilenamePrefixToSlash),
+		)
+		goImportPath = protogen.GoImportPath(path.Join(
 			string(file.GoImportPath),
 			string(file.GoPackageName),
-		)),
+		))
+	}
+	generatedFile := plugin.NewGeneratedFile(
+		file.GeneratedFilenamePrefix+generatedFilenameExtension,
+		goImportPath,
 	)
-	generatedFile.Import(file.GoImportPath)
+	if packageSuffix != "" {
+		generatedFile.Import(file.GoImportPath)
+	}
 	generatePreamble(generatedFile, file)
 	generateServiceNameConstants(generatedFile, file.Services)
-	generateServiceNameVariables(generatedFile, file)
 	for _, service := range file.Services {
-		generateService(generatedFile, service)
+		generateService(generatedFile, file, service)
 	}
 }
 
@@ -213,29 +255,22 @@ func generateServiceNameConstants(g *protogen.GeneratedFile, services []*protoge
 	g.P()
 }
 
-func generateServiceNameVariables(g *protogen.GeneratedFile, file *protogen.File) {
-	wrapComments(g, "These variables are the protoreflect.Descriptor objects for the RPCs defined in this package.")
-	g.P("var (")
-	for _, service := range file.Services {
-		serviceDescName := unexport(fmt.Sprintf("%sServiceDescriptor", service.Desc.Name()))
-		g.P(serviceDescName, ` = `,
-			g.QualifiedGoIdent(file.GoDescriptorIdent),
-			`.Services().ByName("`, service.Desc.Name(), `")`)
-		for _, method := range service.Methods {
-			g.P(procedureVarMethodDescriptor(method), ` = `,
-				serviceDescName,
-				`.Methods().ByName("`, method.Desc.Name(), `")`)
-		}
+func generateServiceMethodsVar(g *protogen.GeneratedFile, file *protogen.File, service *protogen.Service) {
+	if len(service.Methods) == 0 {
+		return
 	}
-	g.P(")")
+	serviceMethodsName := unexport(fmt.Sprintf("%sMethods", service.Desc.Name()))
+	g.P(serviceMethodsName, ` := `,
+		g.QualifiedGoIdent(file.GoDescriptorIdent),
+		`.Services().ByName("`, service.Desc.Name(), `").Methods()`)
 }
 
-func generateService(g *protogen.GeneratedFile, service *protogen.Service) {
+func generateService(g *protogen.GeneratedFile, file *protogen.File, service *protogen.Service) {
 	names := newNames(service)
 	generateClientInterface(g, service, names)
-	generateClientImplementation(g, service, names)
+	generateClientImplementation(g, file, service, names)
 	generateServerInterface(g, service, names)
-	generateServerConstructor(g, service, names)
+	generateServerConstructor(g, file, service, names)
 	generateUnimplementedServerImplementation(g, service, names)
 }
 
@@ -260,7 +295,7 @@ func generateClientInterface(g *protogen.GeneratedFile, service *protogen.Servic
 	g.P()
 }
 
-func generateClientImplementation(g *protogen.GeneratedFile, service *protogen.Service, names names) {
+func generateClientImplementation(g *protogen.GeneratedFile, file *protogen.File, service *protogen.Service, names names) {
 	clientOption := connectPackage.Ident("ClientOption")
 
 	// Client constructor.
@@ -281,6 +316,7 @@ func generateClientImplementation(g *protogen.GeneratedFile, service *protogen.S
 	if len(service.Methods) > 0 {
 		g.P("baseURL = ", stringsPackage.Ident("TrimRight"), `(baseURL, "/")`)
 	}
+	generateServiceMethodsVar(g, file, service)
 	g.P("return &", names.ClientImpl, "{")
 	for _, method := range service.Methods {
 		g.P(unexport(method.GoName), ": ",
@@ -396,7 +432,7 @@ func generateServerInterface(g *protogen.GeneratedFile, service *protogen.Servic
 	g.P()
 }
 
-func generateServerConstructor(g *protogen.GeneratedFile, service *protogen.Service, names names) {
+func generateServerConstructor(g *protogen.GeneratedFile, file *protogen.File, service *protogen.Service, names names) {
 	wrapComments(g, names.ServerConstructor, " builds an HTTP handler from the service implementation.",
 		" It returns the path on which to mount the handler and the handler itself.")
 	g.P("//")
@@ -409,6 +445,7 @@ func generateServerConstructor(g *protogen.GeneratedFile, service *protogen.Serv
 	handlerOption := connectPackage.Ident("HandlerOption")
 	g.P("func ", names.ServerConstructor, "(svc ", names.Server, ", opts ...", handlerOption,
 		") (string, ", httpPackage.Ident("Handler"), ") {")
+	generateServiceMethodsVar(g, file, service)
 	for _, method := range service.Methods {
 		isStreamingServer := method.Desc.IsStreamingServer()
 		isStreamingClient := method.Desc.IsStreamingClient()
@@ -522,7 +559,8 @@ func procedureHandlerName(m *protogen.Method) string {
 }
 
 func procedureVarMethodDescriptor(m *protogen.Method) string {
-	return unexport(fmt.Sprintf("%s%sMethodDescriptor", m.Parent.GoName, m.GoName))
+	serviceMethodsName := unexport(fmt.Sprintf("%sMethods", m.Parent.GoName))
+	return serviceMethodsName + `.ByName("` + string(m.Desc.Name()) + `")`
 }
 
 func isDeprecatedService(service *protogen.Service) bool {