Repository: zeromicro/go-zero
Pull request #1460 (2022-01-17 12:54:29)

Fix/issue#1289
**Description**
Support api code generation with `jwtTransition`

**Usage**
declare jwt transition with key `jwtTransition` as following code:
```api
@server (
	jwt: Auth
	jwtTransition: Trans // new supported
)
service user-api {
	@handler login
	post /user/login (LoginReq) returns (LoginReply)
}
```

**code generation product**
```go
func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
	server.AddRoutes(
		[]rest.Route{
			{
				Method:  http.MethodPost,
				Path:    "/user/login",
				Handler: loginHandler(serverCtx),
			},
		},
		rest.WithJwt(serverCtx.Config.Auth.AccessSecret),
		rest.WithJwtTransition(serverCtx.Config.Trans.PrevSecret, serverCtx.Config.Trans.Secret),// jwt Trans
	)
}
```


Patch:
diff --git a/tools/goctl/api/gogen/genconfig.go b/tools/goctl/api/gogen/genconfig.go
index ce84aba3..1be07dcc 100644
--- a/tools/goctl/api/gogen/genconfig.go
+++ b/tools/goctl/api/gogen/genconfig.go
@@ -19,6 +19,7 @@ import {{.authImport}}
 type Config struct {
 	rest.RestConf
 	{{.auth}}
+	{{.jwtTrans}}
 }
 `
 
@@ -26,6 +27,11 @@ type Config struct {
 		AccessSecret string
 		AccessExpire int64
 	}
+`
+	jwtTransTemplate = ` struct {
+		Secret     string
+		PrevSecret string
+	}
 `
 )
 
