From a636fd52a42e215f265705cd45e17ea4ee0f3f86 Mon Sep 17 00:00:00 2001
From: opensite <lucky9303@163.com>
Date: Wed, 28 Jul 2021 19:20:29 +0800
Subject: [PATCH] feat(cmd): support the third party specify by self and fix
 the warning when third party is not in current directory (#1266)

* feat(cmd): support the third party specify by self and fix the warning when third party is not in current directory

* add env support and fix the directory of specify as full directory

* fix the env and flag variable name and rename the function name of the directory is exist check
---
 cmd/kratos/internal/proto/client/client.go | 33 ++++++++++++++++++----
 1 file changed, 27 insertions(+), 6 deletions(-)

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
+}