nucliadb-utils 6.4.0.post4158__py3-none-any.whl → 6.4.0.post4167__py3-none-any.whl

nucliadb_utils/audit/stream.py

@@ -58,9 +58,9 @@ from nucliadb_protos.kb_usage_pb2 import (
     ClientType as ClientTypeKbUsage,
 )
 from nucliadb_protos.resources_pb2 import FieldID
+from nucliadb_telemetry.jetstream import get_traced_jetstream, get_traced_nats_client
 from nucliadb_utils import logger
 from nucliadb_utils.audit.audit import AuditStorage
-from nucliadb_utils.nats import get_traced_jetstream
 from nucliadb_utils.nuclia_usage.utils.kb_usage_report import KbUsageReportUtility
 
 
@@ -207,7 +207,8 @@ class StreamAuditStorage(AuditStorage):
         if len(self.nats_servers) > 0:
             options["servers"] = self.nats_servers
 
-        self.nc = await nats.connect(**options)
+        nc = await nats.connect(**options)
+        self.nc = get_traced_nats_client(nc, self.service)
 
         self.js = get_traced_jetstream(self.nc, self.service)
         self.task = asyncio.create_task(self.run())

nucliadb_utils/cache/nats.py

@@ -34,6 +34,7 @@ from nats.aio.subscription import Subscription
 from nats.js.client import JetStreamContext
 from nats.js.manager import JetStreamManager
 
+from nucliadb_telemetry.jetstream import get_traced_nats_client
 from nucliadb_utils import logger
 from nucliadb_utils.cache.pubsub import Callback, PubSubDriver
 from nucliadb_utils.nats import get_traced_jetstream
@@ -91,7 +92,7 @@ class NatsPubsub(PubSubDriver):
         # No asyncio loop to run
 
         async with self.lock:
