/
complexity.go
88 lines (69 loc) · 2.17 KB
/
complexity.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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
package extension
import (
"context"
"fmt"
"github.com/99designs/gqlgen/complexity"
"github.com/99designs/gqlgen/graphql"
"github.com/99designs/gqlgen/graphql/errcode"
"github.com/vektah/gqlparser/v2/gqlerror"
)
const errComplexityLimit = "COMPLEXITY_LIMIT_EXCEEDED"
// ComplexityLimit allows you to define a limit on query complexity
//
// If a query is submitted that exceeds the limit, a 422 status code will be returned.
type ComplexityLimit struct {
Func func(ctx context.Context, rc *graphql.OperationContext) int
es graphql.ExecutableSchema
}
var _ interface {
graphql.OperationContextMutator
graphql.HandlerExtension
} = &ComplexityLimit{}
const complexityExtension = "ComplexityLimit"
type ComplexityStats struct {
// The calculated complexity for this request
Complexity int
// The complexity limit for this request returned by the extension func
ComplexityLimit int
}
// FixedComplexityLimit sets a complexity limit that does not change
func FixedComplexityLimit(limit int) *ComplexityLimit {
return &ComplexityLimit{
Func: func(ctx context.Context, rc *graphql.OperationContext) int {
return limit
},
}
}
func (c ComplexityLimit) ExtensionName() string {
return complexityExtension
}
func (c *ComplexityLimit) Validate(schema graphql.ExecutableSchema) error {
if c.Func == nil {
return fmt.Errorf("ComplexityLimit func can not be nil")
}
c.es = schema
return nil
}
func (c ComplexityLimit) MutateOperationContext(ctx context.Context, rc *graphql.OperationContext) *gqlerror.Error {
op := rc.Doc.Operations.ForName(rc.OperationName)
complexity := complexity.Calculate(c.es, op, rc.Variables)
limit := c.Func(ctx, rc)
rc.Stats.SetExtension(complexityExtension, &ComplexityStats{
Complexity: complexity,
ComplexityLimit: limit,
})
if complexity > limit {
err := gqlerror.Errorf("operation has complexity %d, which exceeds the limit of %d", complexity, limit)
errcode.Set(err, errComplexityLimit)
return err
}
return nil
}
func GetComplexityStats(ctx context.Context) *ComplexityStats {
rc := graphql.GetOperationContext(ctx)
if rc == nil {
return nil
}
s, _ := rc.Stats.GetExtension(complexityExtension).(*ComplexityStats)
return s
}