nucliadb-utils 6.3.5.post3972__py3-none-any.whl → 6.3.5.post3985__py3-none-any.whl

nucliadb_utils/audit/stream.py

@@ -50,9 +50,9 @@ from nucliadb_protos.kb_usage_pb2 import (
 )
 from nucliadb_protos.nodereader_pb2 import SearchRequest
 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
 
 
@@ -182,8 +182,7 @@ class StreamAuditStorage(AuditStorage):
         if len(self.nats_servers) > 0:
             options["servers"] = self.nats_servers
 
-        nc = await nats.connect(**options)
-        self.nc = get_traced_nats_client(nc, self.service)
+        self.nc = await nats.connect(**options)
 
         self.js = get_traced_jetstream(self.nc, self.service)
         self.task = asyncio.create_task(self.run())

nucliadb_utils/cache/nats.py

@@ -34,9 +34,9 @@ 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_jetstream
 from nucliadb_utils import logger
 from nucliadb_utils.cache.pubsub import Callback, PubSubDriver
+from nucliadb_utils.nats import get_traced_jetstream
 
 
 async def wait_for_it(future: asyncio.Future, msg):

nucliadb_utils/nats.py

@@ -33,17 +33,23 @@ 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,
-    NatsClientTelemetry,
-    get_traced_nats_client,
-)
+from nucliadb_telemetry.jetstream import JetStreamContextTelemetry
+from nucliadb_telemetry.utils import get_telemetry
 
 logger = logging.getLogger(__name__)
 
-# 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
+
+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
 
 
 class MessageProgressUpdater:
@@ -98,7 +104,7 @@ class MessageProgressUpdater:
 
 
 class NatsConnectionManager:
