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

nucliadb_telemetry/grpc.py

@@ -31,7 +31,7 @@ from opentelemetry.propagate import extract, inject
 from opentelemetry.propagators.textmap import CarrierT, Setter
 from opentelemetry.sdk.trace import Span, TracerProvider
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace import SpanKind, Tracer
+from opentelemetry.trace import SpanKind, Tracer, set_span_in_context
 from opentelemetry.trace.status import Status, StatusCode
 
 from nucliadb_telemetry import grpc_metrics, logger
@@ -52,7 +52,11 @@ _carrier_setter = _CarrierSetter()
 
 def finish_span_grpc(span: Span, result):
     code = result._cython_call._status.code()
-    if code != grpc.StatusCode.OK:
+
+    # grpc.StatusCode value is a tuple like:
+    #    <StatusCode.OK: (0, 'ok')>
+    # so we cannot compare it with the code we get from the result directly
+    if code == grpc.StatusCode.OK.value[0]:  # type: ignore
         span.set_status(
             Status(
                 status_code=StatusCode.OK,
@@ -87,19 +91,25 @@ def start_span_client(
         SpanAttributes.RPC_SERVICE: service,
     }
 
-    # add some attributes from the metadata
-    if client_call_details.metadata is not None:
-        mutable_metadata = OrderedDict(tuple(client_call_details.metadata))
-        inject(mutable_metadata, setter=_carrier_setter)  # type: ignore
-        for key, value in mutable_metadata.items():
-            client_call_details.metadata.add(key=key, value=value)  # type: ignore
-
     span = tracer.start_span(
         name=method_name,
         kind=SpanKind.CLIENT,
         attributes=attributes,
         set_status_on_exception=set_status_on_exception,
     )
+
+    # Create a context containing the new span
+    span_context = set_span_in_context(span)
+
+    if client_call_details.metadata is not None:
+        mutable_metadata = OrderedDict(tuple(client_call_details.metadata))
+    else:
+        mutable_metadata = OrderedDict()
+    inject(mutable_metadata, context=span_context, setter=_carrier_setter)
+    if client_call_details.metadata is not None:
+        for key, value in mutable_metadata.items():
+            client_call_details.metadata.add(key=key, value=value)  # type: ignore
+
     return span
 
 

nucliadb_telemetry/jetstream.py

@@ -19,11 +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
@@ -32,6 +34,7 @@ from opentelemetry.trace import SpanKind, Tracer
 
 from nucliadb_telemetry import logger, metrics
 from nucliadb_telemetry.common import set_span_exception
+from nucliadb_telemetry.utils import get_telemetry
 
 msg_consume_time_histo = metrics.Histogram(
     # time it takes from when msg was queue to when it finished processing
@@ -107,7 +110,13 @@ 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):
@@ -134,7 +143,7 @@ class JetStreamContextTelemetry:
                     )
 
         wrapped_cb = partial(wrapper, cb, tracer)
-        return await self.js.subscribe(cb=wrapped_cb, **kwargs)
+        return await self.js.subscribe(subject, queue=queue, cb=wrapped_cb, **kwargs)
 
     async def publish(
         self,
@@ -215,7 +224,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):
@@ -233,22 +248,22 @@ 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,
         subject: str,
-        body: bytes,
+        body: bytes = b"",
+        reply: str = "",
         headers: Optional[Dict[str, str]] = None,
-        **kwargs,
-    ):
+    ) -> 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:
             try:
-                await self.nc.publish(subject, body, headers=headers, **kwargs)
+                await self.nc.publish(subject, body, reply, headers)
                 msg_sent_counter.inc({"subject": subject, "status": metrics.OK})
             except Exception as error:
                 set_span_exception(span, error)
@@ -274,3 +289,49 @@ class NatsClientTelemetry:
                 raise error
 
         return result
+
+    # Other methods we use but don't need telemetry
+
+    @property
+    def is_connected(self) -> bool:
+        return self.nc.is_connected
+
+    @property
+    def connected_url(self) -> Optional[ParseResult]:
+        return self.nc.connected_url
+
+    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()
+
+    async def flush(self, timeout: int = nats.aio.client.DEFAULT_FLUSH_TIMEOUT) -> None:
+        return await self.nc.flush(timeout)
+
+    async def close(self) -> None:
+        return await self.nc.close()
+
+
+def get_traced_nats_client(nc: Client, service_name: str) -> Union[Client, NatsClientTelemetry]:
+    tracer_provider = get_telemetry(service_name)
+    if tracer_provider is not None:
+        return NatsClientTelemetry(nc, service_name, tracer_provider)
+    else:
+        return nc
+
+
+def get_traced_jetstream(
+    nc: Union[Client, NatsClientTelemetry], service_name: str
+) -> Union[JetStreamContext, JetStreamContextTelemetry]:
+    jetstream = nc.jetstream()
+    tracer_provider = get_telemetry(service_name)
+
+    if tracer_provider is not None and jetstream is not None:  # pragma: no cover
+        logger.info(f"Configuring {service_name} jetstream with telemetry")
+        return JetStreamContextTelemetry(jetstream, service_name, tracer_provider)
+    else:
+        return jetstream

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

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

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

@@ -3,11 +3,11 @@ nucliadb_telemetry/batch_span.py,sha256=5cE7wt-H6wMXnJ5GUgjZyKvBWfaaZ-VWHS1m8d36
 nucliadb_telemetry/common.py,sha256=i5xxHlxeNZ1MPCoBvDDNkRDiYhvq6ghKwuaBnYGyhGU,1222
 nucliadb_telemetry/context.py,sha256=rhLs4NWGMjKzLjr6kTTjq-gaR0sAPJudkS-KjpW9LWc,2505
 nucliadb_telemetry/errors.py,sha256=QTKp-ErpsRb-xN3m6kMpq6tpOrOMQxsHVJAofKPo3Rk,4478
-nucliadb_telemetry/grpc.py,sha256=1CdueeJedxManSUQKwNOvqUuFiVhnMiZViXUJm81bIY,14806
+nucliadb_telemetry/grpc.py,sha256=cNJRJLr0aRedKr5V_ufGqeidnMlBouML3QbFKkugTqg,15162
 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=E2_SY5-8ljhZAY0n_q-qlvewUmi2kF2En3mcqyA26kY,9977
+nucliadb_telemetry/jetstream.py,sha256=w6CLN_oGmQT0FgfMmQdtrASsEw5mCk3KW8TwaWCEl-E,11971
 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.post4161.dist-info/METADATA,sha256=_hP9pNf3JHE4uNMboO7ty3cg8Qlafjxq1Nb5u-u8DdY,11057
-nucliadb_telemetry-6.4.0.post4161.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
-nucliadb_telemetry-6.4.0.post4161.dist-info/top_level.txt,sha256=3qEHI_5ttqQIL2gkNYwSlKsFyBBiEzDiIy9UKISzOaQ,19
-nucliadb_telemetry-6.4.0.post4161.dist-info/RECORD,,
+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.post4161.dist-info → nucliadb_telemetry-6.4.0.post4171.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.0.0)
+Generator: setuptools (80.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any