nucliadb-telemetry 6.4.0.post4167__py3-none-any.whl → 6.4.0.post4176__py3-none-any.whl

nucliadb_telemetry/jetstream.py

@@ -19,12 +19,13 @@
 
 from datetime import datetime
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Union
 from urllib.parse import ParseResult
 
 import nats
 from nats.aio.client import Client
 from nats.aio.msg import Msg
+from nats.aio.subscription import Subscription
 from nats.js.client import JetStreamContext
 from opentelemetry.propagate import extract, inject
 from opentelemetry.sdk.trace import TracerProvider
@@ -71,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}",
@@ -97,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
@@ -109,34 +128,20 @@ class JetStreamContextTelemetry:
     async def add_stream(self, name: str, subjects: List[str]):
         return await self.js.add_stream(name=name, subjects=subjects)
 
-    async def subscribe(self, cb, **kwargs):
+    async def subscribe(
+        self,
+        subject: str,
+        queue: Optional[str] = None,
+        cb: Optional[Callable[[Msg], Awaitable[None]]] = None,
+        **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)
-        return await self.js.subscribe(cb=wrapped_cb, **kwargs)
+        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(
         self,
@@ -159,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
 
@@ -217,7 +226,13 @@ class NatsClientTelemetry:
         self.service_name = service_name
         self.tracer_provider = tracer_provider
 
-    async def subscribe(self, cb, **kwargs):
+    async def subscribe(
+        self,
+        subject: str,
+        queue: str = "",
+        cb: Optional[Callable[[Msg], Awaitable[None]]] = None,
+        **kwargs,
+    ) -> Subscription:
         tracer = self.tracer_provider.get_tracer(f"{self.service_name}_nc_subscriber")
 
         async def wrapper(origin_cb, tracer, msg: Msg):
@@ -235,7 +250,7 @@ class NatsClientTelemetry:
                     raise error
 
         wrapped_cb = partial(wrapper, cb, tracer)
-        return await self.nc.subscribe(cb=wrapped_cb, **kwargs)
+        return await self.nc.subscribe(subject, queue=queue, cb=wrapped_cb, **kwargs)
 
     async def publish(
         self,
@@ -290,6 +305,9 @@ class NatsClientTelemetry:
     def jetstream(self, **opts) -> nats.js.JetStreamContext:
         return self.nc.jetstream(**opts)
 
+    def jsm(self, **opts) -> nats.js.JetStreamManager:
+        return self.nc.jsm(**opts)
+
     async def drain(self) -> None:
         return await self.nc.drain()
 

{nucliadb_telemetry-6.4.0.post4167.dist-info → nucliadb_telemetry-6.4.0.post4176.dist-info}/METADATA

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

{nucliadb_telemetry-6.4.0.post4167.dist-info → nucliadb_telemetry-6.4.0.post4176.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=UXQQdmnocgPGampxRe2Mv7t97qMrfGLO1hf-1pkEGa0,11492
+nucliadb_telemetry/jetstream.py,sha256=PAkteKgKI5zwRzvEn5Uv1sGHykug2wp4bp5H6vSd-PE,12041
 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.post4167.dist-info/METADATA,sha256=xK1l5D9NwFShKHDHAhr8m4Ts8JpNzNiGVQfVVzye2MY,11057
-nucliadb_telemetry-6.4.0.post4167.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
-nucliadb_telemetry-6.4.0.post4167.dist-info/top_level.txt,sha256=3qEHI_5ttqQIL2gkNYwSlKsFyBBiEzDiIy9UKISzOaQ,19
-nucliadb_telemetry-6.4.0.post4167.dist-info/RECORD,,
+nucliadb_telemetry-6.4.0.post4176.dist-info/METADATA,sha256=5NyvhROrahLGQUu7zDTvQ4CAPXRDTF3o78W6gxQBHn4,11057
+nucliadb_telemetry-6.4.0.post4176.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
+nucliadb_telemetry-6.4.0.post4176.dist-info/top_level.txt,sha256=3qEHI_5ttqQIL2gkNYwSlKsFyBBiEzDiIy9UKISzOaQ,19
+nucliadb_telemetry-6.4.0.post4176.dist-info/RECORD,,