44from remoulade import Middleware , broker
55
66from opentelemetry import trace
7+ from opentelemetry .instrumentation .instrumentor import BaseInstrumentor
78from opentelemetry .instrumentation .remoulade import utils
8- from opentelemetry .instrumentation .remoulade .version import __version__
99from opentelemetry .instrumentation .remoulade .package import _instruments
10- from opentelemetry .instrumentation .instrumentor import BaseInstrumentor
10+ from opentelemetry .instrumentation .remoulade . version import __version__
1111from opentelemetry .propagate import extract , inject
1212from opentelemetry .semconv .trace import SpanAttributes
1313
14-
1514_MESSAGE_TAG_KEY = "remoulade.action"
1615_MESSAGE_SEND = "send"
1716_MESSAGE_RUN = "run"
@@ -28,20 +27,32 @@ def before_process_message(self, _broker, message):
2827 trace_ctx = extract (message .options ["trace_ctx" ])
2928 retry_count = message .options .get ("retries" )
3029
31- operation_name = "remoulade/process" if retry_count is None else f"remoulade/process(retry-{ retry_count } )"
30+ operation_name = (
31+ "remoulade/process"
32+ if retry_count is None
33+ else f"remoulade/process(retry-{ retry_count } )"
34+ )
3235
33- span = self ._tracer .start_span (operation_name , kind = trace .SpanKind .CONSUMER , context = trace_ctx )
36+ span = self ._tracer .start_span (
37+ operation_name , kind = trace .SpanKind .CONSUMER , context = trace_ctx
38+ )
3439
3540 if retry_count is not None :
3641 span .set_attribute ("retry_count" , retry_count )
3742
3843 activation = trace .use_span (span , end_on_exit = True )
3944 activation .__enter__ ()
4045
41- utils .attach_span (self ._span_registry , message .message_id , (span , activation ))
46+ utils .attach_span (
47+ self ._span_registry , message .message_id , (span , activation )
48+ )
4249
43- def after_process_message (self , _broker , message , * , result = None , exception = None ):
44- span , activation = utils .retrieve_span (self ._span_registry , message .message_id )
50+ def after_process_message (
51+ self , _broker , message , * , result = None , exception = None
52+ ):
53+ span , activation = utils .retrieve_span (
54+ self ._span_registry , message .message_id
55+ )
4556
4657 if span is None :
4758 # no existing span found for message_id
@@ -58,37 +69,54 @@ def after_process_message(self, _broker, message, *, result=None, exception=None
5869 def before_enqueue (self , _broker , message , delay ):
5970 retry_count = message .options .get ("retries" )
6071
61- operation_name = "remoulade/send" if retry_count is None else f"remoulade/send(retry-{ retry_count } )"
72+ operation_name = (
73+ "remoulade/send"
74+ if retry_count is None
75+ else f"remoulade/send(retry-{ retry_count } )"
76+ )
6277
63- span = self ._tracer .start_span (operation_name , kind = trace .SpanKind .PRODUCER )
78+ span = self ._tracer .start_span (
79+ operation_name , kind = trace .SpanKind .PRODUCER
80+ )
6481
6582 if retry_count is not None :
6683 span .set_attribute ("retry_count" , retry_count )
6784
6885 if span .is_recording ():
6986 span .set_attribute (_MESSAGE_TAG_KEY , _MESSAGE_SEND )
70- span .set_attribute (SpanAttributes .MESSAGING_MESSAGE_ID , message .message_id )
87+ span .set_attribute (
88+ SpanAttributes .MESSAGING_MESSAGE_ID , message .message_id
89+ )
7190 span .set_attribute (_MESSAGE_NAME_KEY , message .actor_name )
7291 pass
7392
7493 activation = trace .use_span (span , end_on_exit = True )
7594 activation .__enter__ ()
7695
77- utils .attach_span (self ._span_registry , message .message_id , (span , activation ), is_publish = True )
96+ utils .attach_span (
97+ self ._span_registry ,
98+ message .message_id ,
99+ (span , activation ),
100+ is_publish = True ,
101+ )
78102
79103 if "trace_ctx" not in message .options :
80104 message .options ["trace_ctx" ] = {}
81105 inject (message .options ["trace_ctx" ])
82106
83107 def after_enqueue (self , _broker , message , delay , exception = None ):
84- _ , activation = utils .retrieve_span (self ._span_registry , message .message_id , is_publish = True )
108+ _ , activation = utils .retrieve_span (
109+ self ._span_registry , message .message_id , is_publish = True
110+ )
85111
86112 if activation is None :
87113 # no existing span found for message_id
88114 return
89115
90116 activation .__exit__ (None , None , None )
91- utils .detach_span (self ._span_registry , message .message_id , is_publish = True )
117+ utils .detach_span (
118+ self ._span_registry , message .message_id , is_publish = True
119+ )
92120
93121
94122class RemouladeInstrumentor (BaseInstrumentor ):
0 commit comments