github.com/cloudwego/kitex@v0.9.0/tool/internal_pkg/generator/generator.go

// Copyright 2021 CloudWeGo Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package generator .
package generator

import (
	"fmt"
	"go/token"
	"path/filepath"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/cloudwego/kitex/tool/internal_pkg/log"
	"github.com/cloudwego/kitex/tool/internal_pkg/tpl"
	"github.com/cloudwego/kitex/tool/internal_pkg/util"
	"github.com/cloudwego/kitex/transport"
)

// Constants .
const (
	KitexGenPath = "kitex_gen"
	DefaultCodec = "thrift"

	BuildFileName       = "build.sh"
	BootstrapFileName   = "bootstrap.sh"
	ToolVersionFileName = "kitex_info.yaml"
	HandlerFileName     = "handler.go"
	MainFileName        = "main.go"
	ClientFileName      = "client.go"
	ServerFileName      = "server.go"
	InvokerFileName     = "invoker.go"
	ServiceFileName     = "*service.go"
	ExtensionFilename   = "extensions.yaml"

	DefaultThriftPluginTimeLimit = time.Minute
)

var (
	kitexImportPath = "github.com/cloudwego/kitex"

	globalMiddlewares  []Middleware
	globalDependencies = map[string]string{
		"kitex":   kitexImportPath,
		"client":  ImportPathTo("client"),
		"server":  ImportPathTo("server"),
		"callopt": ImportPathTo("client/callopt"),
		"frugal":  "github.com/cloudwego/frugal",
	}
)

// SetKitexImportPath sets the import path of kitex.
// Must be called before generating code.
func SetKitexImportPath(path string) {
	for k, v := range globalDependencies {
		globalDependencies[k] = strings.ReplaceAll(v, kitexImportPath, path)
	}
	kitexImportPath = path
}

// ImportPathTo returns an import path to the specified package under kitex.
func ImportPathTo(pkg string) string {
	return util.JoinPath(kitexImportPath, pkg)
}

// AddGlobalMiddleware adds middleware for all generators
func AddGlobalMiddleware(mw Middleware) {
	globalMiddlewares = append(globalMiddlewares, mw)
}

// AddGlobalDependency adds dependency for all generators
func AddGlobalDependency(ref, path string) bool {
	if _, ok := globalDependencies[ref]; !ok {
		globalDependencies[ref] = path
		return true
	}
	return false
}

// Generator generates the codes of main package and scripts for building a server based on kitex.
type Generator interface {
	GenerateService(pkg *PackageInfo) ([]*File, error)
	GenerateMainPackage(pkg *PackageInfo) ([]*File, error)
	GenerateCustomPackage(pkg *PackageInfo) ([]*File, error)
}

// Config .
type Config struct {
	Verbose               bool
	GenerateMain          bool // whether stuff in the main package should be generated
	GenerateInvoker       bool // generate main.go with invoker when main package generate
	Version               string
	NoFastAPI             bool
	ModuleName            string
	ServiceName           string
	Use                   string
	IDLType               string
	Includes              util.StringSlice
	ThriftOptions         util.StringSlice
	ProtobufOptions       util.StringSlice
	Hessian2Options       util.StringSlice
	IDL                   string // the IDL file passed on the command line
	OutputPath            string // the output path for main pkg and kitex_gen
	PackagePrefix         string
	CombineService        bool // combine services to one service
	CopyIDL               bool
	ThriftPlugins         util.StringSlice
	ProtobufPlugins       util.StringSlice
	Features              []feature
	FrugalPretouch        bool
	ThriftPluginTimeLimit time.Duration
	CompilerPath          string // specify the path of thriftgo or protoc

	ExtensionFile string
	tmplExt       *TemplateExtension

	Record    bool
	RecordCmd []string

	TemplateDir string

	GenPath string

	DeepCopyAPI           bool
	Protocol              string
	HandlerReturnKeepResp bool
}

