sorare/graphql/client.go

package graphql

import (
	"context"
	"net/http"
	"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) MaxComplexity() int {
	if c.authenticated {
		return MaxAuthenticatedQueryComplexity
	} else {
		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...)
}