Skip to content

Commit

Permalink
cmd/protoc-gen-go: support protobuf retention feature
Browse files Browse the repository at this point in the history
This change strips out all descriptor option fields marked with [retention = RETENTION_SOURCE] before writing the FileDescriptor to the generated go code bindings.

Change-Id: Ie624d9d4b4f211a256661d80a04199ae8401662b
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/472696
Reviewed-by: Damien Neil <dneil@google.com>
Reviewed-by: Lasse Folger <lassefolger@google.com>
  • Loading branch information
googleberg authored and stapelberg committed Mar 7, 2023
1 parent fcf5f6c commit eba8b09
Show file tree
Hide file tree
Showing 8 changed files with 1,478 additions and 1 deletion.
23 changes: 22 additions & 1 deletion cmd/protoc-gen-go/internal_gengo/reflect.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ import (

"google.golang.org/protobuf/compiler/protogen"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/reflect/protopath"
"google.golang.org/protobuf/reflect/protorange"
"google.golang.org/protobuf/reflect/protoreflect"

"google.golang.org/protobuf/types/descriptorpb"
Expand Down Expand Up @@ -233,10 +235,29 @@ func genReflectFileDescriptor(gen *protogen.Plugin, g *protogen.GeneratedFile, f
g.P("}")
}

// stripSourceRetentionFieldsFromMessage walks the given message tree recursively
// and clears any fields with the field option: [retention = RETENTION_SOURCE]
func stripSourceRetentionFieldsFromMessage(m protoreflect.Message) {
protorange.Range(m, func(ppv protopath.Values) error {
m2, ok := ppv.Index(-1).Value.Interface().(protoreflect.Message)
if !ok {
return nil
}
m2.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
fdo, ok := fd.Options().(*descriptorpb.FieldOptions)
if ok && fdo.GetRetention() == descriptorpb.FieldOptions_RETENTION_SOURCE {
m2.Clear(fd)
}
return true
})
return nil
})
}

func genFileDescriptor(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo) {
descProto := proto.Clone(f.Proto).(*descriptorpb.FileDescriptorProto)
descProto.SourceCodeInfo = nil // drop source code information

stripSourceRetentionFieldsFromMessage(descProto.ProtoReflect())
b, err := proto.MarshalOptions{AllowPartial: true, Deterministic: true}.Marshal(descProto)
if err != nil {
gen.Error(err)
Expand Down
178 changes: 178 additions & 0 deletions cmd/protoc-gen-go/retention_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,178 @@
// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
"testing"

"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/reflect/protoreflect"

retentionpb "google.golang.org/protobuf/cmd/protoc-gen-go/testdata/retention"
)

func TestFileOptionRetention(t *testing.T) {
options := retentionpb.File_cmd_protoc_gen_go_testdata_retention_retention_proto.Options()
tests := []struct {
name string
ext protoreflect.ExtensionType
wantField bool
wantValue int32
}{
{
name: "imported_plain_option",
ext: retentionpb.E_ImportedPlainOption,
wantField: true,
wantValue: 1,
},
{
name: "imported_runtime_option",
ext: retentionpb.E_ImportedRuntimeRetentionOption,
wantField: true,
wantValue: 2,
},
{
name: "imported_source_option",
ext: retentionpb.E_ImportedSourceRetentionOption,
wantField: false,
wantValue: 0,
},
{
name: "plain_option",
ext: retentionpb.E_PlainOption,
wantField: true,
wantValue: 1,
},
{
name: "runtime_option",
ext: retentionpb.E_RuntimeRetentionOption,
wantField: true,
wantValue: 2,
},
{
name: "source_option",
ext: retentionpb.E_SourceRetentionOption,
wantField: false,
wantValue: 0,
},
}

for _, test := range tests {
if test.wantField != proto.HasExtension(options, test.ext) {
t.Errorf("HasExtension(%s): got %v, want %v", test.name, proto.HasExtension(options, test.ext), test.wantField)
}
if test.wantValue != proto.GetExtension(options, test.ext).(int32) {
t.Errorf("GetExtension(%s): got %d, want %d", test.name, proto.GetExtension(options, test.ext).(int32), test.wantValue)
}
}
}

func TestAllEntitiesWithMessageOption(t *testing.T) {
file := retentionpb.File_cmd_protoc_gen_go_testdata_retention_retention_proto
verifyDescriptorOptions(t, string(file.Name()), file.Options())
verifyEnums(t, file.Enums())
verifyMessages(t, file.Messages())
verifyExtensions(t, file.Extensions())
verifyServices(t, file.Services())
}

func verifyExtensions(t *testing.T, extensions protoreflect.ExtensionDescriptors) {
t.Helper()
for i := 0; i < extensions.Len(); i++ {
verifyDescriptorOptions(t, string(extensions.Get(i).Name()), extensions.Get(i).Options())
}
}

func verifyMessages(t *testing.T, messages protoreflect.MessageDescriptors) {
t.Helper()
for i := 0; i < messages.Len(); i++ {
verifyDescriptorOptions(t, string(messages.Get(i).Name()), messages.Get(i).Options())
verifyEnums(t, messages.Get(i).Enums())
verifyMessages(t, messages.Get(i).Messages())
verifyExtensions(t, messages.Get(i).Extensions())
verifyFields(t, messages.Get(i).Fields())
}
}

func verifyFields(t *testing.T, fields protoreflect.FieldDescriptors) {
t.Helper()
for i := 0; i < fields.Len(); i++ {
verifyDescriptorOptions(t, string(fields.Get(i).Name()), fields.Get(i).Options())
}
}

func verifyEnums(t *testing.T, enums protoreflect.EnumDescriptors) {
t.Helper()
for i := 0; i < enums.Len(); i++ {
verifyDescriptorOptions(t, string(enums.Get(i).Name()), enums.Get(i).Options())
verifyEnumValues(t, enums.Get(i).Values())
}
}

func verifyEnumValues(t *testing.T, values protoreflect.EnumValueDescriptors) {
t.Helper()
for i := 0; i < values.Len(); i++ {
verifyDescriptorOptions(t, string(values.Get(i).Name()), values.Get(i).Options())
}
}

func verifyServices(t *testing.T, services protoreflect.ServiceDescriptors) {
t.Helper()
for i := 0; i < services.Len(); i++ {
verifyDescriptorOptions(t, string(services.Get(i).Name()), services.Get(i).Options())
verifyMethods(t, services.Get(i).Methods())
}
}

func verifyMethods(t *testing.T, methods protoreflect.MethodDescriptors) {
t.Helper()
for i := 0; i < methods.Len(); i++ {
verifyDescriptorOptions(t, string(methods.Get(i).Name()), methods.Get(i).Options())
}
}

func verifyDescriptorOptions(t *testing.T, entity string, options protoreflect.ProtoMessage) {
t.Helper()
options.ProtoReflect().Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
maybeVerifyOption(t, fd, v)
return true
})
}

