-
Notifications
You must be signed in to change notification settings - Fork 3.3k
/
middleware.go
40 lines (32 loc) · 1 KB
/
middleware.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
package querylimits
import (
"net/http"
"github.com/go-kit/log"
"github.com/go-kit/log/level"
"github.com/grafana/dskit/middleware"
util_log "github.com/grafana/loki/v3/pkg/util/log"
)
type queryLimitsMiddleware struct {
logger log.Logger
}
// NewQueryLimitsMiddleware creates a middleware that extracts the query limits
// policy from the HTTP header and injects it into the context of the request.
func NewQueryLimitsMiddleware(logger log.Logger) middleware.Interface {
return &queryLimitsMiddleware{
logger: logger,
}
}
// Wrap implements the middleware interface
func (l *queryLimitsMiddleware) Wrap(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
limits, err := ExtractQueryLimitsHTTP(r)
if err != nil {
level.Warn(util_log.Logger).Log("msg", "could not extract query limits from header", "err", err)
limits = nil
}
if limits != nil {
r = r.Clone(InjectQueryLimitsContext(r.Context(), *limits))
}
next.ServeHTTP(w, r)
})
}