feat(hz): enhance the ability for customize template (#767)

cmd/hz/generator/custom_files.go

@@ -216,7 +216,14 @@ func appendUpdateFile(tplInfo *Template, renderInfo interface{}, fileContent []b
 	if err != nil {
 		return []byte(""), fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
 	}
-	_, err = buf.WriteString("\n" + appendContent + "\n")
+	// "\r\n" && "\n" has the same suffix
+	if !bytes.HasSuffix(buf.Bytes(), []byte("\n")) {
+		_, err = buf.WriteString("\n")
+		if err != nil {
+			return []byte(""), fmt.Errorf("write file(%s) line break failed, err: %v", tplInfo.Path, err)
+		}
+	}
+	_, err = buf.WriteString(appendContent)
 	if err != nil {
 		return []byte(""), fmt.Errorf("append file(%s) failed, err: %v", tplInfo.Path, err)
 	}
@@ -348,17 +355,34 @@ func (pkgGen *HttpPackageGenerator) genLoopService(tplInfo *Template, filePathRe
 						return err
 					}
 				}
-				buf := bytes.NewBuffer(nil)
-				_, err = buf.Write(fileContent)
-				if err != nil {
-					return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
-				}
-				_, err = buf.Write(appendContent)
-				if err != nil {
-					return fmt.Errorf("append file(%s) failed, err: %v", tplInfo.Path, err)
+				if len(tplInfo.UpdateBehavior.AppendLocation) == 0 { // default, append to end of file
+					buf := bytes.NewBuffer(nil)
+					_, err = buf.Write(fileContent)
+					if err != nil {
+						return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					_, err = buf.Write(appendContent)
+					if err != nil {
+						return fmt.Errorf("append file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'method'", filePath)
+					pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
+				} else { // 'append location', append new content after 'append location'
+					part := bytes.Split(fileContent, []byte(tplInfo.UpdateBehavior.AppendLocation))
+					if len(part) == 0 {
+						return fmt.Errorf("can not find append location '%s' for file '%s'\n", tplInfo.UpdateBehavior.AppendLocation, filePath)
+					}
+					if len(part) != 2 {
+						return fmt.Errorf("do not support multiple append location '%s' for file '%s'\n", tplInfo.UpdateBehavior.AppendLocation, filePath)
+					}
+					buf := bytes.NewBuffer(nil)
+					err = writeBytes(buf, part[0], []byte(tplInfo.UpdateBehavior.AppendLocation), appendContent, part[1])
+					if err != nil {
+						return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'method'", filePath)
+					pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
 				}
-				logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'method'", filePath)
-				pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
 			} else {
 				logs.Warnf("Loop 'service' field for '%s' only append content by appendKey for 'method', so cannot append content", filePath)
 			}
@@ -505,17 +529,34 @@ func (pkgGen *HttpPackageGenerator) genSingleCustomizedFile(tplInfo *Template, f
 						}
 					}
 				}
-				buf := bytes.NewBuffer(nil)
-				_, err = buf.Write(fileContent)
-				if err != nil {
-					return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
-				}
-				_, err = buf.Write(appendContent)
-				if err != nil {
-					return fmt.Errorf("append file(%s) failed, err: %v", tplInfo.Path, err)
+				if len(tplInfo.UpdateBehavior.AppendLocation) == 0 { // default, append to end of file
+					buf := bytes.NewBuffer(nil)
+					_, err = buf.Write(fileContent)
+					if err != nil {
+						return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					_, err = buf.Write(appendContent)
+					if err != nil {
+						return fmt.Errorf("append file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'method'", filePath)
+					pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
+				} else { // 'append location', append new content after 'append location'
+					part := bytes.Split(fileContent, []byte(tplInfo.UpdateBehavior.AppendLocation))
+					if len(part) == 0 {
+						return fmt.Errorf("can not find append location '%s' for file '%s'\n", tplInfo.UpdateBehavior.AppendLocation, filePath)
+					}
+					if len(part) != 2 {
+						return fmt.Errorf("do not support multiple append location '%s' for file '%s'\n", tplInfo.UpdateBehavior.AppendLocation, filePath)
+					}
+					buf := bytes.NewBuffer(nil)
+					err = writeBytes(buf, part[0], []byte(tplInfo.UpdateBehavior.AppendLocation), appendContent, part[1])
+					if err != nil {
+						return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'method'", filePath)
+					pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
 				}
-				logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'method'", filePath)
-				pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
 			} else if tplInfo.UpdateBehavior.AppendKey == "service" {
 				var appendContent []byte
 				for _, service := range idlPackageRenderInfo.ServiceInfos.Services {
@@ -550,17 +591,34 @@ func (pkgGen *HttpPackageGenerator) genSingleCustomizedFile(tplInfo *Template, f
 						return err
 					}
 				}
-				buf := bytes.NewBuffer(nil)
-				_, err = buf.Write(fileContent)
-				if err != nil {
-					return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
-				}
-				_, err = buf.Write(appendContent)
-				if err != nil {
-					return fmt.Errorf("append file(%s) failed, err: %v", tplInfo.Path, err)
+				if len(tplInfo.UpdateBehavior.AppendLocation) == 0 { // default, append to end of file
+					buf := bytes.NewBuffer(nil)
+					_, err = buf.Write(fileContent)
+					if err != nil {
+						return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					_, err = buf.Write(appendContent)
+					if err != nil {
+						return fmt.Errorf("append file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'service'", filePath)
+					pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
+				} else { // 'append location', append new content after 'append location'
+					part := bytes.Split(fileContent, []byte(tplInfo.UpdateBehavior.AppendLocation))
+					if len(part) == 0 {
+						return fmt.Errorf("can not find append location '%s' for file '%s'\n", tplInfo.UpdateBehavior.AppendLocation, filePath)
+					}
+					if len(part) != 2 {
+						return fmt.Errorf("do not support multiple append location '%s' for file '%s'\n", tplInfo.UpdateBehavior.AppendLocation, filePath)
+					}
+					buf := bytes.NewBuffer(nil)
+					err = writeBytes(buf, part[0], []byte(tplInfo.UpdateBehavior.AppendLocation), appendContent, part[1])
+					if err != nil {
+						return fmt.Errorf("write file(%s) failed, err: %v", tplInfo.Path, err)
+					}
+					logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'service'", filePath)
+					pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
 				}
-				logs.Infof("append content for file '%s', because the update behavior is 'Append' and appendKey is 'service'", filePath)
-				pkgGen.files = append(pkgGen.files, File{filePath, buf.String(), false, ""})
 			} else { // add append content to the file directly
 				data := CustomizedFileForIDL{
 					IDLPackageRenderInfo: &idlPackageRenderInfo,
@@ -581,3 +639,14 @@ func (pkgGen *HttpPackageGenerator) genSingleCustomizedFile(tplInfo *Template, f
 
 	return nil
 }
+
+func writeBytes(buf *bytes.Buffer, bytes ...[]byte) error {
+	for _, b := range bytes {
+		_, err := buf.Write(b)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

cmd/hz/generator/handler.go

@@ -37,6 +37,9 @@ type HttpMethod struct {
 	Path            string
 	Serializer      string
 	OutputDir       string
+	RefPackage      string
+	RefPackageAlias string
+	ModelPackage    map[string]string
 	// Annotations     map[string]string
 	Models map[string]*model.Model
 }
@@ -76,13 +79,24 @@ func (pkgGen *HttpPackageGenerator) genHandler(pkg *HttpPackage, handlerDir, han
 				}
 			}
 		} else { // generate handler service
+			tmpHandlerDir := handlerDir
+			tmpHandlerPackage := handlerPackage
+			if len(s.ServiceGenDir) != 0 {
+				tmpHandlerDir = s.ServiceGenDir
+				tmpHandlerPackage = util.SubPackage(pkgGen.ProjPackage, tmpHandlerDir)
+			}
 			handler = Handler{
-				FilePath:    filepath.Join(handlerDir, util.ToSnakeCase(s.Name)+".go"),
-				PackageName: util.SplitPackage(handlerPackage, ""),
+				FilePath:    filepath.Join(tmpHandlerDir, util.ToSnakeCase(s.Name)+".go"),
+				PackageName: util.SplitPackage(tmpHandlerPackage, ""),
 				Methods:     s.Methods,
 				ProjPackage: pkgGen.ProjPackage,
 			}
 
+			for _, m := range s.Methods {
+				m.RefPackage = tmpHandlerPackage
+				m.RefPackageAlias = util.BaseName(tmpHandlerPackage, "")
+			}
+
 			if err := pkgGen.processHandler(&handler, root, "", "", false); err != nil {
 				return fmt.Errorf("generate handler %s failed, err: %v", handler.FilePath, err.Error())
 			}
@@ -145,6 +159,9 @@ func (pkgGen *HttpPackageGenerator) processHandler(handler *Handler, root *Route
 }
 
 func (pkgGen *HttpPackageGenerator) updateHandler(handler interface{}, handlerTpl, filePath string, noRepeat bool) error {
+	if pkgGen.tplsInfo[handlerTpl].Disable {
+		return nil
+	}
 	isExist, err := util.PathExist(filePath)
 	if err != nil {
 		return err

cmd/hz/generator/package.go

@@ -44,6 +44,8 @@ type Service struct {
 	ClientMethods []*ClientMethod
 	Models        []*model.Model // all dependency models
 	BaseDomain    string         // base domain for client code
+	ServiceGroup  string         // service level router group
+	ServiceGenDir string         // handler_dir for handler_by_service
 }
 
 // HttpPackageGenerator is used to record the configuration related to generating hertz http code.
@@ -61,6 +63,7 @@ type HttpPackageGenerator struct {
 	IdlClientDir   string // client dir for "client" command
 	ForceClientDir string // client dir without namespace for "client" command
 	BaseDomain     string // request domain for "client" command
+	ServiceGenDir  string
 
 	NeedModel            bool
 	HandlerByMethod      bool // generate handler files with method dimension

cmd/hz/generator/router.go

@@ -218,6 +218,8 @@ func (routerNode *RouterNode) Insert(name string, method *HttpMethod, handlerTyp
 				c.HandlerPackageAlias = pkgAlias
 				c.Handler = pkgAlias + "." + method.Name
 				c.HandlerPackage = handlerPkg
+				method.RefPackage = c.HandlerPackage
+				method.RefPackageAlias = c.HandlerPackageAlias
 			} else { // generate handler by service
 				c.Handler = handlerType + "." + method.Name
 			}
@@ -290,6 +292,9 @@ var (
 )
 
 func (pkgGen *HttpPackageGenerator) updateRegister(pkg, rDir, pkgName string) error {
+	if pkgGen.tplsInfo[registerTplName].Disable {
+		return nil
+	}
 	register := RegisterInfo{
 		PackageName: filepath.Base(rDir),
 		DepPkgAlias: strings.ReplaceAll(pkgName, "/", "_"),
@@ -365,8 +370,10 @@ func (pkgGen *HttpPackageGenerator) genRouter(pkg *HttpPackage, root *RouterNode
 	// store router info
 	pkg.RouterInfo = &router
 
-	if err := pkgGen.TemplateGenerator.Generate(router, routerTplName, router.FilePath, false); err != nil {
-		return fmt.Errorf("generate router %s failed, err: %v", router.FilePath, err.Error())
+	if !pkgGen.tplsInfo[routerTplName].Disable {
+		if err := pkgGen.TemplateGenerator.Generate(router, routerTplName, router.FilePath, false); err != nil {
+			return fmt.Errorf("generate router %s failed, err: %v", router.FilePath, err.Error())
+		}
 	}
 	if err := pkgGen.updateMiddlewareReg(router, middlewareTplName, filepath.Join(routerDir, "middleware.go")); err != nil {
 		return fmt.Errorf("generate middleware %s failed, err: %v", filepath.Join(routerDir, "middleware.go"), err.Error())
@@ -379,6 +386,9 @@ func (pkgGen *HttpPackageGenerator) genRouter(pkg *HttpPackage, root *RouterNode
 }
 
 func (pkgGen *HttpPackageGenerator) updateMiddlewareReg(router interface{}, middlewareTpl, filePath string) error {
+	if pkgGen.tplsInfo[middlewareTpl].Disable {
+		return nil
+	}
 	isExist, err := util.PathExist(filePath)
 	if err != nil {
 		return err

cmd/hz/generator/template.go

@@ -43,22 +43,24 @@ const (
 )
 
 type Template struct {
-	Default        bool           // Update command behavior; skip/cover/append
+	Default        bool           // Is it the default template
 	Path           string         `yaml:"path"`            // The generated path and its filename, such as biz/handler/ping.go
 	Delims         [2]string      `yaml:"delims"`          // Template Action Instruction Identifier, default: "{{}}"
 	Body           string         `yaml:"body"`            // Render template, currently only supports go template syntax
+	Disable        bool           `yaml:"disable"`         // Disable generating file, used to disable default package template
 	LoopMethod     bool           `yaml:"loop_method"`     // Loop generate files based on "method"
 	LoopService    bool           `yaml:"loop_service"`    // Loop generate files based on "service"
 	UpdateBehavior UpdateBehavior `yaml:"update_behavior"` // Update command behavior; 0:unchanged, 1:regenerate, 2:append
 }
 
 type UpdateBehavior struct {
-	Type string `yaml:"type"` // Update behavior type; skip/cover/append
+	Type string `yaml:"type"` // Update behavior type: skip/cover/append
 	// the following variables are used for append update
-	AppendKey string   `yaml:"append_key"`         // Append content based in key; for example: 'method'/'service'
-	InsertKey string   `yaml:"insert_key"`         // Insert content by "insert_key"
-	AppendTpl string   `yaml:"append_content_tpl"` // Append content if UpdateBehavior is "append"
-	ImportTpl []string `yaml:"import_tpl"`         // Import insert template
+	AppendKey      string   `yaml:"append_key"`         // Append content based in key; for example: 'method'/'service'
+	InsertKey      string   `yaml:"insert_key"`         // Insert content by "insert_key"
+	AppendTpl      string   `yaml:"append_content_tpl"` // Append content if UpdateBehavior is "append"
+	ImportTpl      []string `yaml:"import_tpl"`         // Import insert template
+	AppendLocation string   `yaml:"append_location"`    // AppendLocation specifies the location of append,  the default is the end of the file
 }
 
 // TemplateGenerator contains information about the output template

cmd/hz/thrift/ast.go

@@ -73,7 +73,20 @@ func astToService(ast *parser.Thrift, resolver *Resolver, args *config.Argument)
 				service.BaseDomain = domainAnno[0]
 			}
 		}
-
+		service.ServiceGroup = ""
+		groupAnno := getAnnotation(s.Annotations, ApiServiceGroup)
+		if len(groupAnno) == 1 {
+			if args.CmdType != meta.CmdClient {
+				service.ServiceGroup = groupAnno[0]
+			}
+		}
+		service.ServiceGenDir = ""
+		serviceGenDirAnno := getAnnotation(s.Annotations, ApiServiceGenDir)
+		if len(serviceGenDirAnno) == 1 {
+			if args.CmdType != meta.CmdClient {
+				service.ServiceGenDir = serviceGenDirAnno[0]
+			}
+		}
 		ms := s.GetFunctions()
 		if len(s.Extends) != 0 && args.EnableExtends {
 			// all the services that are extended to the current service

cmd/hz/thrift/tags.go

@@ -61,7 +61,9 @@ const (
 )
 
 const (
-	ApiBaseDomain = "api.base_domain"
+	ApiBaseDomain    = "api.base_domain"
+	ApiServiceGroup  = "api.service_group"
+	ApiServiceGenDir = "api.service_gen_dir"
 )
 
 var (