func maybeVerifyOption(t *testing.T, fd protoreflect.FieldDescriptor, v protoreflect.Value) {
t.Helper()
if fd.Kind() == protoreflect.MessageKind && string(fd.Message().FullName()) == "goproto.proto.testretention.OptionsMessage" {
if fd.IsList() {
for i := 0; i < v.List().Len(); i++ {
verifyOptionsMessage(t, string(fd.FullName()), v.List().Get(i).Message().Interface().(*retentionpb.OptionsMessage))
}
} else {
verifyOptionsMessage(t, string(fd.FullName()), v.Message().Interface().(*retentionpb.OptionsMessage))
}
}
}

func verifyOptionsMessage(t *testing.T, entity string, msg *retentionpb.OptionsMessage) {
t.Helper()
if msg.PlainField == nil {
t.Errorf("%s.OptionsMessage.HasField(plain_field): got false, want true", entity)
}
if msg.GetPlainField() != 1 {
t.Errorf("%s.OptionsMessage.GetField(plain_field): got %d, want 1", entity, msg.GetPlainField())
}
if msg.RuntimeRetentionField == nil {
t.Errorf("%s.OptionsMessage.HasField(runtime_retention_field): got false, want true", entity)
}
if msg.GetRuntimeRetentionField() != 2 {
t.Errorf("%s.OptionsMessage.GetField(runtime_retention_field): got %d, want 2", entity, msg.GetRuntimeRetentionField())
}
if msg.SourceRetentionField != nil {
t.Errorf("%s.OptionsMessage.HasField(source_retention_field): got true, want false", entity)
}
if msg.GetSourceRetentionField() != 0 {
// Checking that we get 0 even though this was set to 3 in the source file
t.Errorf("%s.OptionsMessage.GetField(source_retention_field): got %d, want 0", entity, msg.GetSourceRetentionField())
}
}
1 change: 1 addition & 0 deletions cmd/protoc-gen-go/testdata/gen_test.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit eba8b09

Please sign in to comment.