go-zero/tools/goctl/rpc/generator/gen.go

package generator

import (
	"os"
	"path/filepath"

	"github.com/zeromicro/go-zero/tools/goctl/rpc/parser"
	"github.com/zeromicro/go-zero/tools/goctl/util/console"
	"github.com/zeromicro/go-zero/tools/goctl/util/ctx"
	"github.com/zeromicro/go-zero/tools/goctl/util/pathx"
)

type ZRpcContext struct {
	// Src is the source file of the proto.
	Src string
	// ProtocCmd is the command to generate proto files.
	ProtocCmd string
	// ProtoGenGrpcDir is the directory to store the generated proto files.
	ProtoGenGrpcDir string
	// ProtoGenGoDir is the directory to store the generated go files.
	ProtoGenGoDir string
	// IsGooglePlugin is the flag to indicate whether the proto file is generated by google plugin.
	IsGooglePlugin bool
	// GoOutput is the output directory of the generated go files.
	GoOutput string
	// GrpcOutput is the output directory of the generated grpc files.
	GrpcOutput string
	// Output is the output directory of the generated files.
	Output string
	// Multiple is the flag to indicate whether the proto file is generated in multiple mode.
	Multiple bool
	// Whether to generate rpc client
	IsGenClient bool
	// Module is the custom module name for go.mod
	Module string
	// NameFromFilename uses proto filename instead of package name for service naming.
	// Default is false (uses package name, which supports multi-proto files).
	NameFromFilename bool
	// ProtoPaths are the directories to search for imported proto files,
	// equivalent to protoc -I flags. When empty the directory of Src is used.
	ProtoPaths []string
}

// Generate generates a rpc service, through the proto file,
// code storage directory, and proto import parameters to control
// the source file and target location of the rpc service that needs to be generated
func (g *Generator) Generate(zctx *ZRpcContext) error {
	abs, err := filepath.Abs(zctx.Output)
	if err != nil {
		return err
	}

	err = pathx.MkdirIfNotExist(abs)
	if err != nil {
		return err
	}

	err = g.Prepare()
	if err != nil {
		return err
	}

	var projectCtx *ctx.ProjectContext
	if len(zctx.Module) > 0 {
		projectCtx, err = ctx.PrepareWithModule(abs, zctx.Module)
	} else {
		projectCtx, err = ctx.Prepare(abs)
	}
	if err != nil {
		return err
	}

	p := parser.NewDefaultProtoParser()
	proto, err := p.Parse(zctx.Src, zctx.Multiple)
	if err != nil {
		return err
	}

	// Populate ImportedProtos so that generators can resolve dotted type references.
	proto.ImportedProtos, err = resolveImportedProtos(zctx)
	if err != nil {
		return err
	}

	dirCtx, err := mkdir(projectCtx, proto, g.cfg, zctx)
	if err != nil {
		return err
	}

	err = g.GenEtc(dirCtx, proto, g.cfg)
	if err != nil {
		return err
	}

	err = g.GenPb(dirCtx, zctx)
	if err != nil {
		return err
	}

	err = g.GenConfig(dirCtx, proto, g.cfg)
	if err != nil {
		return err
	}

	err = g.GenSvc(dirCtx, proto, g.cfg)
	if err != nil {
		return err
	}

	err = g.GenLogic(dirCtx, proto, g.cfg, zctx)
	if err != nil {
		return err
	}

	err = g.GenServer(dirCtx, proto, g.cfg, zctx)
	if err != nil {
		return err
	}

	err = g.GenMain(dirCtx, proto, g.cfg, zctx)
	if err != nil {
		return err
	}

	if zctx.IsGenClient {
		err = g.GenCall(dirCtx, proto, g.cfg, zctx)
	}

	console.NewColorConsole().MarkDone()

	return err
}

// resolveImportedProtos builds the full list of transitively imported proto
// files for zctx and returns their parsed metadata. This mirrors the proto
// path computation in genpb.go so both use the same search paths.
func resolveImportedProtos(zctx *ZRpcContext) ([]parser.ImportedProto, error) {
	pwd, err := os.Getwd()
	if err != nil {
		return nil, err
	}

	protoPaths := make([]string, 0, len(zctx.ProtoPaths)+1)
	srcDir := filepath.Dir(zctx.Src)
	if !filepath.IsAbs(srcDir) {
		srcDir = filepath.Join(pwd, srcDir)
	}
	protoPaths = append(protoPaths, srcDir)
	for _, p := range zctx.ProtoPaths {
		if !filepath.IsAbs(p) {
			p = filepath.Join(pwd, p)
		}
		protoPaths = append(protoPaths, p)
	}
	return parser.ParseImportedProtos(zctx.Src, protoPaths)
}