@@ -40,6 +46,12 @@ func genConfig(dir string, cfg *config.Config, api *spec.ApiSpec) error {
 	for _, item := range authNames {
 		auths = append(auths, fmt.Sprintf("%s %s", item, jwtTemplate))
 	}
+
+	jwtTransNames := getJwtTrans(api)
+	var jwtTransList []string
+	for _, item := range jwtTransNames {
+		jwtTransList = append(jwtTransList, fmt.Sprintf("%s %s", item, jwtTransTemplate))
+	}
 	authImportStr := fmt.Sprintf("\"%s/rest\"", vars.ProjectOpenSourceURL)
 
 	return genFile(fileGenConfig{
@@ -53,6 +65,7 @@ func genConfig(dir string, cfg *config.Config, api *spec.ApiSpec) error {
 		data: map[string]string{
 			"authImport": authImportStr,
 			"auth":       strings.Join(auths, "\n"),
+			"jwtTrans":   strings.Join(jwtTransList, "\n"),
 		},
 	})
 }
diff --git a/tools/goctl/api/gogen/genmiddleware.go b/tools/goctl/api/gogen/genmiddleware.go
index 1200cce7..f6bc06d4 100644
--- a/tools/goctl/api/gogen/genmiddleware.go
+++ b/tools/goctl/api/gogen/genmiddleware.go
@@ -45,6 +45,8 @@ func genMiddleware(dir string, cfg *config.Config, api *spec.ApiSpec) error {
 			subdir:          middlewareDir,
 			filename:        filename + ".go",
 			templateName:    "contextTemplate",
+			category:        category,
+			templateFile:    middlewareImplementCodeFile,
 			builtinTemplate: middlewareImplementCode,
 			data: map[string]string{
 				"name": strings.Title(name),
diff --git a/tools/goctl/api/gogen/genroutes.go b/tools/goctl/api/gogen/genroutes.go
index ed486422..ef12fd8b 100644
--- a/tools/goctl/api/gogen/genroutes.go
+++ b/tools/goctl/api/gogen/genroutes.go
@@ -17,6 +17,7 @@ import (
 )
 
 const (
+	jwtTransKey    = "jwtTransition"
 	routesFilename = "routes"
 	routesTemplate = `// Code generated by goctl. DO NOT EDIT.
 package handler
@@ -58,6 +59,7 @@ type (
 		authName         string
 		middlewares      []string
 		prefix           string
+		jwtTrans         string
 	}
 	route struct {
 		method  string
@@ -73,7 +75,12 @@ func genRoutes(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error
 		return err
 	}
 
-	gt := template.Must(template.New("groupTemplate").Parse(routesAdditionTemplate))
+	templateText, err := pathx.LoadTemplate(category, routesAdditionTemplateFile, routesAdditionTemplate)
+	if err != nil {
+		return err
+	}
+
+	gt := template.Must(template.New("groupTemplate").Parse(templateText))
 	for _, g := range groups {
 		var gbuilder strings.Builder
 		gbuilder.WriteString("[]rest.Route{")
@@ -91,6 +98,9 @@ func genRoutes(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error
 		if g.jwtEnabled {
 			jwt = fmt.Sprintf("\n rest.WithJwt(serverCtx.Config.%s.AccessSecret),", g.authName)
 		}
+		if len(g.jwtTrans) > 0 {
+			jwt = jwt + fmt.Sprintf("\n rest.WithJwtTransition(serverCtx.Config.%s.PrevSecret,serverCtx.Config.%s.Secret),", g.jwtTrans, g.jwtTrans)
+		}
 		var signature, prefix string
 		if g.signatureEnabled {
 			signature = "\n rest.WithSignature(serverCtx.Config.Signature),"
@@ -139,8 +149,8 @@ rest.WithPrefix("%s"),`, g.prefix)
 		subdir:          handlerDir,
 		filename:        routeFilename,
 		templateName:    "routesTemplate",
-		category:        "",
-		templateFile:    "",
+		category:        category,
+		templateFile:    routesTemplateFile,
 		builtinTemplate: routesTemplate,
 		data: map[string]string{
 			"importPackages":  genRouteImports(rootPkg, api),
@@ -200,6 +210,11 @@ func getRoutes(api *spec.ApiSpec) ([]group, error) {
 			groupedRoutes.authName = jwt
 			groupedRoutes.jwtEnabled = true
 		}
+		jwtTrans := g.GetAnnotation(jwtTransKey)
+		if len(jwtTrans) > 0 {
+			groupedRoutes.jwtTrans = jwtTrans
+		}
+
 		signature := g.GetAnnotation("signature")
 		if signature == "true" {
 			groupedRoutes.signatureEnabled = true
diff --git a/tools/goctl/api/gogen/gentypes.go b/tools/goctl/api/gogen/gentypes.go
index ae4dc4d1..5bea40dd 100644
--- a/tools/goctl/api/gogen/gentypes.go
+++ b/tools/goctl/api/gogen/gentypes.go
@@ -63,8 +63,8 @@ func genTypes(dir string, cfg *config.Config, api *spec.ApiSpec) error {
 		subdir:          typesDir,
 		filename:        typeFilename,
 		templateName:    "typesTemplate",
-		category:        "",
-		templateFile:    "",
+		category:        category,
+		templateFile:    typesTemplateFile,
 		builtinTemplate: typesTemplate,
 		data: map[string]interface{}{
 			"types":        val,
diff --git a/tools/goctl/api/gogen/template.go b/tools/goctl/api/gogen/template.go
index 96f514f5..7f80c15d 100644
--- a/tools/goctl/api/gogen/template.go
+++ b/tools/goctl/api/gogen/template.go
@@ -8,22 +8,30 @@ import (
 )
 
 const (
-	category            = "api"
-	configTemplateFile  = "config.tpl"
-	contextTemplateFile = "context.tpl"
-	etcTemplateFile     = "etc.tpl"
-	handlerTemplateFile = "handler.tpl"
-	logicTemplateFile   = "logic.tpl"
-	mainTemplateFile    = "main.tpl"
+	category                    = "api"
+	configTemplateFile          = "config.tpl"
+	contextTemplateFile         = "context.tpl"
+	etcTemplateFile             = "etc.tpl"
+	handlerTemplateFile         = "handler.tpl"
+	logicTemplateFile           = "logic.tpl"
+	mainTemplateFile            = "main.tpl"
+	middlewareImplementCodeFile = "middleware.tpl"
+	routesTemplateFile          = "routes.tpl"
+	routesAdditionTemplateFile  = "route-addition.tpl"
+	typesTemplateFile           = "types.tpl"
 )
 
 var templates = map[string]string{
-	configTemplateFile:  configTemplate,
-	contextTemplateFile: contextTemplate,
-	etcTemplateFile:     etcTemplate,
-	handlerTemplateFile: handlerTemplate,
-	logicTemplateFile:   logicTemplate,
-	mainTemplateFile:    mainTemplate,
+	configTemplateFile:          configTemplate,
+	contextTemplateFile:         contextTemplate,
+	etcTemplateFile:             etcTemplate,
+	handlerTemplateFile:         handlerTemplate,
+	logicTemplateFile:           logicTemplate,
+	mainTemplateFile:            mainTemplate,
+	middlewareImplementCodeFile: middlewareImplementCode,
+	routesTemplateFile:          routesTemplate,
+	routesAdditionTemplateFile:  routesAdditionTemplate,
+	typesTemplateFile:           typesTemplate,
 }
 
 // Category returns the category of the api files.
diff --git a/tools/goctl/api/gogen/util.go b/tools/goctl/api/gogen/util.go
index 3eef54ba..cd2f7cce 100644
--- a/tools/goctl/api/gogen/util.go
+++ b/tools/goctl/api/gogen/util.go
@@ -111,6 +111,17 @@ func getAuths(api *spec.ApiSpec) []string {
 	return authNames.KeysStr()
 }
 
+func getJwtTrans(api *spec.ApiSpec) []string {
+	jwtTransList := collection.NewSet()
+	for _, g := range api.Service.Groups {
+		jt := g.GetAnnotation(jwtTransKey)
+		if len(jt) > 0 {
+			jwtTransList.Add(jt)
+		}
+	}
+	return jwtTransList.KeysStr()
+}
+
 func getMiddleware(api *spec.ApiSpec) []string {
 	result := collection.NewSet()
 	for _, g := range api.Service.Groups {
diff --git a/tools/goctl/model/sql/gen/delete.go b/tools/goctl/model/sql/gen/delete.go
index 47e9fb31..9f401fb4 100644
--- a/tools/goctl/model/sql/gen/delete.go
+++ b/tools/goctl/model/sql/gen/delete.go
@@ -38,6 +38,7 @@ func genDelete(table Table, withCache, postgreSql bool) (string, string, error)
 			"originalPrimaryKey":        wrapWithRawString(table.PrimaryKey.Name.Source(), postgreSql),
 			"keyValues":                 strings.Join(keyVariableSet.KeysStr(), ", "),
 			"postgreSql":                postgreSql,
+			"data":                      table,
 		})
 	if err != nil {
 		return "", "", err
@@ -54,6 +55,7 @@ func genDelete(table Table, withCache, postgreSql bool) (string, string, error)
 		Execute(map[string]interface{}{
 			"lowerStartCamelPrimaryKey": stringx.From(table.PrimaryKey.Name.ToCamel()).Untitle(),
 			"dataType":                  table.PrimaryKey.DataType,
+			"data":                      table,
 		})
 	if err != nil {
 		return "", "", err
diff --git a/tools/goctl/model/sql/gen/field.go b/tools/goctl/model/sql/gen/field.go
index bca2f17d..b4397dbe 100644
--- a/tools/goctl/model/sql/gen/field.go
+++ b/tools/goctl/model/sql/gen/field.go
@@ -9,11 +9,11 @@ import (
 	"github.com/tal-tech/go-zero/tools/goctl/util/pathx"
 )
 
-func genFields(fields []*parser.Field) (string, error) {
+func genFields(table Table, fields []*parser.Field) (string, error) {
 	var list []string
 
 	for _, field := range fields {
-		result, err := genField(field)
+		result, err := genField(table, field)
 		if err != nil {
 			return "", err
 		}
@@ -24,8 +24,8 @@ func genFields(fields []*parser.Field) (string, error) {
 	return strings.Join(list, "\n"), nil
 }
 
-func genField(field *parser.Field) (string, error) {
-	tag, err := genTag(field.NameOriginal)
+func genField(table Table, field *parser.Field) (string, error) {
+	tag, err := genTag(table, field.NameOriginal)
 	if err != nil {
 		return "", err
 	}
@@ -43,6 +43,7 @@ func genField(field *parser.Field) (string, error) {
 			"tag":        tag,
 			"hasComment": field.Comment != "",
 			"comment":    field.Comment,
+			"data":       table,
 		})
 	if err != nil {
 		return "", err
diff --git a/tools/goctl/model/sql/gen/findone.go b/tools/goctl/model/sql/gen/findone.go
index a5acd62a..a4a0d0cd 100644
--- a/tools/goctl/model/sql/gen/findone.go
+++ b/tools/goctl/model/sql/gen/findone.go
@@ -26,6 +26,7 @@ func genFindOne(table Table, withCache, postgreSql bool) (string, string, error)
 			"cacheKey":                  table.PrimaryCacheKey.KeyExpression,
 			"cacheKeyVariable":          table.PrimaryCacheKey.KeyLeft,
 			"postgreSql":                postgreSql,
+			"data":                      table,
 		})
 	if err != nil {
 		return "", "", err
@@ -42,6 +43,7 @@ func genFindOne(table Table, withCache, postgreSql bool) (string, string, error)
 			"upperStartCamelObject":     camel,
 			"lowerStartCamelPrimaryKey": stringx.From(table.PrimaryKey.Name.ToCamel()).Untitle(),
 			"dataType":                  table.PrimaryKey.DataType,
+			"data":                      table,
 		})
 	if err != nil {
 		return "", "", err
diff --git a/tools/goctl/model/sql/gen/findonebyfield.go b/tools/goctl/model/sql/gen/findonebyfield.go
index 35373d9a..029eb498 100644
--- a/tools/goctl/model/sql/gen/findonebyfield.go
+++ b/tools/goctl/model/sql/gen/findonebyfield.go
@@ -40,6 +40,7 @@ func genFindOneByField(table Table, withCache, postgreSql bool) (*findOneCode, e
 			"upperStartCamelPrimaryKey": table.PrimaryKey.Name.ToCamel(),
 			"originalField":             originalFieldString,
 			"postgreSql":                postgreSql,
+			"data":                      table,
 		})
 		if err != nil {
 			return nil, err
@@ -71,6 +72,7 @@ func genFindOneByField(table Table, withCache, postgreSql bool) (*findOneCode, e
 			"upperStartCamelObject": camelTableName,
 			"upperField":            key.FieldNameJoin.Camel().With("").Source(),
 			"in":                    in,
+			"data":                  table,
 		})
 		if err != nil {
 			return nil, err
@@ -91,6 +93,7 @@ func genFindOneByField(table Table, withCache, postgreSql bool) (*findOneCode, e
 			"lowerStartCamelObject": stringx.From(camelTableName).Untitle(),
 			"originalPrimaryField":  wrapWithRawString(table.PrimaryKey.Name.Source(), postgreSql),
 			"postgreSql":            postgreSql,
+			"data":                  table,
 		})
 		if err != nil {
 			return nil, err
diff --git a/tools/goctl/model/sql/gen/gen.go b/tools/goctl/model/sql/gen/gen.go
index 8f2b50b6..940d48e4 100644
--- a/tools/goctl/model/sql/gen/gen.go
+++ b/tools/goctl/model/sql/gen/gen.go
@@ -217,17 +217,17 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 
 	primaryKey, uniqueKey := genCacheKeys(in)
 
-	importsCode, err := genImports(withCache, in.ContainsTime())
-	if err != nil {
-		return "", err
-	}
-
 	var table Table
 	table.Table = in
 	table.PrimaryCacheKey = primaryKey
 	table.UniqueCacheKey = uniqueKey
 	table.ContainsUniqueCacheKey = len(uniqueKey) > 0
 
+	importsCode, err := genImports(withCache, in.ContainsTime(), table)
+	if err != nil {
+		return "", err
+	}
+
 	varsCode, err := genVars(table, withCache, g.isPostgreSql)
 	if err != nil {
 		return "", err
@@ -284,7 +284,7 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 		cacheExtra:  ret.cacheExtra,
 	}
 
-	output, err := g.executeModel(code)
+	output, err := g.executeModel(table, code)
 	if err != nil {
 		return "", err
 	}
@@ -292,7 +292,7 @@ func (g *defaultGenerator) genModel(in parser.Table, withCache bool) (string, er
 	return output.String(), nil
 }
 
-func (g *defaultGenerator) executeModel(code *code) (*bytes.Buffer, error) {
+func (g *defaultGenerator) executeModel(table Table, code *code) (*bytes.Buffer, error) {
 	text, err := pathx.LoadTemplate(category, modelTemplateFile, template.Model)
 	if err != nil {
 		return nil, err
@@ -311,6 +311,7 @@ func (g *defaultGenerator) executeModel(code *code) (*bytes.Buffer, error) {
 		"update":      code.updateCode,
 		"delete":      code.deleteCode,
 		"extraMethod": code.cacheExtra,
+		"data":        table,
 	})
 	if err != nil {
 		return nil, err
diff --git a/tools/goctl/model/sql/gen/imports.go b/tools/goctl/model/sql/gen/imports.go
index a89cb21b..a458147a 100644
--- a/tools/goctl/model/sql/gen/imports.go
+++ b/tools/goctl/model/sql/gen/imports.go
@@ -6,7 +6,7 @@ import (
 	"github.com/tal-tech/go-zero/tools/goctl/util/pathx"
 )
 
-func genImports(withCache, timeImport bool) (string, error) {
+func genImports(withCache, timeImport bool, table Table) (string, error) {
 	if withCache {
 		text, err := pathx.LoadTemplate(category, importsTemplateFile, template.Imports)
 		if err != nil {
@@ -15,6 +15,7 @@ func genImports(withCache, timeImport bool) (string, error) {
 
 		buffer, err := util.With("import").Parse(text).Execute(map[string]interface{}{
 			"time": timeImport,
+			"data": table,
 		})
 		if err != nil {
 			return "", err
@@ -30,6 +31,7 @@ func genImports(withCache, timeImport bool) (string, error) {
 
 	buffer, err := util.With("import").Parse(text).Execute(map[string]interface{}{
 		"time": timeImport,
+		"data": table,
 	})
 	if err != nil {
 		return "", err
diff --git a/tools/goctl/model/sql/gen/insert.go b/tools/goctl/model/sql/gen/insert.go
index c2bd2cf0..a8ecee43 100644
--- a/tools/goctl/model/sql/gen/insert.go
+++ b/tools/goctl/model/sql/gen/insert.go
@@ -62,6 +62,7 @@ func genInsert(table Table, withCache, postgreSql bool) (string, string, error)
 			"expressionValues":      strings.Join(expressionValues, ", "),
 			"keys":                  strings.Join(keySet.KeysStr(), "\n"),
 			"keyValues":             strings.Join(keyVariableSet.KeysStr(), ", "),
+			"data":                  table,
 		})
 	if err != nil {
 		return "", "", err
@@ -75,6 +76,7 @@ func genInsert(table Table, withCache, postgreSql bool) (string, string, error)
 
 	insertMethodOutput, err := util.With("insertMethod").Parse(text).Execute(map[string]interface{}{
 		"upperStartCamelObject": camel,
+		"data":                  table,
 	})
 	if err != nil {
 		return "", "", err
diff --git a/tools/goctl/model/sql/gen/new.go b/tools/goctl/model/sql/gen/new.go
index fec03df3..fb58812f 100644
--- a/tools/goctl/model/sql/gen/new.go
+++ b/tools/goctl/model/sql/gen/new.go
@@ -25,6 +25,7 @@ func genNew(table Table, withCache, postgreSql bool) (string, error) {
 			"table":                 t,
 			"withCache":             withCache,
 			"upperStartCamelObject": table.Name.ToCamel(),
+			"data":                  table,
 		})
 	if err != nil {
 		return "", err
diff --git a/tools/goctl/model/sql/gen/tag.go b/tools/goctl/model/sql/gen/tag.go
index 612089ac..247b140b 100644
--- a/tools/goctl/model/sql/gen/tag.go
+++ b/tools/goctl/model/sql/gen/tag.go
@@ -6,7 +6,7 @@ import (
 	"github.com/tal-tech/go-zero/tools/goctl/util/pathx"
 )
 
-func genTag(in string) (string, error) {
+func genTag(table Table, in string) (string, error) {
 	if in == "" {
 		return in, nil
 	}
@@ -18,6 +18,7 @@ func genTag(in string) (string, error) {
 
 	output, err := util.With("tag").Parse(text).Execute(map[string]interface{}{
 		"field": in,
+		"data":  table,
 	})
 	if err != nil {
 		return "", err
diff --git a/tools/goctl/model/sql/gen/types.go b/tools/goctl/model/sql/gen/types.go
index 57d3a37f..c9f5363e 100644
--- a/tools/goctl/model/sql/gen/types.go
+++ b/tools/goctl/model/sql/gen/types.go
@@ -8,7 +8,7 @@ import (
 
 func genTypes(table Table, methods string, withCache bool) (string, error) {
 	fields := table.Fields
-	fieldsString, err := genFields(fields)
+	fieldsString, err := genFields(table, fields)
 	if err != nil {
 		return "", err
 	}
@@ -25,6 +25,7 @@ func genTypes(table Table, methods string, withCache bool) (string, error) {
 			"method":                methods,
 			"upperStartCamelObject": table.Name.ToCamel(),
 			"fields":                fieldsString,
+			"data":                  table,
 		})
 	if err != nil {
 		return "", err
diff --git a/tools/goctl/model/sql/gen/update.go b/tools/goctl/model/sql/gen/update.go
index de2c6230..041c41c1 100644
--- a/tools/goctl/model/sql/gen/update.go
+++ b/tools/goctl/model/sql/gen/update.go
@@ -58,6 +58,7 @@ func genUpdate(table Table, withCache, postgreSql bool) (string, string, error)
 			"originalPrimaryKey":    wrapWithRawString(table.PrimaryKey.Name.Source(), postgreSql),
 			"expressionValues":      strings.Join(expressionValues, ", "),
 			"postgreSql":            postgreSql,
+			"data":                  table,
 		})
 	if err != nil {
 		return "", "", nil
@@ -73,6 +74,7 @@ func genUpdate(table Table, withCache, postgreSql bool) (string, string, error)
 		Parse(text).
 		Execute(map[string]interface{}{
 			"upperStartCamelObject": camelTableName,
+			"data":                  table,
 		})
 	if err != nil {
 		return "", "", nil
diff --git a/tools/goctl/model/sql/gen/vars.go b/tools/goctl/model/sql/gen/vars.go
index 58deff06..8688da22 100644
--- a/tools/goctl/model/sql/gen/vars.go
+++ b/tools/goctl/model/sql/gen/vars.go
@@ -31,6 +31,7 @@ func genVars(table Table, withCache, postgreSql bool) (string, error) {
 		"originalPrimaryKey":    wrapWithRawString(table.PrimaryKey.Name.Source(), postgreSql),
 		"withCache":             withCache,
 		"postgreSql":            postgreSql,
+		"data":                  table,
 	})
 	if err != nil {
 		return "", err


Test patch:
diff --git a/tools/goctl/api/gogen/gen_test.go b/tools/goctl/api/gogen/gen_test.go
index ba93856a..eeab8730 100644
--- a/tools/goctl/api/gogen/gen_test.go
+++ b/tools/goctl/api/gogen/gen_test.go
@@ -178,6 +178,7 @@ type Response struct {
 
 @server(
 	jwt: Auth
+	jwtTransition: Trans
 	middleware: TokenValidate
 )
 service A-api {