-    _nc: Union[NATSClient, NatsClientTelemetry]
+    _nc: NATSClient
     _subscriptions: list[tuple[Subscription, Callable[[], Awaitable[None]]]]
     _pull_subscriptions: list[
         tuple[
@@ -160,8 +166,7 @@ class NatsConnectionManager:
             options["servers"] = self._nats_servers
 
         async with self._lock:
-            nc = await nats.connect(**options)
-            self._nc = get_traced_nats_client(nc, service_name=self._service_name)
+            self._nc = await nats.connect(**options)
 
         self._expected_subscription_task = asyncio.create_task(self._verify_expected_subscriptions())
 
@@ -302,7 +307,7 @@ class NatsConnectionManager:
         logger.info("Connection is closed on NATS")
 
     @property
-    def nc(self) -> Union[NATSClient, NatsClientTelemetry]:
+    def nc(self) -> NATSClient:
         return self._nc
 
     @cached_property

nucliadb_utils/transaction.py

@@ -35,14 +35,10 @@ from nucliadb_protos.writer_pb2 import (
     Notification,
     OpStatusWriter,
 )
-from nucliadb_telemetry.jetstream import (
-    JetStreamContextTelemetry,
-    NatsClientTelemetry,
-    get_traced_jetstream,
-    get_traced_nats_client,
-)
+from nucliadb_telemetry.jetstream import JetStreamContextTelemetry
 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
 
 
@@ -99,7 +95,7 @@ class LocalTransactionUtility:
 
 
 class TransactionUtility:
-    nc: Union[Client, NatsClientTelemetry]
+    nc: Client
     js: Union[JetStreamContext, JetStreamContextTelemetry]
     pubsub: PubSubDriver
 
@@ -166,7 +162,7 @@ class TransactionUtility:
         )
         return waiting_event
 
-    async def initialize(self, service_name: str):
+    async def initialize(self, service_name: Optional[str] = None):
         self.pubsub = await get_pubsub()  # type: ignore
 
         options: dict[str, Any] = {
@@ -181,9 +177,8 @@ class TransactionUtility:
         if len(self.nats_servers) > 0:
             options["servers"] = self.nats_servers
 
-        nc = await nats.connect(**options)
-        self.nc = get_traced_nats_client(nc, service_name)
-        self.js = get_traced_jetstream(self.nc, service_name)
+        self.nc = await nats.connect(**options)
+        self.js = get_traced_jetstream(self.nc, service_name or "nucliadb")
 
     async def finalize(self):
         try:

{nucliadb_utils-6.3.5.post3972.dist-info → nucliadb_utils-6.3.5.post3985.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb_utils
-Version: 6.3.5.post3972
+Version: 6.3.5.post3985
 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.3.5.post3972
-Requires-Dist: nucliadb-telemetry>=6.3.5.post3972
+Requires-Dist: nucliadb-protos>=6.3.5.post3985
+Requires-Dist: nucliadb-telemetry>=6.3.5.post3985
 Provides-Extra: cache
 Requires-Dist: redis>=4.3.4; extra == "cache"
 Requires-Dist: orjson>=3.6.7; extra == "cache"

{nucliadb_utils-6.3.5.post3972.dist-info → nucliadb_utils-6.3.5.post3985.dist-info}/RECORD

@@ -7,14 +7,14 @@ nucliadb_utils/exceptions.py,sha256=y_3wk77WLVUtdo-5FtbBsdSkCtK_DsJkdWb5BoPn3qo,
 nucliadb_utils/featureflagging.py,sha256=YxDiXzWuiDlHtqgeTVeyakmbAdzBePzJpgJv53ELbmI,2259
 nucliadb_utils/grpc.py,sha256=apu0uePnkGHCAT7GRQ9YZfRYyFj26kJ440i8jitbM3U,3314
 nucliadb_utils/helpers.py,sha256=nPw8yod3hP-pxq80VF8QC36s7ygSg0dBUdfI-LatvCs,1600
-nucliadb_utils/nats.py,sha256=aeqtkNPFsBbTZ9AMqYy2W4-AJIt4s8QqEtu46BrFqkw,15110
+nucliadb_utils/nats.py,sha256=-ehKk0u5ZhiJEyg67yV26YysBGEGp4TDSQKq3aydJn0,15275
 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=h8nt_qYXDlDxGbqI130Okaat2n4w1W05GDzd7akDVDI,8006
+nucliadb_utils/transaction.py,sha256=z_VeiTIta48rosS2SXMqx86XaavprSMRWf6s6zWIeEs,7920
 nucliadb_utils/utilities.py,sha256=Vc4zLpDf-FQh9bs0ZyAfGAjzlbbMTMWf3VWt2Ao5V3k,15379
 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=VDUqcSToEoz5m7iG-4omjTlsBSEJGcp0_xl51idnU3Q,3384
 nucliadb_utils/audit/basic.py,sha256=pVf1oKp7T_oTDCyt8QbSsfC_uWjF66ooF6rqY-SWSrc,4032
-nucliadb_utils/audit/stream.py,sha256=AtXTERd6yEYkIc02fn09m716pnEpIszD39-nmVmxsbI,16122
+nucliadb_utils/audit/stream.py,sha256=SovPj5iTzXka2cg4NnwAZwnOYljhsT6ynhDC2OsRVPc,16035
 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=NH-QxN1ImUV--Ox34QFbx8Wf6IIxDVxxtJTOs_pexek,7046
+nucliadb_utils/cache/nats.py,sha256=-AjCfkFgKVdJUlGR0hT9JDSNkPVFg4S6w9eW-ZIcXPM,7037
 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.3.5.post3972.dist-info/METADATA,sha256=84aI-bPVWNv3h6aPP6k3pNKD-BFHlEDH4Fmara0zUP0,2205
-nucliadb_utils-6.3.5.post3972.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-nucliadb_utils-6.3.5.post3972.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.3.5.post3972.dist-info/RECORD,,
+nucliadb_utils-6.3.5.post3985.dist-info/METADATA,sha256=PHCoEtl7WwdYbBmfkRIa09XMsIsjetg3uKLeDTr4ynY,2205
+nucliadb_utils-6.3.5.post3985.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+nucliadb_utils-6.3.5.post3985.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.3.5.post3985.dist-info/RECORD,,