Skip to content

Commit

Permalink
feat: Add functionality to export proto files (#475)
Browse files Browse the repository at this point in the history
* 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.
  • Loading branch information
Eitol authored Jul 12, 2024
1 parent a05d48d commit 46c38b3
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 13 deletions.
7 changes: 7 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
24 changes: 24 additions & 0 deletions cmd/grpcurl/grpcurl.go
Original file line number Diff line number Diff line change
Expand Up @@ -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(`
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand Down
91 changes: 78 additions & 13 deletions desc_source.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
}

0 comments on commit 46c38b3

Please sign in to comment.