feat: add requestid middleware

This commit is contained in:
tiglog 2023-10-17 17:21:53 +08:00
parent e2ac9a7704
commit 7f5da8e5e1
2 changed files with 103 additions and 0 deletions

41
gweb/requestid/options.go Normal file
View File

@ -0,0 +1,41 @@
//
// options.go
// Copyright (C) 2023 tiglog <me@tiglog.com>
//
// Distributed under terms of the MIT license.
//
package requestid
import "github.com/gin-gonic/gin"
// Option for queue system
type Option func(*config)
type (
Generator func() string
Handler func(c *gin.Context, requestID string)
)
type HeaderStrKey string
// WithGenerator set generator function
func WithGenerator(g Generator) Option {
return func(cfg *config) {
cfg.generator = g
}
}
// WithCustomHeaderStrKey set custom header key for request id
func WithCustomHeaderStrKey(s HeaderStrKey) Option {
return func(cfg *config) {
cfg.headerKey = s
}
}
// WithHandler set handler function for request id with context
func WithHandler(handler Handler) Option {
return func(cfg *config) {
cfg.handler = handler
}
}

62
gweb/requestid/reqid.go Normal file
View File

@ -0,0 +1,62 @@
//
// reqid.go
// Copyright (C) 2023 tiglog <me@tiglog.com>
//
// Distributed under terms of the MIT license.
//
package requestid
import (
"github.com/gin-gonic/gin"
"github.com/rs/xid"
)
var headerXRequestID string
// Config defines the config for RequestID middleware
type config struct {
// Generator defines a function to generate an ID.
// Optional. Default: func() string {
// return xid.New().String()
// }
generator Generator
headerKey HeaderStrKey
handler Handler
}
// New initializes the RequestID middleware.
func New(opts ...Option) gin.HandlerFunc {
cfg := &config{
generator: func() string {
return xid.New().String()
},
headerKey: "X-Request-ID",
}
for _, opt := range opts {
opt(cfg)
}
headerXRequestID = string(cfg.headerKey)
return func(c *gin.Context) {
// Get id from request
rid := c.GetHeader(headerXRequestID)
if rid == "" {
rid = cfg.generator()
c.Request.Header.Add(headerXRequestID, rid)
}
if cfg.handler != nil {
cfg.handler(c, rid)
}
// Set the id to ensure that the requestid is in the response
c.Header(headerXRequestID, rid)
c.Next()
}
}
// Get returns the request identifier
func Get(c *gin.Context) string {
return c.GetHeader(headerXRequestID)
}