// Pack packs the Config into a slice of "key=val" strings.
func (c *Config) Pack() (res []string) {
	t := reflect.TypeOf(c).Elem()
	v := reflect.ValueOf(c).Elem()
	for i := 0; i < t.NumField(); i++ {
		f := t.Field(i)
		x := v.Field(i)
		n := f.Name

		// skip the plugin arguments to avoid the 'strings in strings' trouble
		if f.Name == "ThriftPlugins" || !token.IsExported(f.Name) {
			continue
		}

		if str, ok := x.Interface().(interface{ String() string }); ok {
			res = append(res, n+"="+str.String())
			continue
		}

		switch x.Kind() {
		case reflect.Bool:
			res = append(res, n+"="+fmt.Sprint(x.Bool()))
		case reflect.String:
			res = append(res, n+"="+x.String())
		case reflect.Slice:
			var ss []string
			if x.Type().Elem().Kind() == reflect.Int {
				for i := 0; i < x.Len(); i++ {
					ss = append(ss, strconv.Itoa(int(x.Index(i).Int())))
				}
			} else {
				for i := 0; i < x.Len(); i++ {
					ss = append(ss, x.Index(i).String())
				}
			}
			res = append(res, n+"="+strings.Join(ss, ";"))
		default:
			panic(fmt.Errorf("unsupported field type: %+v", f))
		}
	}
	return res
}

// Unpack restores the Config from a slice of "key=val" strings.
func (c *Config) Unpack(args []string) error {
	t := reflect.TypeOf(c).Elem()
	v := reflect.ValueOf(c).Elem()
	for _, a := range args {
		parts := strings.SplitN(a, "=", 2)
		if len(parts) != 2 {
			return fmt.Errorf("invalid argument: '%s'", a)
		}
		name, value := parts[0], parts[1]
		f, ok := t.FieldByName(name)
		if ok && value != "" {
			x := v.FieldByName(name)
			if _, ok := x.Interface().(time.Duration); ok {
				if d, err := time.ParseDuration(value); err != nil {
					return fmt.Errorf("invalid time duration '%s' for %s", value, name)
				} else {
					x.SetInt(int64(d))
				}
				continue
			}
			switch x.Kind() {
			case reflect.Bool:
				x.SetBool(value == "true")
			case reflect.String:
				x.SetString(value)
			case reflect.Slice:
				ss := strings.Split(value, ";")
				if x.Type().Elem().Kind() == reflect.Int {
					n := reflect.MakeSlice(x.Type(), len(ss), len(ss))
					for i, s := range ss {
						val, err := strconv.ParseInt(s, 10, 64)
						if err != nil {
							return err
						}
						n.Index(i).SetInt(val)
					}
					x.Set(n)
				} else {
					for _, s := range ss {
						val := reflect.Append(x, reflect.ValueOf(s))
						x.Set(val)
					}
				}
			default:
				return fmt.Errorf("unsupported field type: %+v", f)
			}
		}
	}
	log.Verbose = c.Verbose
	return c.ApplyExtension()
}

// AddFeature add registered feature to config
func (c *Config) AddFeature(key string) bool {
	if f, ok := getFeature(key); ok {
		c.Features = append(c.Features, f)
		return true
	}
	return false
}

// ApplyExtension applies template extension.
func (c *Config) ApplyExtension() error {
	templateExtExist := false
	path := util.JoinPath(c.TemplateDir, ExtensionFilename)
	if c.TemplateDir != "" && util.Exists(path) {
		templateExtExist = true
	}

	if c.ExtensionFile == "" && !templateExtExist {
		return nil
	}

	ext := new(TemplateExtension)
	if c.ExtensionFile != "" {
		if err := ext.FromYAMLFile(c.ExtensionFile); err != nil {
			return fmt.Errorf("read template extension %q failed: %s", c.ExtensionFile, err.Error())
		}
	}

	if templateExtExist {
		yamlExt := new(TemplateExtension)
		if err := yamlExt.FromYAMLFile(path); err != nil {
			return fmt.Errorf("read template extension %q failed: %s", path, err.Error())
		}
		ext.Merge(yamlExt)
	}

	for _, fn := range ext.FeatureNames {
		RegisterFeature(fn)
	}
	for _, fn := range ext.EnableFeatures {
		c.AddFeature(fn)
	}
	for path, alias := range ext.Dependencies {
		AddGlobalDependency(alias, path)
	}

	c.tmplExt = ext
	return nil
}

// NewGenerator .
func NewGenerator(config *Config, middlewares []Middleware) Generator {
	mws := append(globalMiddlewares, middlewares...)
	g := &generator{Config: config, middlewares: mws}
	if g.IDLType == "" {
		g.IDLType = DefaultCodec
	}
	return g
}

// Middleware used generator
type Middleware func(HandleFunc) HandleFunc

// HandleFunc used generator
type HandleFunc func(*Task, *PackageInfo) (*File, error)

type generator struct {
	*Config
	middlewares []Middleware
}

func (g *generator) chainMWs(handle HandleFunc) HandleFunc {
	for i := len(g.middlewares) - 1; i > -1; i-- {
		handle = g.middlewares[i](handle)
	}
	return handle
}

