From 1f441de351e700a6ab705c05a0425b15ce8c8651 Mon Sep 17 00:00:00 2001 From: bufdev Date: Thu, 19 Dec 2024 10:41:59 -0500 Subject: [PATCH] Add package_suffix option Signed-off-by: bufdev --- cmd/protoc-gen-connect-go/main.go | 19 +++++++----- cmd/protoc-gen-connect-go/main_test.go | 40 ++++++++++++++++++++++++-- 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/cmd/protoc-gen-connect-go/main.go b/cmd/protoc-gen-connect-go/main.go index 4b322ce3..917e1df4 100644 --- a/cmd/protoc-gen-connect-go/main.go +++ b/cmd/protoc-gen-connect-go/main.go @@ -65,7 +65,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." @@ -91,7 +92,11 @@ func main() { os.Exit(1) } var flagSet flag.FlagSet - samePackage := flagSet.Bool("same_package", false, "Generate files into the same Go package as the .pb.go base files.") + 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( @@ -101,7 +106,7 @@ func main() { plugin.SupportedEditionsMaximum = descriptorpb.Edition_EDITION_2023 for _, file := range plugin.Files { if file.Generate { - generate(plugin, file, *samePackage) + generate(plugin, file, *packageSuffix) } } return nil @@ -109,14 +114,14 @@ func main() { ) } -func generate(plugin *protogen.Plugin, file *protogen.File, samePackage bool) { +func generate(plugin *protogen.Plugin, file *protogen.File, packageSuffix string) { if len(file.Services) == 0 { return } goImportPath := file.GoImportPath - if !samePackage { - file.GoPackageName += generatedPackageSuffix + if packageSuffix != "" { + file.GoPackageName += protogen.GoPackageName(packageSuffix) generatedFilenamePrefixToSlash := filepath.ToSlash(file.GeneratedFilenamePrefix) file.GeneratedFilenamePrefix = path.Join( path.Dir(generatedFilenamePrefixToSlash), @@ -132,7 +137,7 @@ func generate(plugin *protogen.Plugin, file *protogen.File, samePackage bool) { file.GeneratedFilenamePrefix+generatedFilenameExtension, goImportPath, ) - if !samePackage { + if packageSuffix != "" { generatedFile.Import(file.GoImportPath) } generatePreamble(generatedFile, file) diff --git a/cmd/protoc-gen-connect-go/main_test.go b/cmd/protoc-gen-connect-go/main_test.go index 17a2b9b2..afc36ef4 100644 --- a/cmd/protoc-gen-connect-go/main_test.go +++ b/cmd/protoc-gen-connect-go/main_test.go @@ -69,11 +69,12 @@ func TestGenerate(t *testing.T) { assert.Equal(t, file.GetName(), "connectrpc.com/connect/internal/gen/connect/ping/v1/pingv1connect/ping.connect.go") assert.NotZero(t, file.GetContent()) }) - t.Run("ping.proto:same_package=true", func(t *testing.T) { + // Make sure empty package_suffix works alone. + t.Run("ping.proto:package_suffix=", func(t *testing.T) { t.Parallel() req := &pluginpb.CodeGeneratorRequest{ FileToGenerate: []string{"connect/ping/v1/ping.proto"}, - Parameter: ptr("same_package=true"), + Parameter: ptr("package_suffix="), ProtoFile: []*descriptorpb.FileDescriptorProto{pingFileDesc}, SourceFileDescriptors: []*descriptorpb.FileDescriptorProto{pingFileDesc}, CompilerVersion: compilerVersion, @@ -86,6 +87,41 @@ func TestGenerate(t *testing.T) { assert.Equal(t, file.GetName(), "connectrpc.com/connect/internal/gen/connect/ping/v1/ping.connect.go") assert.NotZero(t, file.GetContent()) }) + // Make sure empty package_suffix works with another option. + t.Run("ping.proto:package_suffix=,paths=source_relative", func(t *testing.T) { + t.Parallel() + req := &pluginpb.CodeGeneratorRequest{ + FileToGenerate: []string{"connect/ping/v1/ping.proto"}, + Parameter: ptr("package_suffix="), + ProtoFile: []*descriptorpb.FileDescriptorProto{pingFileDesc}, + SourceFileDescriptors: []*descriptorpb.FileDescriptorProto{pingFileDesc}, + CompilerVersion: compilerVersion, + } + rsp := testGenerate(t, req) + assert.Nil(t, rsp.Error) + + assert.Equal(t, len(rsp.File), 1) + file := rsp.File[0] + assert.Equal(t, file.GetName(), "connectrpc.com/connect/internal/gen/connect/ping/v1/ping.connect.go") + assert.NotZero(t, file.GetContent()) + }) + t.Run("ping.proto:package_suffix=baz", func(t *testing.T) { + t.Parallel() + req := &pluginpb.CodeGeneratorRequest{ + FileToGenerate: []string{"connect/ping/v1/ping.proto"}, + Parameter: ptr("package_suffix=baz"), + ProtoFile: []*descriptorpb.FileDescriptorProto{pingFileDesc}, + SourceFileDescriptors: []*descriptorpb.FileDescriptorProto{pingFileDesc}, + CompilerVersion: compilerVersion, + } + rsp := testGenerate(t, req) + assert.Nil(t, rsp.Error) + + assert.Equal(t, len(rsp.File), 1) + file := rsp.File[0] + assert.Equal(t, file.GetName(), "connectrpc.com/connect/internal/gen/connect/ping/v1/pingv1baz/ping.connect.go") + assert.NotZero(t, file.GetContent()) + }) } func testGenerate(t *testing.T, req *pluginpb.CodeGeneratorRequest) *pluginpb.CodeGeneratorResponse {