nucliadb-telemetry 6.4.0.post4171__py3-none-any.whl → 6.4.0.post4184__py3-none-any.whl

nucliadb_telemetry/jetstream.py

@@ -72,7 +72,7 @@ def start_span_message_receiver(tracer: Tracer, msg: Msg):
     }
 
     # add some attributes from the metadata
-    ctx = extract(msg.headers)
+    ctx = extract(msg.headers or {})
 
     span = tracer.start_as_current_span(
         name=f"Received from {msg.subject}",
@@ -98,6 +98,24 @@ def start_span_message_publisher(tracer: Tracer, subject: str):
     return span
 
 
+async def _traced_callback(origin_cb: Callable[[Msg], Awaitable[None]], tracer: Tracer, msg: Msg):
+    with start_span_message_receiver(tracer, msg) as span:
+        try:
+            await origin_cb(msg)
+        except Exception as error:
+            set_span_exception(span, error)
+            raise error
+        finally:
+            msg_consume_time_histo.observe(
+                (datetime.now() - msg.metadata.timestamp).total_seconds(),
+                {
+                    "stream": msg.metadata.stream,
+                    "consumer": msg.metadata.consumer or "",
+                    "acked": "yes" if msg._ackd else "no",
+                },
+            )
+
+
 class JetStreamContextTelemetry:
     def __init__(self, js: JetStreamContext, service_name: str, tracer_provider: TracerProvider):
         self.js = js
@@ -118,31 +136,11 @@ class JetStreamContextTelemetry:
         **kwargs,
     ):
         tracer = self.tracer_provider.get_tracer(f"{self.service_name}_js_subscriber")
-
-        async def wrapper(origin_cb, tracer, msg: Msg):
-            # Execute the callback without tracing
-            if msg.headers is None:
-                logger.debug("Message received without headers, skipping span")
-                await origin_cb(msg)
-                return
-
-            with start_span_message_receiver(tracer, msg) as span:
-                try:
-                    await origin_cb(msg)
-                except Exception as error:
-                    set_span_exception(span, error)
-                    raise error
-                finally:
-                    msg_consume_time_histo.observe(
-                        (datetime.now() - msg.metadata.timestamp).total_seconds(),
-                        {
-                            "stream": msg.metadata.stream,
-                            "consumer": msg.metadata.consumer or "",
-                            "acked": "yes" if msg._ackd else "no",  # type: ignore
-                        },
-                    )
-
-        wrapped_cb = partial(wrapper, cb, tracer)
+        wrapped_cb: Optional[Callable[[Msg], Awaitable[None]]]
+        if cb is not None:
+            wrapped_cb = partial(_traced_callback, cb, tracer)
+        else:
+            wrapped_cb = cb
         return await self.js.subscribe(subject, queue=queue, cb=wrapped_cb, **kwargs)
 
     async def publish(
@@ -154,8 +152,8 @@ class JetStreamContextTelemetry:
     ):
         tracer = self.tracer_provider.get_tracer(f"{self.service_name}_js_publisher")
         headers = {} if headers is None else headers
-        inject(headers)
         with start_span_message_publisher(tracer, subject) as span:
+            inject(headers)
             try:
                 result = await self.js.publish(subject, body, headers=headers, **kwargs)
                 msg_sent_counter.inc({"subject": subject, "status": metrics.OK})
@@ -166,6 +164,10 @@ class JetStreamContextTelemetry:
 
         return result
 
+    async def trace_pull_subscriber_message(self, cb: Callable[[Msg], Awaitable[None]], msg: Msg):
+        tracer = self.tracer_provider.get_tracer(f"{self.service_name}_js_pull_subscriber")
+        await _traced_callback(cb, tracer, msg)
+
     # Just for convenience, to wrap all we use in the context of
     # telemetry-instrumented stuff using the JetStreamContextTelemetry class
 
@@ -259,9 +261,9 @@ class NatsClientTelemetry:
     ) -> None:
         tracer = self.tracer_provider.get_tracer(f"{self.service_name}_nc_publisher")
         headers = {} if headers is None else headers
-        inject(headers)
 
         with start_span_message_publisher(tracer, subject) as span:
+            inject(headers)
             try:
                 await self.nc.publish(subject, body, reply, headers)
                 msg_sent_counter.inc({"subject": subject, "status": metrics.OK})