func (g *generator) GenerateMainPackage(pkg *PackageInfo) (fs []*File, err error) {
	g.updatePackageInfo(pkg)

	tasks := []*Task{
		{
			Name: BuildFileName,
			Path: util.JoinPath(g.OutputPath, BuildFileName),
			Text: tpl.BuildTpl,
		},
		{
			Name: BootstrapFileName,
			Path: util.JoinPath(g.OutputPath, "script", BootstrapFileName),
			Text: tpl.BootstrapTpl,
		},
		{
			Name: ToolVersionFileName,
			Path: util.JoinPath(g.OutputPath, ToolVersionFileName),
			Text: tpl.ToolVersionTpl,
		},
	}
	if !g.Config.GenerateInvoker {
		tasks = append(tasks, &Task{
			Name: MainFileName,
			Path: util.JoinPath(g.OutputPath, MainFileName),
			Text: tpl.MainTpl,
		})
	}
	for _, t := range tasks {
		if util.Exists(t.Path) {
			log.Info(t.Path, "exists. Skipped.")
			continue
		}
		g.setImports(t.Name, pkg)
		handle := func(task *Task, pkg *PackageInfo) (*File, error) {
			return task.Render(pkg)
		}
		f, err := g.chainMWs(handle)(t, pkg)
		if err != nil {
			return nil, err
		}
		fs = append(fs, f)
	}

	handlerFilePath := filepath.Join(g.OutputPath, HandlerFileName)
	if util.Exists(handlerFilePath) {
		comp := newCompleter(
			pkg.ServiceInfo.AllMethods(),
			handlerFilePath,
			pkg.ServiceInfo.ServiceName)
		f, err := comp.CompleteMethods()
		if err != nil {
			if err == errNoNewMethod {
				return fs, nil
			}
			return nil, err
		}
		fs = append(fs, f)
	} else {
		task := Task{
			Name: HandlerFileName,
			Path: handlerFilePath,
			Text: tpl.HandlerTpl + "\n" + tpl.HandlerMethodsTpl,
		}
		g.setImports(task.Name, pkg)
		handle := func(task *Task, pkg *PackageInfo) (*File, error) {
			return task.Render(pkg)
		}
		f, err := g.chainMWs(handle)(&task, pkg)
		if err != nil {
			return nil, err
		}
		fs = append(fs, f)
	}
	return
}

func (g *generator) GenerateService(pkg *PackageInfo) ([]*File, error) {
	g.updatePackageInfo(pkg)
	output := util.JoinPath(g.OutputPath, util.CombineOutputPath(g.GenPath, pkg.Namespace))
	svcPkg := strings.ToLower(pkg.ServiceName)
	output = util.JoinPath(output, svcPkg)
	ext := g.tmplExt
	if ext == nil {
		ext = new(TemplateExtension)
	}

	tasks := []*Task{
		{
			Name: ClientFileName,
			Path: util.JoinPath(output, ClientFileName),
			Text: tpl.ClientTpl,
			Ext:  ext.ExtendClient,
		},
		{
			Name: ServerFileName,
			Path: util.JoinPath(output, ServerFileName),
			Text: tpl.ServerTpl,
			Ext:  ext.ExtendServer,
		},
		{
			Name: InvokerFileName,
			Path: util.JoinPath(output, InvokerFileName),
			Text: tpl.InvokerTpl,
			Ext:  ext.ExtendInvoker,
		},
		{
			Name: ServiceFileName,
			Path: util.JoinPath(output, svcPkg+".go"),
			Text: tpl.ServiceTpl,
		},
	}

	var fs []*File
	for _, t := range tasks {
		if err := t.Build(); err != nil {
			err = fmt.Errorf("build %s failed: %w", t.Name, err)
			return nil, err
		}
		g.setImports(t.Name, pkg)
		if t.Ext != nil {
			for _, path := range t.Ext.ImportPaths {
				if alias, exist := ext.Dependencies[path]; exist {
					pkg.AddImports(alias)
				}
			}
		}
		handle := func(task *Task, pkg *PackageInfo) (*File, error) {
			return task.Render(pkg)
		}
		f, err := g.chainMWs(handle)(t, pkg)
		if err != nil {
			err = fmt.Errorf("render %s failed: %w", t.Name, err)
			return nil, err
		}
		fs = append(fs, f)
	}
	return fs, nil
}

