diff --git a/cmd/kratos/internal/proto/client/client.go b/cmd/kratos/internal/proto/client/client.go index 61b6651ff..6ca4433bf 100644 --- a/cmd/kratos/internal/proto/client/client.go +++ b/cmd/kratos/internal/proto/client/client.go @@ -17,14 +17,22 @@ import ( var ( // CmdClient represents the source command. CmdClient = &cobra.Command{ - Use: "client", - Short: "Generate the proto client code", - Long: "Generate the proto client code. Example: kratos proto client helloworld.proto", - DisableFlagParsing: true, - Run: run, + Use: "client", + Short: "Generate the proto client code", + Long: "Generate the proto client code. Example: kratos proto client helloworld.proto", + Run: run, } ) +var protoPath string + +func init() { + if protoPath = os.Getenv("KRATOS_PROTO_PATH"); protoPath == "" { + protoPath = "./third_party" + } + CmdClient.Flags().StringVarP(&protoPath, "proto_path", "p", protoPath, "proto path") +} + func run(cmd *cobra.Command, args []string) { if len(args) == 0 { fmt.Println("Please enter the proto file or directory") @@ -79,7 +87,11 @@ func walk(dir string, args []string) error { func generate(proto string, args []string) error { input := []string{ "--proto_path=.", - "--proto_path=./third_party", + } + if pathExists(protoPath) { + input = append(input, "--proto_path="+protoPath) + } + inputExt := []string{ "--proto_path=" + base.KratosMod(), "--proto_path=" + filepath.Join(base.KratosMod(), "third_party"), "--go_out=paths=source_relative:.", @@ -87,6 +99,7 @@ func generate(proto string, args []string) error { "--go-http_out=paths=source_relative:.", "--go-errors_out=paths=source_relative:.", } + input = append(input, inputExt...) protoBytes, err := ioutil.ReadFile(proto) if err == nil && len(protoBytes) > 0 { if ok, _ := regexp.Match(`\n[^/]*(import)\s+"validate/validate.proto"`, protoBytes); ok { @@ -109,3 +122,11 @@ func generate(proto string, args []string) error { fmt.Printf("proto: %s\n", proto) return nil } + +func pathExists(path string) bool { + _, err := os.Stat(path) + if err != nil { + return os.IsExist(err) + } + return true +}