@@ -279,9 +281,9 @@ class NatsClientTelemetry:
         headers: Optional[Dict[str, Any]] = None,
     ) -> Msg:
         headers = {} if headers is None else headers
-        inject(headers)
         tracer = self.tracer_provider.get_tracer(f"{self.service_name}_nc_request")
         with start_span_message_publisher(tracer, subject) as span:
+            inject(headers)
             try:
                 result = await self.nc.request(subject, payload, timeout, old_style, headers)
             except Exception as error:

{nucliadb_telemetry-6.4.0.post4171.dist-info → nucliadb_telemetry-6.4.0.post4184.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb_telemetry
-Version: 6.4.0.post4171
+Version: 6.4.0.post4184
 Summary: NucliaDB Telemetry Library Python process
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL

{nucliadb_telemetry-6.4.0.post4171.dist-info → nucliadb_telemetry-6.4.0.post4184.dist-info}/RECORD

@@ -7,7 +7,7 @@ nucliadb_telemetry/grpc.py,sha256=cNJRJLr0aRedKr5V_ufGqeidnMlBouML3QbFKkugTqg,15
 nucliadb_telemetry/grpc_metrics.py,sha256=L2XMcJDxioNUYPX_IbKl_zvrG8uOFpF_a67c6yJq-Jw,5174
 nucliadb_telemetry/grpc_sentry.py,sha256=3eS_DVIX5MzvbJBkykf8JyZUf1b_b9w1KLhjoiW7oV0,2659
 nucliadb_telemetry/jaeger.py,sha256=HMQK8x_oaol5tCx3hCjaMG5Uzvsy_WWhKnD-FEZoYA4,7036
-nucliadb_telemetry/jetstream.py,sha256=w6CLN_oGmQT0FgfMmQdtrASsEw5mCk3KW8TwaWCEl-E,11971
+nucliadb_telemetry/jetstream.py,sha256=hhlAv0sEottFjRp16vjwf2aWQqsoGVttVgwjzh1jNVs,12053
 nucliadb_telemetry/logs.py,sha256=ag1e7tKEoNFgB1FuAmyTxdIHtTFwOJlGxPl2s1bhUic,9311
 nucliadb_telemetry/metrics.py,sha256=euexqkmlZip6QFbMfvJW1fSuha96J4ZHnMWixeVJnAs,8652
 nucliadb_telemetry/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -30,7 +30,7 @@ nucliadb_telemetry/tests/grpc/helloworld_pb2.py,sha256=ht4dmi0pAy6qDrwcjkbtSf_hO
 nucliadb_telemetry/tests/grpc/helloworld_pb2.pyi,sha256=mBZCQE6z7riQw6fOegJgZ5lHmXqO23_zxL8iQMqSqms,1191
 nucliadb_telemetry/tests/grpc/helloworld_pb2_grpc.py,sha256=_jxUNxl4Fx-JztK9RO5R6osjNP2sVNVPAxLnmczEYOc,2677
 nucliadb_telemetry/tests/grpc/helloworld_pb2_grpc.pyi,sha256=Y6teCx-PhPU-rI6w5ItLBKaTb34FLpngPnuDVWtNve4,958
-nucliadb_telemetry-6.4.0.post4171.dist-info/METADATA,sha256=K1Paf3JNS3nY907ZhPn1BZ3EiVDURjsEdLmyNeoSVFk,11057
-nucliadb_telemetry-6.4.0.post4171.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
-nucliadb_telemetry-6.4.0.post4171.dist-info/top_level.txt,sha256=3qEHI_5ttqQIL2gkNYwSlKsFyBBiEzDiIy9UKISzOaQ,19
-nucliadb_telemetry-6.4.0.post4171.dist-info/RECORD,,
+nucliadb_telemetry-6.4.0.post4184.dist-info/METADATA,sha256=cM_ssIHA_F4TjqTkd9FqeM_VkcpYfqx5VJ9XtsunbyU,11057
+nucliadb_telemetry-6.4.0.post4184.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
+nucliadb_telemetry-6.4.0.post4184.dist-info/top_level.txt,sha256=3qEHI_5ttqQIL2gkNYwSlKsFyBBiEzDiIy9UKISzOaQ,19
+nucliadb_telemetry-6.4.0.post4184.dist-info/RECORD,,