func (g *generator) updatePackageInfo(pkg *PackageInfo) {
	pkg.NoFastAPI = g.NoFastAPI
	pkg.Codec = g.IDLType
	pkg.Version = g.Version
	pkg.RealServiceName = g.ServiceName
	pkg.Features = g.Features
	pkg.ExternalKitexGen = g.Use
	pkg.FrugalPretouch = g.FrugalPretouch
	pkg.Module = g.ModuleName
	if strings.EqualFold(g.Protocol, transport.HESSIAN2.String()) {
		pkg.Protocol = transport.HESSIAN2
	}
	if pkg.Dependencies == nil {
		pkg.Dependencies = make(map[string]string)
	}

	for ref, path := range globalDependencies {
		if _, ok := pkg.Dependencies[ref]; !ok {
			pkg.Dependencies[ref] = path
		}
	}
}

func (g *generator) setImports(name string, pkg *PackageInfo) {
	pkg.Imports = make(map[string]map[string]bool)
	switch name {
	case ClientFileName:
		pkg.AddImports("client")
		if pkg.HasStreaming {
			pkg.AddImport("streaming", "github.com/cloudwego/kitex/pkg/streaming")
			pkg.AddImport("transport", "github.com/cloudwego/kitex/transport")
		}
		if len(pkg.AllMethods()) > 0 {
			if needCallOpt(pkg) {
				pkg.AddImports("callopt")
			}
			pkg.AddImports("context")
		}
		fallthrough
	case HandlerFileName:
		for _, m := range pkg.ServiceInfo.AllMethods() {
			if !m.ServerStreaming && !m.ClientStreaming {
				pkg.AddImports("context")
			}
			for _, a := range m.Args {
				for _, dep := range a.Deps {
					pkg.AddImport(dep.PkgRefName, dep.ImportPath)
				}
			}
			if !m.Void && m.Resp != nil {
				for _, dep := range m.Resp.Deps {
					pkg.AddImport(dep.PkgRefName, dep.ImportPath)
				}
			}
		}
	case ServerFileName, InvokerFileName:
		if len(pkg.CombineServices) == 0 {
			pkg.AddImport(pkg.ServiceInfo.PkgRefName, pkg.ServiceInfo.ImportPath)
		}
		pkg.AddImports("server")
	case ServiceFileName:
		pkg.AddImports("errors")
		pkg.AddImports("client")
		pkg.AddImport("kitex", "github.com/cloudwego/kitex/pkg/serviceinfo")
		pkg.AddImport(pkg.ServiceInfo.PkgRefName, pkg.ServiceInfo.ImportPath)
		if len(pkg.AllMethods()) > 0 {
			pkg.AddImports("context")
		}
		for _, m := range pkg.ServiceInfo.AllMethods() {
			if m.ClientStreaming || m.ServerStreaming {
				pkg.AddImports("fmt")
			}
			if m.GenArgResultStruct {
				pkg.AddImports("proto")
			} else {
				// for method Arg and Result
				pkg.AddImport(m.PkgRefName, m.ImportPath)
			}
			for _, a := range m.Args {
				for _, dep := range a.Deps {
					pkg.AddImport(dep.PkgRefName, dep.ImportPath)
				}
			}
			if m.Streaming.IsStreaming || pkg.Codec == "protobuf" {
				// protobuf handler support both PingPong and Unary (streaming) requests
				pkg.AddImport("streaming", "github.com/cloudwego/kitex/pkg/streaming")
			}
			if !m.Void && m.Resp != nil {
				for _, dep := range m.Resp.Deps {
					pkg.AddImport(dep.PkgRefName, dep.ImportPath)
				}
			}
			for _, e := range m.Exceptions {
				for _, dep := range e.Deps {
					pkg.AddImport(dep.PkgRefName, dep.ImportPath)
				}
			}
		}
		if pkg.FrugalPretouch {
			pkg.AddImports("sync")
			if len(pkg.AllMethods()) > 0 {
				pkg.AddImports("frugal")
				pkg.AddImports("reflect")
			}
		}
	case MainFileName:
		pkg.AddImport("log", "log")
		pkg.AddImport(pkg.PkgRefName, util.JoinPath(pkg.ImportPath, strings.ToLower(pkg.ServiceName)))
	}
}

func needCallOpt(pkg *PackageInfo) bool {
	// callopt is referenced only by non-streaming methods
	needCallOpt := false
	switch pkg.Codec {
	case "thrift":
		for _, m := range pkg.ServiceInfo.AllMethods() {
			if !m.Streaming.IsStreaming {
				needCallOpt = true
				break
			}
		}
	case "protobuf":
		needCallOpt = true
	}
	return needCallOpt
}