github.com/mloves0824/enron/cmd/enron@v0.0.0-20230830012320-113bbf6be3c8/internal/proto/client/client.go (about)

     1  package client
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"os/exec"
     7  	"path/filepath"
     8  	"regexp"
     9  	"strings"
    10  
    11  	"github.com/spf13/cobra"
    12  
    13  	"github.com/mloves0824/enron/cmd/enron/internal/base"
    14  )
    15  
    16  // CmdClient represents the source command.
    17  var CmdClient = &cobra.Command{
    18  	Use:   "client",
    19  	Short: "Generate the proto client code",
    20  	Long:  "Generate the proto client code. Example: enron proto client helloworld.proto",
    21  	Run:   run,
    22  }
    23  
    24  var protoPath string
    25  
    26  func init() {
    27  	if protoPath = os.Getenv("ENRON_PROTO_PATH"); protoPath == "" {
    28  		protoPath = "./third_party"
    29  	}
    30  	CmdClient.Flags().StringVarP(&protoPath, "proto_path", "p", protoPath, "proto path")
    31  }
    32  
    33  func run(_ *cobra.Command, args []string) {
    34  	if len(args) == 0 {
    35  		fmt.Println("Please enter the proto file or directory")
    36  		return
    37  	}
    38  	var (
    39  		err   error
    40  		proto = strings.TrimSpace(args[0])
    41  	)
    42  	if err = look("protoc-gen-go", "protoc-gen-go-grpc", "protoc-gen-go-http", "protoc-gen-go-errors", "protoc-gen-openapi"); err != nil {
    43  		// update the enron plugins
    44  		cmd := exec.Command("enron", "upgrade")
    45  		cmd.Stdout = os.Stdout
    46  		cmd.Stderr = os.Stderr
    47  		if err = cmd.Run(); err != nil {
    48  			fmt.Println(err)
    49  			return
    50  		}
    51  	}
    52  	if strings.HasSuffix(proto, ".proto") {
    53  		err = generate(proto, args)
    54  	} else {
    55  		err = walk(proto, args)
    56  	}
    57  	if err != nil {
    58  		fmt.Println(err)
    59  	}
    60  }
    61  
    62  func look(name ...string) error {
    63  	for _, n := range name {
    64  		if _, err := exec.LookPath(n); err != nil {
    65  			return err
    66  		}
    67  	}
    68  	return nil
    69  }
    70  
    71  func walk(dir string, args []string) error {
    72  	if dir == "" {
    73  		dir = "."
    74  	}
    75  	return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
    76  		if ext := filepath.Ext(path); ext != ".proto" || strings.HasPrefix(path, "third_party") {
    77  			return nil
    78  		}
    79  		return generate(path, args)
    80  	})
    81  }
    82  
    83  // generate is used to execute the generate command for the specified proto file
    84  func generate(proto string, args []string) error {
    85  	input := []string{
    86  		"--proto_path=.",
    87  	}
    88  	if pathExists(protoPath) {
    89  		input = append(input, "--proto_path="+protoPath)
    90  	}
    91  	inputExt := []string{
    92  		"--proto_path=" + base.EnronMod(),
    93  		"--proto_path=" + filepath.Join(base.EnronMod(), "third_party"),
    94  		"--go_out=paths=source_relative:.",
    95  		"--go-grpc_out=paths=source_relative:.",
    96  		"--go-http_out=paths=source_relative:.",
    97  		"--go-errors_out=paths=source_relative:.",
    98  		"--openapi_out=paths=source_relative:.",
    99  	}
   100  	input = append(input, inputExt...)
   101  	protoBytes, err := os.ReadFile(proto)
   102  	if err == nil && len(protoBytes) > 0 {
   103  		if ok, _ := regexp.Match(`\n[^/]*(import)\s+"validate/validate.proto"`, protoBytes); ok {
   104  			input = append(input, "--validate_out=lang=go,paths=source_relative:.")
   105  		}
   106  	}
   107  	input = append(input, proto)
   108  	for _, a := range args {
   109  		if strings.HasPrefix(a, "-") {
   110  			input = append(input, a)
   111  		}
   112  	}
   113  	fd := exec.Command("protoc", input...)
   114  	fd.Stdout = os.Stdout
   115  	fd.Stderr = os.Stderr
   116  	fd.Dir = "."
   117  	if err := fd.Run(); err != nil {
   118  		return err
   119  	}
   120  	fmt.Printf("proto: %s\n", proto)
   121  	return nil
   122  }
   123  
   124  func pathExists(path string) bool {
   125  	_, err := os.Stat(path)
   126  	if err != nil {
   127  		return os.IsExist(err)
   128  	}
   129  	return true
   130  }