Bläddra i källkod

Added support for setting the parameter size accepted by the interface and custom timeout and maxbytes in API file (#1713)

* Added support for setting the parameter size accepted by the interface

* support custom timeout and maxbytes in API file

* support timeout used unit

* remove goctl maxBytes
Xiaoju Jiang 3 år sedan
förälder
incheckning
321dc2d410
5 ändrade filer med 74 tillägg och 2 borttagningar
  1. 10 1
      rest/engine.go
  2. 35 0
      rest/engine_test.go
  3. 7 0
      rest/server.go
  4. 1 0
      rest/types.go
  5. 21 1
      tools/goctl/api/gogen/genroutes.go

+ 10 - 1
rest/engine.go

@@ -94,7 +94,7 @@ func (ng *engine) bindRoute(fr featuredRoutes, router httpx.Router, metrics *sta
 		handler.TimeoutHandler(ng.checkedTimeout(fr.timeout)),
 		handler.RecoverHandler,
 		handler.MetricHandler(metrics),
-		handler.MaxBytesHandler(ng.conf.MaxBytes),
+		handler.MaxBytesHandler(ng.checkedMaxBytes(fr.maxBytes)),
 		handler.GunzipHandler,
 	)
 	chain = ng.appendAuthHandler(fr, chain, verifier)
@@ -127,6 +127,15 @@ func (ng *engine) checkedTimeout(timeout time.Duration) time.Duration {
 	return time.Duration(ng.conf.Timeout) * time.Millisecond
 }
 
+func (ng *engine) checkedMaxBytes(bytes int64) int64 {
+
+	if bytes > 0 {
+		return bytes
+	}
+
+	return ng.conf.MaxBytes
+}
+
 func (ng *engine) createMetrics() *stat.Metrics {
 	var metrics *stat.Metrics
 

+ 35 - 0
rest/engine_test.go

@@ -194,6 +194,41 @@ func TestEngine_checkedTimeout(t *testing.T) {
 	}
 }
 
+func TestEngine_checkedMaxBytes(t *testing.T) {
+	tests := []struct {
+		name     string
+		maxBytes int64
+		expect   int64
+	}{
+		{
+			name:   "not set",
+			expect: 1000,
+		},
+		{
+			name:     "less",
+			maxBytes: 500,
+			expect:   500,
+		},
+		{
+			name:     "equal",
+			maxBytes: 1000,
+			expect:   1000,
+		},
+		{
+			name:     "more",
+			maxBytes: 1500,
+			expect:   1500,
+		},
+	}
+
+	ng := newEngine(RestConf{
+		MaxBytes: 1000,
+	})
+	for _, test := range tests {
+		assert.Equal(t, test.expect, ng.checkedMaxBytes(test.maxBytes))
+	}
+}
+
 func TestEngine_notFoundHandler(t *testing.T) {
 	logx.Disable()
 

+ 7 - 0
rest/server.go

@@ -223,6 +223,13 @@ func WithTimeout(timeout time.Duration) RouteOption {
 	}
 }
 
+// WithMaxBytes returns a RouteOption to set maxBytes with given value.
+func WithMaxBytes(maxBytes int64) RouteOption {
+	return func(r *featuredRoutes) {
+		r.maxBytes = maxBytes
+	}
+}
+
 // WithTLSConfig returns a RunOption that with given tls config.
 func WithTLSConfig(cfg *tls.Config) RunOption {
 	return func(svr *Server) {

+ 1 - 0
rest/types.go

@@ -36,5 +36,6 @@ type (
 		jwt       jwtSetting
 		signature signatureSetting
 		routes    []Route
+		maxBytes  int64
 	}
 )

+ 21 - 1
tools/goctl/api/gogen/genroutes.go

@@ -7,6 +7,7 @@ import (
 	"sort"
 	"strings"
 	"text/template"
+	"time"
 
 	"github.com/zeromicro/go-zero/core/collection"
 	"github.com/zeromicro/go-zero/tools/goctl/api/spec"
@@ -34,7 +35,7 @@ func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
 `
 	routesAdditionTemplate = `
 	server.AddRoutes(
-		{{.routes}} {{.jwt}}{{.signature}} {{.prefix}}
+		{{.routes}} {{.jwt}}{{.signature}} {{.prefix}} {{.timeout}}
 	)
 `
 )
@@ -57,6 +58,8 @@ type (
 		jwtEnabled       bool
 		signatureEnabled bool
 		authName         string
+		timeout          string
+		timeoutEnable    bool
 		middlewares      []string
 		prefix           string
 		jwtTrans         string
@@ -110,6 +113,15 @@ func genRoutes(dir, rootPkg string, cfg *config.Config, api *spec.ApiSpec) error
 rest.WithPrefix("%s"),`, g.prefix)
 		}
 
+		var timeout string
+		if g.timeoutEnable {
+			duration, err := time.ParseDuration(g.timeout)
+			if err != nil {
+				panic(err)
+			}
+			timeout = fmt.Sprintf("rest.WithTimeout(%d),", duration)
+		}
+
 		var routes string
 		if len(g.middlewares) > 0 {
 			gbuilder.WriteString("\n}...,")
@@ -130,6 +142,7 @@ rest.WithPrefix("%s"),`, g.prefix)
 			"jwt":       jwt,
 			"signature": signature,
 			"prefix":    prefix,
+			"timeout":   timeout,
 		}); err != nil {
 			return err
 		}
@@ -205,6 +218,13 @@ func getRoutes(api *spec.ApiSpec) ([]group, error) {
 			})
 		}
 
+		timeout := g.GetAnnotation("timeout")
+
+		if len(timeout) > 0 {
+			groupedRoutes.timeoutEnable = true
+			groupedRoutes.timeout = timeout
+		}
+
 		jwt := g.GetAnnotation("jwt")
 		if len(jwt) > 0 {
 			groupedRoutes.authName = jwt