Add package_suffix option

Signed-off-by: bufdev <[email protected]>
connectrpc · Dec 19, 2024 · 1f441de · 1f441de
1 parent 70ee2a2
commit 1f441de
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 9 deletions.
diff --git 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,22 +106,22 @@ 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
 		},
 	)
 }
 
-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
@@ -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 {