kratos-utils/http/middleware/page.go

package middleware

import (
	"context"
	"strconv"

	"github.com/go-kratos/kratos/v2/middleware"
	"github.com/go-kratos/kratos/v2/transport/http"
)

var pageURI = make(map[string]bool)

func Page(handler middleware.Handler) middleware.Handler {
	return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
		if info, ok := http.FromServerContext(ctx); ok {
			uri := info.Request.RequestURI
			if _, ok1 := pageURI[uri]; !ok1 {
				return handler(ctx, req)
			}
			page := 0
			pageSize := 10
			query := info.Request.URL.Query()
			pageStr := query.Get("page")
			if pageStr != "" {
				page, err = strconv.Atoi(pageStr)
				if err != nil {
					return nil, err
				}
			}
			pageSizeStr := query.Get("page_size")
			if pageSizeStr != "" {
				pageSize, err = strconv.Atoi(pageSizeStr)
				if err != nil {
					return nil, err
				}
			}
			if pageSize == 0 {
				pageSize = 10
			}

			ctx = context.WithValue(ctx, "page", page)
			ctx = context.WithValue(ctx, "page_size", pageSize)
		}

		return handler(ctx, req)
	}
}

func AddPageURI(r string) {
	pageURI[r] = true
}