package graphql import ( "context" "net/http" "reflect" "sync" "time" "github.com/llehouerou/go-graphql-client" "github.com/pkg/errors" "golang.org/x/time/rate" ) const ( rateLimitPeriod = time.Second rateLimitBurst = 60 MaxAuthenticatedQueryComplexity = 30000 MaxAnonymousQueryComplexity = 500 ) type Client struct { httpClient *http.Client gql *graphql.Client rl *rate.Limiter apirl *rate.Limiter lock *sync.Mutex authenticated bool token JwtToken apiKey string } func NewClient(httpclient *http.Client, baseUrl string) *Client { return &Client{ httpClient: httpclient, gql: graphql.NewClient(baseUrl, httpclient), rl: rate.NewLimiter(rate.Every(rateLimitPeriod), rateLimitBurst), apirl: rate.NewLimiter(rate.Every(rateLimitPeriod/10), rateLimitBurst*10), lock: &sync.Mutex{}, authenticated: false, } } func (c *Client) GetCurrentToken() JwtToken { return c.token } func (c *Client) SetJWTToken(token JwtToken, audience string) { c.gql = c.gql.WithRequestModifier(func(request *http.Request) { request.Header.Set("Authorization", "Bearer "+token.Token) request.Header.Set("JWT-AUD", audience) }) c.token = token c.authenticated = true } func (c *Client) SetApiKey(apiKey string) { c.gql = c.gql.WithRequestModifier(func(request *http.Request) { request.Header.Set("APIKEY", apiKey) }) c.apiKey = apiKey } func (c *Client) MaxQueryComplexity() int { if c.apiKey != "" || c.authenticated { return MaxAuthenticatedQueryComplexity } return MaxAnonymousQueryComplexity } func (c *Client) ConstructRawQuery( q interface{}, variables map[string]interface{}, ) (string, error) { return graphql.ConstructQuery(q, variables) } func (c *Client) Query( ctx context.Context, q interface{}, variables interface{}, options ...graphql.Option, ) error { if c.apiKey != "" { err := c.apirl.Wait(ctx) if err != nil { return errors.Wrap(err, "waiting for rate limit") } } else { err := c.rl.Wait(ctx) if err != nil { return errors.Wrap(err, "waiting for rate limit") } } c.lock.Lock() defer c.lock.Unlock() return c.gql.Query(ctx, q, variables, options...) } func (c *Client) QueryRaw( ctx context.Context, q interface{}, variables interface{}, options ...graphql.Option, ) ([]byte, error) { if c.apiKey != "" { err := c.apirl.Wait(ctx) if err != nil { return nil, errors.Wrap(err, "waiting for rate limit") } } else { err := c.rl.Wait(ctx) if err != nil { return nil, errors.Wrap(err, "waiting for rate limit") } } c.lock.Lock() defer c.lock.Unlock() return c.gql.QueryRaw(ctx, q, variables, options...) } func (c *Client) Mutate( ctx context.Context, q interface{}, variables interface{}, options ...graphql.Option, ) error { if c.apiKey != "" { err := c.apirl.Wait(ctx) if err != nil { return errors.Wrap(err, "waiting for rate limit") } } else { err := c.rl.Wait(ctx) if err != nil { return errors.Wrap(err, "waiting for rate limit") } } c.lock.Lock() defer c.lock.Unlock() return c.gql.Mutate(ctx, q, variables, options...) } func GetTypeComplexity(t reflect.Type) int { var checkStruct func(t reflect.Type, complexity *int) checkStruct = func(t reflect.Type, complexity *int) { if t.Kind() != reflect.Struct { return } for i := 0; i < t.NumField(); i++ { field := t.Field(i) tag := field.Tag.Get("graphql") if field.Type.Kind() == reflect.Struct { *complexity++ if len(tag) > 6 && tag[:6] == "... on" { *complexity-- } checkStruct(field.Type, complexity) continue } if field.Type.Kind() == reflect.Slice { if isSimpleType(field.Type.Elem()) { *complexity++ } else { *complexity++ tmpcomplexity := 0 checkStruct(field.Type.Elem(), &tmpcomplexity) *complexity += tmpcomplexity * 10 } continue } if tag != "" && tag != "__typename" { *complexity++ } } } complexity := 0 checkStruct(t, &complexity) return complexity } func isSimpleType(t reflect.Type) bool { switch t.Kind() { case reflect.Bool, reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr, reflect.Float32, reflect.Float64, reflect.Complex64, reflect.Complex128, reflect.String: return true default: return false } }