11import {
22 Attributes ,
3+ GatewayConfigContext ,
34 getRetryInfo ,
45 isRetryExecutionRequest ,
56 Logger ,
@@ -214,10 +215,31 @@ const HeadersTextMapGetter: TextMapGetter<Headers> = {
214215 } ,
215216} ;
216217
218+ export type ContextMatcher = {
219+ request ?: Request ;
220+ context ?: any ;
221+ executionRequest ?: ExecutionRequest ;
222+ } ;
223+
224+ export type OpenTelemetryPluginUtils = {
225+ tracer ?: Tracer ;
226+ getActiveContext : ( payload : ContextMatcher ) => Context ;
227+ getHttpContext : ( request : Request ) => Context | undefined ;
228+ getOperationContext : ( context : any ) => Context | undefined ;
229+ getExecutionRequestContext : (
230+ ExecutionRequest : ExecutionRequest ,
231+ ) => Context | undefined ;
232+ } ;
233+
217234export type OpenTelemetryContextExtension = {
218235 openTelemetry : {
219236 tracer : Tracer ;
220- activeContext : ( ) => Context ;
237+ getActiveContext : ( payload ?: ContextMatcher ) => Context ;
238+ getHttpContext : ( request ?: Request ) => Context | undefined ;
239+ getOperationContext : ( context ?: any ) => Context | undefined ;
240+ getExecutionRequestContext : (
241+ ExecutionRequest : ExecutionRequest ,
242+ ) => Context | undefined ;
221243 } ;
222244} ;
223245
@@ -229,30 +251,26 @@ type State = Partial<
229251 HttpState < OtelState > & GraphQLState < OtelState > & GatewayState < OtelState >
230252> ;
231253
232- export type OpenTelemetryPlugin =
233- GatewayPlugin < OpenTelemetryContextExtension > & {
234- getOtelContext : ( payload : {
235- request ?: Request ;
236- context ?: any ;
237- executionRequest ?: ExecutionRequest ;
238- } ) => Context ;
239- getTracer ( ) : Tracer ;
240- } ;
254+ export type OpenTelemetryPlugin = GatewayPlugin < OpenTelemetryContextExtension > &
255+ OpenTelemetryPluginUtils ;
241256
242257export function useOpenTelemetry (
243- options : OpenTelemetryGatewayPluginOptions & {
244- log : Logger ;
245- } ,
258+ options : OpenTelemetryGatewayPluginOptions &
259+ // We ask for a Partial context to still allow the usage as a Yoga plugin
260+ Partial < GatewayConfigContext > ,
246261) : OpenTelemetryPlugin {
247262 const inheritContext = options . inheritContext ?? true ;
248263 const propagateContext = options . propagateContext ?? true ;
249264 let useContextManager : boolean ;
250265 const traces = typeof options . traces === 'object' ? options . traces : { } ;
251266
252267 let tracer : Tracer ;
253- let pluginLogger : Logger ;
254268 let initSpan : Context | null ;
255269
270+ // TODO: Make it const once Yoga has the Hive Logger
271+ let pluginLogger : Logger | undefined =
272+ options . log && options . log . child ( '[OpenTelemetry] ' ) ;
273+
256274 function isParentEnabled ( state : State ) : boolean {
257275 const parentState = getMostSpecificState ( state ) ;
258276 return ! parentState || ! ! parentState . otel ;
@@ -304,7 +322,7 @@ export function useOpenTelemetry(
304322
305323 if ( ! useContextManager ) {
306324 if ( traces . spans ?. schema ) {
307- pluginLogger . warn (
325+ pluginLogger ? .warn (
308326 'Schema loading spans are disabled because no context manager is available' ,
309327 ) ;
310328 }
@@ -314,14 +332,25 @@ export function useOpenTelemetry(
314332 }
315333 }
316334
317- return withState <
335+ const plugin = withState <
318336 OpenTelemetryPlugin ,
319337 OtelState ,
320338 OtelState & { skipExecuteSpan ?: true ; subgraphNames : string [ ] } ,
321339 OtelState
322340 > ( ( getState ) => ( {
323- getTracer : ( ) => tracer ,
324- getOtelContext : ( { state } ) => getContext ( state ) ,
341+ get tracer ( ) {
342+ return tracer ;
343+ } ,
344+ getActiveContext : ( { state } ) => getContext ( state ) ,
345+ getHttpContext : ( request ) => {
346+ return getState ( { request } ) . forRequest . otel ?. root ;
347+ } ,
348+ getOperationContext : ( context ) => {
349+ return getState ( { context } ) . forOperation . otel ?. root ;
350+ } ,
351+ getExecutionRequestContext : ( executionRequest ) => {
352+ return getState ( { executionRequest } ) . forSubgraphExecution . otel ?. root ;
353+ } ,
325354 instrumentation : {
326355 request ( { state : { forRequest } , request } , wrapped ) {
327356 if ( ! shouldTrace ( traces . spans ?. http , { request } ) ) {
@@ -648,21 +677,17 @@ export function useOpenTelemetry(
648677 } ,
649678
650679 onYogaInit ( { yoga } ) {
651- const log =
652- options . log ??
653- //TODO remove this when Yoga will also use the new Logger API
654- new Logger ( {
655- writers : [
656- {
657- write ( level , attrs , msg ) {
658- level = level === 'trace' ? 'debug' : level ;
659- yoga . logger [ level ] ( msg , attrs ) ;
660- } ,
680+ //TODO remove this when Yoga will also use the new Logger API
681+ pluginLogger ??= new Logger ( {
682+ writers : [
683+ {
684+ write ( level , attrs , msg ) {
685+ level = level === 'trace' ? 'debug' : level ;
686+ yoga . logger [ level ] ( msg , attrs ) ;
661687 } ,
662- ] ,
663- } ) ;
664-
665- pluginLogger = log . child ( '[OpenTelemetry] ' ) ;
688+ } ,
689+ ] ,
690+ } ) . child ( '[OpenTelemetry] ' ) ;
666691
667692 if ( options . configureDiagLogger !== false ) {
668693 const logLevel = diagLogLevelFromEnv ( ) ; // We enable the diag only if it is explicitly enabled, as NodeSDK does
@@ -685,7 +710,9 @@ export function useOpenTelemetry(
685710 try {
686711 const requestId = requestIdByRequest . get ( request ) ;
687712 if ( requestId ) {
688- loggerForRequest ( options . log . child ( { requestId } ) , request ) ;
713+ if ( options . log ) {
714+ loggerForRequest ( options . log . child ( { requestId } ) , request ) ;
715+ }
689716
690717 // When running in a runtime without a context manager, we have to keep track of the
691718 // span correlated to a log manually. For now, we just link all logs for a request to
@@ -695,7 +722,7 @@ export function useOpenTelemetry(
695722 }
696723 }
697724 } catch ( error ) {
698- pluginLogger . error (
725+ pluginLogger ! . error (
699726 { error } ,
700727 'Error while setting up logger for request' ,
701728 ) ;
@@ -706,7 +733,20 @@ export function useOpenTelemetry(
706733 extendContext ( {
707734 openTelemetry : {
708735 tracer,
709- activeContext : ( ) => getContext ( state ) ,
736+ getHttpContext : ( request ) => {
737+ const { forRequest } = request ? getState ( { request } ) : state ;
738+ return forRequest . otel ?. root ;
739+ } ,
740+ getOperationContext : ( context ) => {
741+ const { forOperation } = context ? getState ( { context } ) : state ;
742+ return forOperation . otel ?. root ;
743+ } ,
744+ getExecutionRequestContext : ( executionRequest ) => {
745+ return getState ( { executionRequest } ) . forSubgraphExecution . otel
746+ ?. root ;
747+ } ,
748+ getActiveContext : ( contextMatcher ?: Parameters < typeof getState > [ 0 ] ) =>
749+ getContext ( contextMatcher ? getState ( contextMatcher ) : state ) ,
710750 } ,
711751 } ) ;
712752 } ,
@@ -903,6 +943,16 @@ export function useOpenTelemetry(
903943 }
904944 } ,
905945 } ) ) ;
946+
947+ if ( options . openTelemetry ) {
948+ if ( options . openTelemetry . register ) {
949+ options . openTelemetry ?. register ?.( plugin ) ;
950+ } else {
951+ options . log ?. warn ( 'An OpenTelemetry plugin is already registered' ) ;
952+ }
953+ }
954+
955+ return plugin ;
906956}
907957
908958function shouldTrace < Args > (
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?v=4&size=40){kind=avatar}
0 commit comments