genmain.go 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. package generator
  2. import (
  3. "fmt"
  4. "path/filepath"
  5. "strings"
  6. conf "github.com/tal-tech/go-zero/tools/goctl/config"
  7. "github.com/tal-tech/go-zero/tools/goctl/rpc/parser"
  8. "github.com/tal-tech/go-zero/tools/goctl/util"
  9. "github.com/tal-tech/go-zero/tools/goctl/util/format"
  10. "github.com/tal-tech/go-zero/tools/goctl/util/stringx"
  11. )
  12. const mainTemplate = `package main
  13. import (
  14. "flag"
  15. "fmt"
  16. {{.imports}}
  17. "github.com/tal-tech/go-zero/core/conf"
  18. "github.com/tal-tech/go-zero/zrpc"
  19. "google.golang.org/grpc"
  20. )
  21. var configFile = flag.String("f", "etc/{{.serviceName}}.yaml", "the config file")
  22. func main() {
  23. flag.Parse()
  24. var c config.Config
  25. conf.MustLoad(*configFile, &c)
  26. ctx := svc.NewServiceContext(c)
  27. srv := server.New{{.serviceNew}}Server(ctx)
  28. s := zrpc.MustNewServer(c.RpcServerConf, func(grpcServer *grpc.Server) {
  29. {{.pkg}}.Register{{.service}}Server(grpcServer, srv)
  30. })
  31. defer s.Stop()
  32. fmt.Printf("Starting rpc server at %s...\n", c.ListenOn)
  33. s.Start()
  34. }
  35. `
  36. // GenMain generates the main file of the rpc service, which is an rpc service program call entry
  37. func (g *DefaultGenerator) GenMain(ctx DirContext, proto parser.Proto, cfg *conf.Config) error {
  38. mainFilename, err := format.FileNamingFormat(cfg.NamingFormat, ctx.GetServiceName().Source())
  39. if err != nil {
  40. return err
  41. }
  42. fileName := filepath.Join(ctx.GetMain().Filename, fmt.Sprintf("%v.go", mainFilename))
  43. imports := make([]string, 0)
  44. pbImport := fmt.Sprintf(`"%v"`, ctx.GetPb().Package)
  45. svcImport := fmt.Sprintf(`"%v"`, ctx.GetSvc().Package)
  46. remoteImport := fmt.Sprintf(`"%v"`, ctx.GetServer().Package)
  47. configImport := fmt.Sprintf(`"%v"`, ctx.GetConfig().Package)
  48. imports = append(imports, configImport, pbImport, remoteImport, svcImport)
  49. text, err := util.LoadTemplate(category, mainTemplateFile, mainTemplate)
  50. if err != nil {
  51. return err
  52. }
  53. return util.With("main").GoFmt(true).Parse(text).SaveTo(map[string]interface{}{
  54. "serviceName": strings.ToLower(ctx.GetServiceName().ToCamel()),
  55. "imports": strings.Join(imports, util.NL),
  56. "pkg": proto.PbPackage,
  57. "serviceNew": stringx.From(proto.Service.Name).ToCamel(),
  58. "service": parser.CamelCase(proto.Service.Name),
  59. }, fileName, false)
  60. }