github.com/senomas/gqlgen@v0.17.11-0.20220626120754-9aee61b0716a/graphql/handler/apollofederatedtracingv1/tracing.go (about)

     1  package apollofederatedtracingv1
     2  
     3  import (
     4  	"context"
     5  	"encoding/base64"
     6  	"fmt"
     7  
     8  	"github.com/99designs/gqlgen/graphql"
     9  	"github.com/vektah/gqlparser/v2/gqlerror"
    10  	"google.golang.org/protobuf/proto"
    11  )
    12  
    13  type (
    14  	Tracer struct {
    15  		ClientName  string
    16  		Version     string
    17  		Hostname    string
    18  		TreeBuilder *TreeBuilder
    19  		ShouldTrace bool
    20  	}
    21  )
    22  
    23  var _ interface {
    24  	graphql.HandlerExtension
    25  	graphql.ResponseInterceptor
    26  	graphql.FieldInterceptor
    27  	graphql.OperationInterceptor
    28  	graphql.OperationParameterMutator
    29  } = &Tracer{}
    30  
    31  // ExtensionName returns the name of the extension
    32  func (Tracer) ExtensionName() string {
    33  	return "ApolloFederatedTracingV1"
    34  }
    35  
    36  // Validate returns errors based on the schema; since this extension doesn't require validation, we return nil
    37  func (Tracer) Validate(graphql.ExecutableSchema) error {
    38  	return nil
    39  }
    40  
    41  func (t *Tracer) MutateOperationParameters(ctx context.Context, request *graphql.RawParams) *gqlerror.Error {
    42  	t.ShouldTrace = request.Headers.Get("apollo-federation-include-trace") == "ftv1" // check for header
    43  	return nil
    44  }
    45  
    46  // InterceptOperation acts on each Graph operation; on each operation, start a tree builder and start the tree's timer for tracing
    47  func (t *Tracer) InterceptOperation(ctx context.Context, next graphql.OperationHandler) graphql.ResponseHandler {
    48  	if !t.ShouldTrace {
    49  		return next(ctx)
    50  	}
    51  
    52  	t.TreeBuilder = NewTreeBuilder()
    53  
    54  	return next(ctx)
    55  }
    56  
    57  // InterceptField is called on each field's resolution, including information about the path and parent node.
    58  // This information is then used to build the relevant Node Tree used in the FTV1 tracing format
    59  func (t *Tracer) InterceptField(ctx context.Context, next graphql.Resolver) (interface{}, error) {
    60  	if !t.ShouldTrace {
    61  		return next(ctx)
    62  	}
    63  
    64  	t.TreeBuilder.WillResolveField(ctx)
    65  
    66  	return next(ctx)
    67  }
    68  
    69  // InterceptResponse is called before the overall response is sent, but before each field resolves; as a result
    70  // the final marshaling is deferred to happen at the end of the operation
    71  func (t *Tracer) InterceptResponse(ctx context.Context, next graphql.ResponseHandler) *graphql.Response {
    72  	if !t.ShouldTrace {
    73  		return next(ctx)
    74  	}
    75  	t.TreeBuilder.StartTimer(ctx)
    76  
    77  	// because we need to update the ftv1 string at a later time (as fields resolve before the response is sent),
    78  	// we instantiate the string and use a pointer to be able to update later
    79  	var ftv1 string
    80  	graphql.RegisterExtension(ctx, "ftv1", &ftv1)
    81  
    82  	// now that fields have finished resolving, it stops the timer to calculate trace duration
    83  	defer func() {
    84  		t.TreeBuilder.StopTimer(ctx)
    85  
    86  		// marshal the protobuf ...
    87  		p, err := proto.Marshal(t.TreeBuilder.Trace)
    88  		if err != nil {
    89  			fmt.Print(err)
    90  		}
    91  
    92  		// ... then set the previously instantiated string as the base64 formatted string as required
    93  		ftv1 = base64.StdEncoding.EncodeToString(p)
    94  	}()
    95  
    96  	resp := next(ctx)
    97  	return resp
    98  }