-            self.nc = Client()
+            nc = Client()
             options = {
                 "servers": self._hosts,
                 "disconnected_cb": self.disconnected_cb,
@@ -104,11 +105,12 @@ class NatsPubsub(PubSubDriver):
             if self.user_credentials_file is not None:
                 options["user_credentials"] = self.user_credentials_file
             try:
-                await self.nc.connect(**options)
+                await nc.connect(**options)
             except ErrNoServers:
                 logger.exception("No servers found")
                 raise
 
+            self.nc = get_traced_nats_client(nc, "nucliadb_pubsub")
             logger.info("Connected to nats")
 
         self.initialized = True

nucliadb_utils/nats.py

@@ -33,24 +33,18 @@ from nats.aio.subscription import Subscription
 from nats.js.client import JetStreamContext
 
 from nucliadb_telemetry.errors import capture_exception
-from nucliadb_telemetry.jetstream import JetStreamContextTelemetry
+from nucliadb_telemetry.jetstream import (
+    JetStreamContextTelemetry,
+    NatsClientTelemetry,
+    get_traced_nats_client,
+)
 from nucliadb_telemetry.metrics import Counter
-from nucliadb_telemetry.utils import get_telemetry
 
 logger = logging.getLogger(__name__)
 
-
-def get_traced_jetstream(
-    nc: NATSClient, 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
+# Re-export for bw/c. This function was defined here but makes more sense in the
+# telemetry library
+from nucliadb_telemetry.jetstream import get_traced_jetstream  # noqa
 
 
 class MessageProgressUpdater:
@@ -105,7 +99,7 @@ class MessageProgressUpdater:
 
 
 class NatsConnectionManager:
-    _nc: NATSClient
+    _nc: Union[NATSClient, NatsClientTelemetry]
     _subscriptions: list[tuple[Subscription, Callable[[], Awaitable[None]]]]
     _pull_subscriptions: list[
         tuple[
@@ -169,7 +163,8 @@ class NatsConnectionManager:
             options["servers"] = self._nats_servers
 
         async with self._lock:
-            self._nc = await nats.connect(**options)
+            nc = await nats.connect(**options)
+            self._nc = get_traced_nats_client(nc, self._service_name)
 
         self._expected_subscription_task = asyncio.create_task(self._verify_expected_subscriptions())
 
@@ -310,7 +305,7 @@ class NatsConnectionManager:
         logger.info("Connection is closed on NATS")
 
     @property
-    def nc(self) -> NATSClient:
+    def nc(self) -> Union[NATSClient, NatsClientTelemetry]:
         return self._nc
 
     @cached_property

nucliadb_utils/transaction.py

@@ -35,10 +35,14 @@ from nucliadb_protos.writer_pb2 import (
     Notification,
     OpStatusWriter,
 )
-from nucliadb_telemetry.jetstream import JetStreamContextTelemetry
+from nucliadb_telemetry.jetstream import (
+    JetStreamContextTelemetry,
+    NatsClientTelemetry,
+    get_traced_jetstream,
+    get_traced_nats_client,
+)
 from nucliadb_utils import const, logger
 from nucliadb_utils.cache.pubsub import PubSubDriver
-from nucliadb_utils.nats import get_traced_jetstream
 from nucliadb_utils.utilities import get_pubsub
 
 
@@ -95,7 +99,7 @@ class LocalTransactionUtility:
 
 
 class TransactionUtility:
-    nc: Client
+    nc: Union[Client, NatsClientTelemetry]
     js: Union[JetStreamContext, JetStreamContextTelemetry]
     pubsub: PubSubDriver
 
@@ -177,7 +181,8 @@ class TransactionUtility:
         if len(self.nats_servers) > 0:
             options["servers"] = self.nats_servers
 
-        self.nc = await nats.connect(**options)
+        nc = await nats.connect(**options)
+        self.nc = get_traced_nats_client(nc, service_name or "nucliadb")
         self.js = get_traced_jetstream(self.nc, service_name or "nucliadb")
 
     async def finalize(self):

{nucliadb_utils-6.4.0.post4158.dist-info → nucliadb_utils-6.4.0.post4167.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb_utils
-Version: 6.4.0.post4158
+Version: 6.4.0.post4167
 Summary: NucliaDB util library
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -27,8 +27,8 @@ Requires-Dist: nats-py[nkeys]>=2.6.0
 Requires-Dist: PyNaCl
 Requires-Dist: pyjwt>=2.4.0
 Requires-Dist: mrflagly>=0.2.9
-Requires-Dist: nucliadb-protos>=6.4.0.post4158
-Requires-Dist: nucliadb-telemetry>=6.4.0.post4158
+Requires-Dist: nucliadb-protos>=6.4.0.post4167
+Requires-Dist: nucliadb-telemetry>=6.4.0.post4167
 Provides-Extra: cache
 Requires-Dist: redis>=4.3.4; extra == "cache"
 Requires-Dist: orjson>=3.6.7; extra == "cache"

{nucliadb_utils-6.4.0.post4158.dist-info → nucliadb_utils-6.4.0.post4167.dist-info}/RECORD

@@ -7,14 +7,14 @@ nucliadb_utils/exceptions.py,sha256=y_3wk77WLVUtdo-5FtbBsdSkCtK_DsJkdWb5BoPn3qo,
 nucliadb_utils/featureflagging.py,sha256=SYqr31e0_NeOMx6dZgSiDO9VJiV6xxYhM_gWkk8ITTQ,2386
 nucliadb_utils/grpc.py,sha256=apu0uePnkGHCAT7GRQ9YZfRYyFj26kJ440i8jitbM3U,3314
 nucliadb_utils/helpers.py,sha256=nPw8yod3hP-pxq80VF8QC36s7ygSg0dBUdfI-LatvCs,1600
-nucliadb_utils/nats.py,sha256=LrUNBPffjK3e_EuF_s2MGmmE7ExOVD9cdi0yZjHNSv4,16167
+nucliadb_utils/nats.py,sha256=T5A_pH0koGPNH9IPNPv1LsYGnAS5hnJgN8VuNCqO_p8,15989
 nucliadb_utils/partition.py,sha256=jBgy4Hu5Iwn4gjbPPcthSykwf-qNx-GcLAIwbzPd1d0,1157
 nucliadb_utils/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb_utils/run.py,sha256=Es0_Bu5Yc-LWczvwL6gzWqSwC85RjDCk-0oFQAJi9g4,1827
 nucliadb_utils/settings.py,sha256=RnGhEUvwv6faNqALiqDCivvzNOyyXVBflYh_37uNkow,8193
 nucliadb_utils/signals.py,sha256=lo_Mk12NIX5Au--3H3WObvDOXq_OMurql2qiC2TnAao,2676
 nucliadb_utils/store.py,sha256=kQ35HemE0v4_Qg6xVqNIJi8vSFAYQtwI3rDtMsNy62Y,890
-nucliadb_utils/transaction.py,sha256=z_VeiTIta48rosS2SXMqx86XaavprSMRWf6s6zWIeEs,7920
+nucliadb_utils/transaction.py,sha256=K_eXwbbC_exc1DKqce-eIYkMEwbLISmR4dqVOb8geAo,8051
 nucliadb_utils/utilities.py,sha256=w4nsdyiFgZVEq5aBhjzVg176UDJNxfhIEfDbEi8bJdA,15604
 nucliadb_utils/aiopynecone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/aiopynecone/client.py,sha256=MPyHnDXwhukJr7U3CJh7BpsekfSuOkyM4g5b9LLtzc8,22941
@@ -23,10 +23,10 @@ nucliadb_utils/aiopynecone/models.py,sha256=XkNIZx4bxdbVo9zYVn8IRp70q4DWUMWN79yb
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/audit/audit.py,sha256=QNrohgWcWkO0exDOi-U7wp1YY7s8ezJqJESWGx4KnYM,3493
 nucliadb_utils/audit/basic.py,sha256=4NX3IkZtTM_F2jD9TGE8-1eoS-hc-_zWwFGYVahGNSs,4141
-nucliadb_utils/audit/stream.py,sha256=lyQXTPswlbGlQLvZsHorCKkMSiLk0_tXToZqPyjHbq0,17174
+nucliadb_utils/audit/stream.py,sha256=XEY4wRxEPdrwaBtepT2-4fmeNSaCcyhMf4sZjiJVA6A,17261
 nucliadb_utils/cache/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
 nucliadb_utils/cache/exceptions.py,sha256=Zu-O_-0-yctOEgoDGI92gPzWfBMRrpiAyESA62ld6MA,975
-nucliadb_utils/cache/nats.py,sha256=-AjCfkFgKVdJUlGR0hT9JDSNkPVFg4S6w9eW-ZIcXPM,7037
+nucliadb_utils/cache/nats.py,sha256=zCZNJ2kxT6L0y-KPwqB7IIRyssVQBTtpKM2-jQslln4,7159
 nucliadb_utils/cache/pubsub.py,sha256=l8i_RwRf7OPzfmPy-gyn66xgYFs5aHidCIjEaU9VOHE,1654
 nucliadb_utils/cache/settings.py,sha256=WVeHOE6Re5i4k2hUHdFKfkoL4n83v_Z6UPBK6GHYb8g,1059
 nucliadb_utils/encryption/__init__.py,sha256=oav6jFOTGgmIen88sdmy-bCK-uj1tyDt2hr7lB2YKik,2690
@@ -57,7 +57,7 @@ nucliadb_utils/tests/gcs.py,sha256=MBMzn_UHU5SU6iILuCsB5zU4umhNcaCw_MKrxZhwvOc,4
 nucliadb_utils/tests/local.py,sha256=cxIfPrKuqs5Ef0nbrVYQQAH2mwc4E0iD9bC2sWegS-c,1934
 nucliadb_utils/tests/nats.py,sha256=RWHjwqq5esuO7OFbP24yYX1cXnpPLcWJwDUdmwCpH28,1897
 nucliadb_utils/tests/s3.py,sha256=DACUh3HvgH3BchKFZ9R7RFUzsrg3v9A-cxTcXx4nmvA,3734
-nucliadb_utils-6.4.0.post4158.dist-info/METADATA,sha256=KPkRuBKajcHcxE-iszjsutjEtfweRGqp8gE2xvAPoUA,2205
-nucliadb_utils-6.4.0.post4158.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
-nucliadb_utils-6.4.0.post4158.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.4.0.post4158.dist-info/RECORD,,
+nucliadb_utils-6.4.0.post4167.dist-info/METADATA,sha256=OxK4wMK1RRNHI10gstT-ItclGBL_0bpw66kyqo5-GwU,2205
+nucliadb_utils-6.4.0.post4167.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
+nucliadb_utils-6.4.0.post4167.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.4.0.post4167.dist-info/RECORD,,

{nucliadb_utils-6.4.0.post4158.dist-info → nucliadb_utils-6.4.0.post4167.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