support custom maxBytes in API file

zeromicro · Jan 26, 2023 · a12ff89 · a12ff89
1 parent ab9eeff
commit a12ff89
Showing 1 changed file with 16 additions and 2 deletions.
diff --git a/tools/goctl/api/gogen/genroutes.go b/tools/goctl/api/gogen/genroutes.go
@@ -5,6 +5,7 @@ import (
 	"os"
 	"path"
 	"sort"
+	"strconv"
 	"strings"
 	"text/template"
 	"time"
@@ -36,7 +37,7 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
 `
 	routesAdditionTemplate = `
 	server.AddRoutes(
-		{{.routes}} {{.jwt}}{{.signature}} {{.prefix}} {{.timeout}}
+		{{.routes}} {{.jwt}}{{.signature}} {{.prefix}} {{.timeout}} {{.maxBytes}}
 	)
 `
 	timeoutThreshold = time.Millisecond
@@ -64,6 +65,7 @@ type (
 		middlewares      []string
 		prefix           string
 		jwtTrans         string
+		maxBytes         string
 	}
 	route struct {
 		method  string
@@ -127,10 +129,20 @@ rest.WithPrefix("%s"),`, g.prefix)
 				return fmt.Errorf("timeout should not less than 1ms, now %v", duration)
 			}
 
-			timeout = fmt.Sprintf("rest.WithTimeout(%d * time.Millisecond),", duration/time.Millisecond)
+			timeout = fmt.Sprintf("\n rest.WithTimeout(%d * time.Millisecond),", duration/time.Millisecond)
 			hasTimeout = true
 		}
 
+		var maxBytes string
+		if len(g.maxBytes) > 0 {
+			_, err := strconv.ParseInt(g.maxBytes, 10, 64)
+			if err != nil {
+				return fmt.Errorf("maxBytes %s parse error,it is an invalid number", g.maxBytes)
+			}
+
+			maxBytes = fmt.Sprintf("\n rest.WithMaxBytes(%s),", g.maxBytes)
+		}
+
 		var routes string
 		if len(g.middlewares) > 0 {
 			gbuilder.WriteString("\n}...,")
@@ -152,6 +164,7 @@ rest.WithPrefix("%s"),`, g.prefix)
 			"signature": signature,
 			"prefix":    prefix,
 			"timeout":   timeout,
+			"maxBytes":  maxBytes,
 		}); err != nil {
 			return err
 		}
@@ -230,6 +243,7 @@ func getRoutes(api *spec.ApiSpec) ([]group, error) {
 		}
 
 		groupedRoutes.timeout = g.GetAnnotation("timeout")
+		groupedRoutes.maxBytes = g.GetAnnotation("maxBytes")
 
 		jwt := g.GetAnnotation("jwt")
 		if len(jwt) > 0 {