Add --exclude-source-retention-options flag to buf build (#2807)

jhump · web-flow · commit fa6f9de4e245 · 2024-03-07T16:54:44.000-05:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,7 +7,11 @@
   field of the `CodeGeneratorRequest` message. This provides the plugin with access to options
   that are configured to only be retained in source and not at runtime (via
   [field option](https://github.com/protocolbuffers/protobuf/blob/v24.0/src/google/protobuf/descriptor.proto#L693-L702)).
-  Descriptors in the `proto_file` field will not include any options configured this way.
+  Descriptors in the `proto_file` field will not include any options configured this way
+  for the files named in `file_to_generate` field.
+- Add `--exclude-source-retention-options` flag to `buf build`, which
+  causes options configured to only be retained in source to be stripped
+  from the output descriptors.
 
 ## [v1.29.0] - 2024-01-24
 
diff --git a/private/buf/cmd/buf/command/build/build.go b/private/buf/cmd/buf/command/build/build.go
@@ -31,17 +31,18 @@ import (
 )
 
 const (
-	asFileDescriptorSetFlagName = "as-file-descriptor-set"
-	errorFormatFlagName         = "error-format"
-	excludeImportsFlagName      = "exclude-imports"
-	excludeSourceInfoFlagName   = "exclude-source-info"
-	pathsFlagName               = "path"
-	outputFlagName              = "output"
-	outputFlagShortName         = "o"
-	configFlagName              = "config"
-	excludePathsFlagName        = "exclude-path"
-	disableSymlinksFlagName     = "disable-symlinks"
-	typeFlagName                = "type"
+	asFileDescriptorSetFlagName           = "as-file-descriptor-set"
+	errorFormatFlagName                   = "error-format"
+	excludeImportsFlagName                = "exclude-imports"
+	excludeSourceInfoFlagName             = "exclude-source-info"
+	excludeSourceRetentionOptionsFlagName = "exclude-source-retention-options"
+	pathsFlagName                         = "path"
+	outputFlagName                        = "output"
+	outputFlagShortName                   = "o"
+	configFlagName                        = "config"
+	excludePathsFlagName                  = "exclude-path"
+	disableSymlinksFlagName               = "disable-symlinks"
+	typeFlagName                          = "type"
 )
 
 // NewCommand returns a new Command.
@@ -66,16 +67,17 @@ func NewCommand(
 }
 
 type flags struct {
-	AsFileDescriptorSet bool
-	ErrorFormat         string
-	ExcludeImports      bool
-	ExcludeSourceInfo   bool
-	Paths               []string
-	Output              string
-	Config              string
-	ExcludePaths        []string
-	DisableSymlinks     bool
-	Types               []string
+	AsFileDescriptorSet           bool
+	ErrorFormat                   string
+	ExcludeImports                bool
+	ExcludeSourceInfo             bool
+	ExcludeSourceRetentionOptions bool
+	Paths                         []string
+	Output                        string
+	Config                        string
+	ExcludePaths                  []string
+	DisableSymlinks               bool
+	Types                         []string
 	// special
 	InputHashtag string
 }
@@ -92,6 +94,12 @@ func (f *flags) Bind(flagSet *pflag.FlagSet) {
 	bufcli.BindPaths(flagSet, &f.Paths, pathsFlagName)
 	bufcli.BindExcludePaths(flagSet, &f.ExcludePaths, excludePathsFlagName)
 	bufcli.BindDisableSymlinks(flagSet, &f.DisableSymlinks, disableSymlinksFlagName)
+	flagSet.BoolVar(
+		&f.ExcludeSourceRetentionOptions,
+		excludeSourceRetentionOptionsFlagName,
+		false,
+		"Exclude options whose retention is source",
+	)
 	flagSet.StringVar(
 		&f.ErrorFormat,
 		errorFormatFlagName,
@@ -165,6 +173,12 @@ func run(
 			return err
 		}
 	}
+	if flags.ExcludeSourceRetentionOptions {
+		image, err = bufimageutil.StripSourceRetentionOptions(image)
+		if err != nil {
+			return err
+		}
+	}
 	return bufcli.NewWireImageWriter(
 		container.Logger(),
 	).PutImage(
diff --git a/private/bufpkg/bufimage/bufimageutil/bufimageutil.go b/private/bufpkg/bufimage/bufimageutil/bufimageutil.go
@@ -22,6 +22,7 @@ import (
 
 	"github.com/bufbuild/buf/private/bufpkg/bufimage"
 	"github.com/bufbuild/buf/private/pkg/protosource"
+	"github.com/bufbuild/protocompile/options"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/types/descriptorpb"
@@ -372,6 +373,21 @@ func ImageFilteredByTypesWithOptions(image bufimage.Image, types []string, opts
 	return bufimage.NewImage(includedFiles)
 }
 
+// StripSourceRetentionOptions strips any options with a retention of "source" from
+// the descriptors in the given image. The image is not mutated but instead a new
+// image is returned. The returned image may share state with the original.
+func StripSourceRetentionOptions(image bufimage.Image) (bufimage.Image, error) {
+	updatedFiles := make([]bufimage.ImageFile, len(image.Files()))
+	for i, inputFile := range image.Files() {
+		updatedFile, err := stripSourceRetentionOptionsFromFile(inputFile)
+		if err != nil {
+			return nil, fmt.Errorf("failed to strip source-retention options from file %q: %w", inputFile.Path(), err)
+		}
+		updatedFiles[i] = updatedFile
+	}
+	return bufimage.NewImage(updatedFiles)
+}
+
 // trimMessageDescriptors removes (nested) messages and nested enums from a slice
 // of message descriptors if their type names are not found in the toKeep map.
 func trimMessageDescriptors(
@@ -917,3 +933,19 @@ func newImageFilterOptions() *imageFilterOptions {
 		allowImportedTypes:     false,
 	}
 }
+
+func stripSourceRetentionOptionsFromFile(imageFile bufimage.ImageFile) (bufimage.ImageFile, error) {
+	updatedFileDescriptor, err := options.StripSourceRetentionOptionsFromFile(imageFile.FileDescriptorProto())
+	if err != nil {
+		return nil, err
+	}
+	return bufimage.NewImageFile(
+		updatedFileDescriptor,
+		imageFile.ModuleIdentity(),
+		imageFile.Commit(),
+		imageFile.ExternalPath(),
+		imageFile.IsImport(),
+		imageFile.IsSyntaxUnspecified(),
+		imageFile.UnusedDependencyIndexes(),
+	)
+}
