refactor: refactor kitex tool

Author: HeyJavaBean

tool/cmd/kitex/args/args.go

@@ -25,6 +25,10 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/cloudwego/kitex/tool/cmd/kitex/versions"
+
+	"github.com/cloudwego/kitex/tool/cmd/kitex/code"
+
 	"github.com/cloudwego/kitex/tool/internal_pkg/generator"
 	"github.com/cloudwego/kitex/tool/internal_pkg/log"
 	"github.com/cloudwego/kitex/tool/internal_pkg/pluginmode/protoc"
@@ -50,7 +54,9 @@ type ExtraFlag struct {
 // Arguments .
 type Arguments struct {
 	generator.Config
-	extends []*ExtraFlag
+	extends           []*ExtraFlag
+	queryVersion      bool
+	dependencyChecker *versions.DependencyChecker
 }
 
 const (
@@ -77,6 +83,8 @@ func (a *Arguments) AddExtraFlag(e *ExtraFlag) {
 
 func (a *Arguments) buildFlags(version string) *flag.FlagSet {
 	f := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
+	f.BoolVar(&a.queryVersion, "version", false,
+		"Show the version of kitex")
 	f.BoolVar(&a.NoFastAPI, "no-fast-api", false,
 		"Generate codes without injecting fast method.")
 	f.StringVar(&a.ModuleName, "module", "",
@@ -175,6 +183,15 @@ func (a *Arguments) ParseArgs(version, curpath string, kitexArgs []string) (err
 	if err = f.Parse(kitexArgs); err != nil {
 		return err
 	}
+	if a.queryVersion {
+		println(a.Version)
+		return code.ErrExitZeroInterrupted
+	}
+	if !a.NoDependencyCheck {
+		if ok, _ := versions.CheckVersion(); !ok {
+			return code.ErrVersionCheckFailed
+		}
+	}
 	if a.StreamX {
 		a.ThriftOptions = append(a.ThriftOptions, "streamx")
 	}

tool/cmd/kitex/code/code.go

@@ -0,0 +1,51 @@
+package code
+
+import (
+	"errors"
+	"flag"
+	"fmt"
+)
+
+// ToolExitCode defines the enumeration type for tool exit codes.
+type ToolExitCode int
+
+const (
+	// ToolSuccess indicates that the tool has run successfully and returned.
+	ToolSuccess ToolExitCode = iota
+
+	ToolExecuteFailed
+
+	ToolArgsError
+	// ToolVersionCheckFailed indicates that the tool's version check has failed.
+	ToolVersionCheckFailed
+	ToolInterrupted
+	ToolInitFailed
+)
+
+var (
+	ErrExitZeroInterrupted = fmt.Errorf("os.Exit(%d)", ToolInterrupted)
+	ErrVersionCheckFailed  = fmt.Errorf("os.Exit(%d)", ToolVersionCheckFailed)
+)
+
+func WrapExitCode(ret int) (error, ToolExitCode) {
+	switch ret {
+	case 0:
+		return nil, ToolSuccess
+	case 1:
+		return fmt.Errorf("tool run failed"), ToolExecuteFailed
+	default:
+		return fmt.Errorf("unexpected return value: %d", ret), ToolExecuteFailed
+	}
+}
+
+func Good(ret ToolExitCode) bool {
+	return ret == ToolSuccess || ret == ToolInterrupted
+}
+
+func IsInterrupted(err error) bool {
+	return errors.Is(err, flag.ErrHelp) || errors.Is(err, ErrExitZeroInterrupted)
+}
+
+func IsVersionCheckFailed(err error) bool {
+	return errors.Is(err, ErrVersionCheckFailed)
+}

tool/cmd/kitex/generator/generator.go

@@ -0,0 +1,84 @@
+package generator
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+
+	"github.com/cloudwego/kitex/tool/internal_pkg/thriftgo"
+
+	"github.com/cloudwego/kitex/tool/cmd/kitex/code"
+
+	kargs "github.com/cloudwego/kitex/tool/cmd/kitex/args"
+	"github.com/cloudwego/kitex/tool/internal_pkg/pluginmode/protoc"
+	thriftgo_plugin "github.com/cloudwego/kitex/tool/internal_pkg/pluginmode/thriftgo"
+	"github.com/cloudwego/kitex/tool/internal_pkg/prutal"
+	"github.com/cloudwego/kitex/tool/internal_pkg/util/env"
+
+	"github.com/cloudwego/kitex"
+	t "github.com/cloudwego/kitex/tool/internal_pkg/thriftgo"
+	"github.com/cloudwego/thriftgo/plugin"
+)
+
+// RunKitexThriftgoGen run kitex tool to generate thrift as sdk
+func RunKitexThriftgoGen(wd string, plugins []plugin.SDKPlugin, kitexOsArgs ...string) (error, code.ToolExitCode) {
+	var args kargs.Arguments
+	err := args.ParseArgs(kitex.Version, wd, kitexOsArgs)
+	if err != nil {
+		return err, code.ToolArgsError
+	}
+	return t.RunKitexThriftgoGenFromArgs(wd, plugins, &args)
+}
+
+func RunKitexToolCmd(curpath string, osArgs []string, args *kargs.Arguments, toolVersion string) (err error, retCode code.ToolExitCode) {
+	// try run as thriftgo/protoc plugin process
+	mode := os.Getenv(kargs.EnvPluginMode)
+	if len(osArgs) <= 1 && mode != "" {
+		switch mode {
+		case thriftgo_plugin.PluginName:
+			return code.WrapExitCode(thriftgo_plugin.Run())
+		case protoc.PluginName:
+			return code.WrapExitCode(protoc.Run())
+		}
+		return fmt.Errorf("unknown plugin mode %s", mode), code.ToolArgsError
+	}
+
+	// parse arguments
+	err = args.ParseArgs(toolVersion, curpath, osArgs[1:])
+	if err != nil {
+		if code.IsInterrupted(err) {
+			return nil, code.ToolInterrupted
+		}
+		if code.IsVersionCheckFailed(err) {
+			return err, code.ToolVersionCheckFailed
+		}
+		return err, code.ToolArgsError
+	}
+
+	if args.IsProtobuf() && !env.UseProtoc() {
+		if err = prutal.NewPrutalGen(args.Config).Process(); err != nil {
+			return err, code.ToolExecuteFailed
+		}
+	}
+
+	if args.IsThrift() && !args.LocalThriftgo {
+		return thriftgo.RunKitexThriftgoGenFromArgs(curpath, nil, args)
+	}
+
+	out := new(bytes.Buffer)
+	cmd, err := args.BuildCmd(out)
+	if err != nil {
+		return err, code.ToolArgsError
+	}
+
+	err = kargs.ValidateCMD(cmd.Path, args.IDLType)
+	if err != nil {
+		return err, code.ToolArgsError
+	}
+
+	if err = cmd.Run(); err != nil {
+		return err, code.ToolExecuteFailed
+	}
+
+	return nil, code.ToolSuccess
+}

tool/cmd/kitex/main.go

@@ -15,128 +15,41 @@
 package main
 
 import (
-	"bytes"
-	"errors"
-	"flag"
 	"os"
 	"path/filepath"
-	"strings"
 
-	"github.com/cloudwego/kitex/tool/cmd/kitex/utils"
+	"github.com/cloudwego/kitex/tool/cmd/kitex/code"
 
-	"github.com/cloudwego/kitex/tool/cmd/kitex/sdk"
+	"github.com/cloudwego/kitex/tool/cmd/kitex/utils"
 
 	"github.com/cloudwego/kitex"
 	kargs "github.com/cloudwego/kitex/tool/cmd/kitex/args"
-	"github.com/cloudwego/kitex/tool/cmd/kitex/versions"
+	"github.com/cloudwego/kitex/tool/cmd/kitex/generator"
 	"github.com/cloudwego/kitex/tool/internal_pkg/log"
-	"github.com/cloudwego/kitex/tool/internal_pkg/pluginmode/protoc"
-	"github.com/cloudwego/kitex/tool/internal_pkg/pluginmode/thriftgo"
-	"github.com/cloudwego/kitex/tool/internal_pkg/prutal"
-	"github.com/cloudwego/kitex/tool/internal_pkg/util/env"
 )
 
-var args kargs.Arguments
-
-func init() {
-	var queryVersion bool
-	args.AddExtraFlag(&kargs.ExtraFlag{
-		Apply: func(f *flag.FlagSet) {
-			f.BoolVar(&queryVersion, "version", false,
-				"Show the version of kitex")
-		},
-		Check: func(a *kargs.Arguments) error {
-			if queryVersion {
-				println(a.Version)
-				os.Exit(0)
-			}
-			return nil
-		},
-	})
-	if err := versions.RegisterMinDepVersion(
-		&versions.MinDepVersion{
-			RefPath: "github.com/cloudwego/kitex",
-			Version: "v0.11.0",
-		},
-	); err != nil {
-		log.Error(err)
-		os.Exit(versions.CompatibilityCheckExitCode)
-	}
-}
-
 func main() {
-	mode := os.Getenv(kargs.EnvPluginMode)
-	if len(os.Args) <= 1 && mode != "" {
-		// run as a plugin
-		switch mode {
-		case thriftgo.PluginName:
-			os.Exit(thriftgo.Run())
-		case protoc.PluginName:
-			os.Exit(protoc.Run())
-		}
-		panic(mode)
-	}
+	var args kargs.Arguments
 
 	curpath, err := filepath.Abs(".")
 	if err != nil {
-		log.Errorf("Get current path failed: %s", err)
-		os.Exit(1)
-	}
-	// run as kitex
-	err = args.ParseArgs(kitex.Version, curpath, os.Args[1:])
-	if err != nil {
-		if errors.Is(err, flag.ErrHelp) {
-			os.Exit(0)
-		}
-		log.Error(err)
-		os.Exit(2)
-	}
-	if !args.NoDependencyCheck {
-		// check dependency compatibility between kitex cmd tool and dependency in go.mod
-		if err := versions.DefaultCheckDependencyAndProcess(); err != nil {
-			os.Exit(versions.CompatibilityCheckExitCode)
-		}
-	}
-	if args.IsProtobuf() && !env.UseProtoc() {
-		g := prutal.NewPrutalGen(args.Config)
-		if err := g.Process(); err != nil {
-			log.Errorf("%s", err)
-			os.Exit(1)
-		}
-		return
-	}
-
-	out := new(bytes.Buffer)
-	cmd, err := args.BuildCmd(out)
-	if err != nil {
-		log.Warn(err)
-		os.Exit(1)
+		log.Errorf("get current path failed: %s", err.Error())
+		os.Exit(int(code.ToolExecuteFailed))
 	}
 
-	if args.IsThrift() && !args.LocalThriftgo {
-		if err = sdk.InvokeThriftgoBySDK(curpath, cmd); err != nil {
-			// todo: optimize -use and remove error returned from thriftgo
-			out.WriteString(err.Error())
-		}
-	} else {
-		err = kargs.ValidateCMD(cmd.Path, args.IDLType)
-		if err != nil {
-			log.Warn(err)
-			os.Exit(1)
-		}
-		err = cmd.Run()
+	// run kitex tool:
+	// 1. check and execute process plugin mode callback
+	// 2. parse user input to arguments
+	// 3. check dependency version ( details: versions/dependencies_v2.go)
+	// 4. execute code generation by choosing specific tool (thriftgo、prutal、protoc?)
+	err, retCode := generator.RunKitexToolCmd(curpath, os.Args, &args, kitex.Version)
+	if err != nil || !code.Good(retCode) {
+		log.Errorf(err.Error())
+		os.Exit(int(retCode))
 	}
-
-	if err != nil {
-		if args.Use != "" {
-			out := strings.TrimSpace(out.String())
-			if strings.HasSuffix(out, thriftgo.TheUseOptionMessage) {
-				goto NormalExit
-			}
-		}
-		log.Warn(err)
-		os.Exit(1)
+	// only execute after tool finishes code generation (not -help、-version)
+	if retCode != code.ToolInterrupted {
+		utils.OnKitexToolNormalExit(args)
 	}
-NormalExit:
-	utils.OnKitexToolNormalExit(args)
+	os.Exit(int(code.ToolSuccess))
 }