gqlgen-contrib/extensions/hive/reporter.go

package extensions

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/99designs/gqlgen/graphql"
	"github.com/google/uuid"
	"github.com/vektah/gqlparser/v2/ast"
	"go.uber.org/zap"
)

type Report struct {
	Size       int                           `json:"size"`
	Map        map[string]OperationMapRecord `json:"map"`
	Operations []Operation                   `json:"operations"`
}

type OperationMapRecord struct {
	Operation     string   `json:"operation"`
	OperationName string   `json:"operationName,omitempty"`
	Fields        []string `json:"fields"`
}

type Operation struct {
	OperationMapKey string    `json:"operationMapKey"`
	Timestamp       int64     `json:"timestamp"`
	Execution       Execution `json:"execution"`
	Metadata        *Metadata `json:"metadata"`
}

type Execution struct {
	Ok          bool  `json:"ok"`
	Duration    int64 `json:"duration"`
	ErrorsTotal int64 `json:"errorsTotal"`
}

type Metadata struct {
	Client Client `json:"client"`
}

type Client struct {
	Name    string `json:"name,omitempty"`
	Version string `json:"version,omitempty"`
}

type Reporter struct {
	Endpoint  string
	ApiToken  string
	BatchSize int
	report    *Report
	logger    *zap.Logger
}

var _ interface {
	graphql.HandlerExtension
	graphql.ResponseInterceptor
} = Reporter{}

func (Reporter) ExtensionName() string {
	return "HiveReporter"
}

func (Reporter) Validate(graphql.ExecutableSchema) error {
	return nil
}

func (r Reporter) InterceptResponse(ctx context.Context, next graphql.ResponseHandler) *graphql.Response {
	if !graphql.HasOperationContext(ctx) {
		return next(ctx)
	}

	rc := graphql.GetOperationContext(ctx)

	start := rc.Stats.OperationStart

	resp := next(ctx)

	if rc.Operation == nil {
		return resp
	}

	fields := getPreloads(ctx, rc.Operation.SelectionSet)

	var mapKey string

	for key, omr := range r.report.Map {
		if omr.Operation != rc.RawQuery {
			continue
		}

		if omr.OperationName != rc.OperationName {
			continue
		}

		if len(omr.Fields) != len(fields) {
			continue
		}

		found := false

		for _, f := range omr.Fields {
			for _, _f := range fields {
				if f == _f {
					found = true
				}
			}
		}

		if !found {
			continue
		}

		mapKey = key

		break
	}

	if mapKey == "" {
		id := uuid.New()

		opr := OperationMapRecord{
			Operation:     rc.RawQuery,
			OperationName: rc.OperationName,
			Fields:        fields,
		}

		r.report.Map[id.String()] = opr
		mapKey = id.String()
	}

	end := graphql.Now()
	duration := end.Sub(start)

	var (
		clientName    = rc.Headers.Get("User-Agent")
		clientVersion string
	)

	if parts := strings.Split(clientName, "/"); len(parts) > 1 {
		clientName, clientVersion = parts[0], strings.Split(parts[1], " ")[0]
	}

	op := Operation{
		OperationMapKey: mapKey,
		Timestamp:       start.UnixMilli(),
		Execution: Execution{
			Ok:          len(resp.Errors) == 0,
			Duration:    duration.Nanoseconds(),
			ErrorsTotal: int64(len(resp.Errors)),
		},
		Metadata: &Metadata{
			Client: Client{
				Name:    clientName,
				Version: clientVersion,
			},
		},
	}

	r.report.Operations[r.report.Size] = op
	r.report.Size++

	if r.report.Size >= r.BatchSize {
		r.logger.Info("Sending report to Hive", zap.String("Endpoint", r.Endpoint), zap.Int("Operations", len(r.report.Operations)))
		jsonData, err := json.Marshal(r.report)

		if err != nil {
			r.logger.Error("Error marshaling report", zap.Error(err))
			graphql.AddError(ctx, err)
		}

		request, err := http.NewRequest(http.MethodPost, r.Endpoint, bytes.NewBuffer(jsonData))
		if err != nil {
			r.logger.Error("Error sending report to Hive", zap.Error(err))
			graphql.AddError(ctx, err)
		}

		request.Header.Set("Content-Type", "application/json; charset=UTF-8")
		request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", r.ApiToken))

		client := &http.Client{}

		response, err := client.Do(request)
		if err != nil {
			r.logger.Error("Error sending report to Hive", zap.Error(err))
			graphql.AddError(ctx, err)
		}

		r.logger.Info("Sent report to Hive", zap.Any("statusCode", response.StatusCode))

		defer response.Body.Close()

		*r.report = Report{
			Map:        map[string]OperationMapRecord{},
			Operations: make([]Operation, r.BatchSize),
		}
	}

	return resp
}

func getPreloads(ctx context.Context, selectionSet ast.SelectionSet) []string {
	return getNestedPreloads(
		graphql.GetOperationContext(ctx),
		graphql.CollectFields(graphql.GetOperationContext(ctx), selectionSet, nil),
		"",
	)
}

func getNestedPreloads(ctx *graphql.OperationContext, fields []graphql.CollectedField, prefix string) (preloads []string) {
	for _, column := range fields {
		prefixColumn := getPreloadString(prefix, column.Name)
		preloads = append(preloads, prefixColumn)
		preloads = append(preloads, getNestedPreloads(ctx, graphql.CollectFields(ctx, column.Selections, nil), prefixColumn)...)
	}

	return
}

func getPreloadString(prefix, name string) string {
	if len(prefix) > 0 {
		return prefix + "." + name
	}

	return name
}

func NewReporter(endpoint, apiToken string, batchSize int, logger *zap.Logger) *Reporter {
	return &Reporter{
		Endpoint:  endpoint,
		ApiToken:  apiToken,
		BatchSize: batchSize,
		report: &Report{
			Map:        map[string]OperationMapRecord{},
			Operations: make([]Operation, batchSize),
		},
		logger: logger,
	}
}