Add package_suffix option to protoc-gen-connect-go

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 {