From 46c38b351a226b11ec734b7ac2454065aa496693 Mon Sep 17 00:00:00 2001 From: Hector Oliveros Date: Fri, 12 Jul 2024 12:47:32 -0400 Subject: [PATCH] feat: Add functionality to export proto files (#475) * Add functionality to export proto files Added a new function, `WriteProtoFiles` in `desc_source.go` which is used to generate .proto files. The process involves resolving symbols from the descriptor source and writing their definitions to a designated output directory. The corresponding flag `--proto-out` has been included in `grpcurl.go` to allow users to specify the directory path. * Refactor file creation error handling The code for file creation and error handling in desc_source.go has been refactored. Previously, the file closure operation was executed irrespective of whether the file was created successfully or not. Now, the file will only be closed if it was successfully created, improving error handling. * Update 'proto-out' command to 'proto-out-dir' The command for exporting proto files and setting the output directory has been updated from 'proto-out' to 'proto-out-dir'. This change has been made both in the README and the grpcurl go file. This makes the command name more descriptive, accurately reflecting its functionality. * fix import sort * Reorder file close operation in error handling The file close operation has been moved within the error handling of the 'PrintProtoFile' function. Previously, it was being executed before this function, now it's executed immediately after. Moreover, an additional close operation has been added after the function success ensuring the file is properly closed in all scenarios. * Update grpcurl commands in README The grpcurl commands for exporting proto files and protoset files in the README are updated. The IP address has been changed to localhost and port number to '8787'. Also, the service name is adjusted to 'my.custom.server.Service'. Instructions for use of specific command options are added for enhanced clarity. * Refactor error handling in file creation The code responsible for error handling during file creation in the desc_source.go file has been streamlined. This modification simplifies the code by reducing unnecessary condition checks and redundant file closure action after an error has occurred. * Refactor proto file writing into separate function The file writing process for protobuf files has been extracted into a new function called writeProtoFile(). This refactoring simplifies the main function. The code is cleaner and more manageable this way, improving maintainability and readability. * Refactor writeProtoFile function for better error handling Streamlined the writeProtoFile function in desc_source.go file. Simplified path calculations and improved error messages for file-creation functions, making it easier to trace the exact point of failure and enhance the debugging process. --- README.md | 7 ++++ cmd/grpcurl/grpcurl.go | 24 +++++++++++ desc_source.go | 91 ++++++++++++++++++++++++++++++++++++------ 3 files changed, 109 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 49960186..85e923b0 100644 --- a/README.md +++ b/README.md @@ -159,6 +159,13 @@ grpcurl -protoset my-protos.bin list # Using proto sources grpcurl -import-path ../protos -proto my-stuff.proto list + +# Export proto files (use -proto-out-dir to specify the output directory) +grpcurl -plaintext -proto-out-dir "out_protos" "localhost:8787" describe my.custom.server.Service + +# Export protoset file (use -protoset-out to specify the output file) +grpcurl -plaintext -protoset-out "out.protoset" "localhost:8787" describe my.custom.server.Service + ``` The "list" verb also lets you see all methods in a particular service: diff --git a/cmd/grpcurl/grpcurl.go b/cmd/grpcurl/grpcurl.go index 70d079bc..80f87bfc 100644 --- a/cmd/grpcurl/grpcurl.go +++ b/cmd/grpcurl/grpcurl.go @@ -151,6 +151,14 @@ var ( file if this option is given. When invoking an RPC and this option is given, the method being invoked and its transitive dependencies will be included in the output file.`)) + protoOut = flags.String("proto-out-dir", "", prettify(` + The name of a directory where the generated .proto files will be written. + With the list and describe verbs, the listed or described elements and + their transitive dependencies will be written as .proto files in the + specified directory if this option is given. When invoking an RPC and + this option is given, the method being invoked and its transitive + dependencies will be included in the generated .proto files in the + output directory.`)) msgTemplate = flags.Bool("msg-template", false, prettify(` When describing messages, show a template of input data.`)) verbose = flags.Bool("v", false, prettify(` @@ -645,6 +653,9 @@ func main() { if err := writeProtoset(descSource, svcs...); err != nil { fail(err, "Failed to write protoset to %s", *protosetOut) } + if err := writeProtos(descSource, svcs...); err != nil { + fail(err, "Failed to write protos to %s", *protoOut) + } } else { methods, err := grpcurl.ListMethods(descSource, symbol) if err != nil { @@ -660,6 +671,9 @@ func main() { if err := writeProtoset(descSource, symbol); err != nil { fail(err, "Failed to write protoset to %s", *protosetOut) } + if err := writeProtos(descSource, symbol); err != nil { + fail(err, "Failed to write protos to %s", *protoOut) + } } } else if describe { @@ -764,6 +778,9 @@ func main() { if err := writeProtoset(descSource, symbols...); err != nil { fail(err, "Failed to write protoset to %s", *protosetOut) } + if err := writeProtos(descSource, symbol); err != nil { + fail(err, "Failed to write protos to %s", *protoOut) + } } else { // Invoke an RPC @@ -923,6 +940,13 @@ func writeProtoset(descSource grpcurl.DescriptorSource, symbols ...string) error return grpcurl.WriteProtoset(f, descSource, symbols...) } +func writeProtos(descSource grpcurl.DescriptorSource, symbols ...string) error { + if *protoOut == "" { + return nil + } + return grpcurl.WriteProtoFiles(*protoOut, descSource, symbols...) +} + type optionalBoolFlag struct { set, val bool } diff --git a/desc_source.go b/desc_source.go index d12fae0d..87f1403b 100644 --- a/desc_source.go +++ b/desc_source.go @@ -6,11 +6,13 @@ import ( "fmt" "io" "os" + "path/filepath" "sync" "github.com/golang/protobuf/proto" //lint:ignore SA1019 we have to import this because it appears in exported API "github.com/jhump/protoreflect/desc" "github.com/jhump/protoreflect/desc/protoparse" + "github.com/jhump/protoreflect/desc/protoprint" "github.com/jhump/protoreflect/dynamic" "github.com/jhump/protoreflect/grpcreflect" "google.golang.org/grpc/codes" @@ -258,19 +260,9 @@ func reflectionSupport(err error) error { // given output. The output will include descriptors for all files in which the // symbols are defined as well as their transitive dependencies. func WriteProtoset(out io.Writer, descSource DescriptorSource, symbols ...string) error { - // compute set of file descriptors - filenames := make([]string, 0, len(symbols)) - fds := make(map[string]*desc.FileDescriptor, len(symbols)) - for _, sym := range symbols { - d, err := descSource.FindSymbol(sym) - if err != nil { - return fmt.Errorf("failed to find descriptor for %q: %v", sym, err) - } - fd := d.GetFile() - if _, ok := fds[fd.GetName()]; !ok { - fds[fd.GetName()] = fd - filenames = append(filenames, fd.GetName()) - } + filenames, fds, err := getFileDescriptors(symbols, descSource) + if err != nil { + return err } // now expand that to include transitive dependencies in topologically sorted // order (such that file always appears after its dependencies) @@ -302,3 +294,76 @@ func addFilesToSet(allFiles []*descriptorpb.FileDescriptorProto, expanded map[st } return append(allFiles, fd.AsFileDescriptorProto()) } + +// WriteProtoFiles will use the given descriptor source to resolve all the given +// symbols and write proto files with their definitions to the given output directory. +func WriteProtoFiles(outProtoDirPath string, descSource DescriptorSource, symbols ...string) error { + filenames, fds, err := getFileDescriptors(symbols, descSource) + if err != nil { + return err + } + // now expand that to include transitive dependencies in topologically sorted + // order (such that file always appears after its dependencies) + expandedFiles := make(map[string]struct{}, len(fds)) + allFileDescriptors := make([]*desc.FileDescriptor, 0, len(fds)) + for _, filename := range filenames { + allFileDescriptors = addFilesToFileDescriptorList(allFileDescriptors, expandedFiles, fds[filename]) + } + pr := protoprint.Printer{} + // now we can serialize to files + for i := range allFileDescriptors { + if err := writeProtoFile(outProtoDirPath, allFileDescriptors[i], &pr); err != nil { + return err + } + } + return nil +} + +func writeProtoFile(outProtoDirPath string, fd *desc.FileDescriptor, pr *protoprint.Printer) error { + outFile := filepath.Join(outProtoDirPath, fd.GetFullyQualifiedName()) + outDir := filepath.Dir(outFile) + if err := os.MkdirAll(outDir, 0777); err != nil { + return fmt.Errorf("failed to create directory %q: %w", outDir, err) + } + + f, err := os.Create(outFile) + if err != nil { + return fmt.Errorf("failed to create proto file %q: %w", outFile, err) + } + defer f.Close() + if err := pr.PrintProtoFile(fd, f); err != nil { + return fmt.Errorf("failed to write proto file %q: %w", outFile, err) + } + return nil +} + +func getFileDescriptors(symbols []string, descSource DescriptorSource) ([]string, map[string]*desc.FileDescriptor, error) { + // compute set of file descriptors + filenames := make([]string, 0, len(symbols)) + fds := make(map[string]*desc.FileDescriptor, len(symbols)) + for _, sym := range symbols { + d, err := descSource.FindSymbol(sym) + if err != nil { + return nil, nil, fmt.Errorf("failed to find descriptor for %q: %v", sym, err) + } + fd := d.GetFile() + if _, ok := fds[fd.GetName()]; !ok { + fds[fd.GetName()] = fd + filenames = append(filenames, fd.GetName()) + } + } + return filenames, fds, nil +} + +func addFilesToFileDescriptorList(allFiles []*desc.FileDescriptor, expanded map[string]struct{}, fd *desc.FileDescriptor) []*desc.FileDescriptor { + if _, ok := expanded[fd.GetName()]; ok { + // already seen this one + return allFiles + } + expanded[fd.GetName()] = struct{}{} + // add all dependencies first + for _, dep := range fd.GetDependencies() { + allFiles = addFilesToFileDescriptorList(allFiles, expanded, dep) + } + return append(allFiles, fd) +}