sorare/graphql/client.go

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
	}
}