nucliadb-utils 6.4.0.post4161__py3-none-any.whl → 6.4.0.post4171__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

@@ -23,7 +23,7 @@ import functools
 import os
 import uuid
 from inspect import iscoroutinefunction
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import nats
 import nats.errors
@@ -34,9 +34,14 @@ from nats.aio.subscription import Subscription
 from nats.js.client import JetStreamContext
 from nats.js.manager import JetStreamManager
 
+from nucliadb_telemetry.jetstream import (
+    JetStreamContextTelemetry,
+    NatsClientTelemetry,
+    get_traced_jetstream,
+    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
 
 
 async def wait_for_it(future: asyncio.Future, msg):
@@ -46,7 +51,7 @@ async def wait_for_it(future: asyncio.Future, msg):
 # Configuration Utility
 
 
-class NatsPubsub(PubSubDriver):
+class NatsPubsub(PubSubDriver[Msg]):
     _jetstream = None
     _jsm = None
     _subscriptions: Dict[str, Subscription]
@@ -68,11 +73,11 @@ class NatsPubsub(PubSubDriver):
         self._uuid = os.environ.get("HOSTNAME", uuid.uuid4().hex)
         self.initialized = False
         self.lock = asyncio.Lock()
-        self.nc = None
+        self.nc: Union[Client, NatsClientTelemetry, None] = None
         self.user_credentials_file = user_credentials_file
 
     @property
-    def jetstream(self) -> JetStreamContext:
+    def jetstream(self) -> Union[JetStreamContext, JetStreamContextTelemetry]:
         if self.nc is None:
             raise AttributeError("NC not initialized")
         if self._jetstream is None:
@@ -91,7 +96,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 +109,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
@@ -199,7 +205,7 @@ class NatsPubsub(PubSubDriver):
         else:
             raise ErrConnectionClosed("Could not publish")
 
-    def parse(self, data: Msg):
+    def parse(self, data: Msg) -> bytes:
         return data.data
 
 

nucliadb_utils/cache/pubsub.py

@@ -17,12 +17,14 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
-from typing import Any, Awaitable, Callable, Optional, Union
+from typing import Awaitable, Callable, Generic, Optional, TypeVar
 
-Callback = Union[Callable, Awaitable]
+T = TypeVar("T")
 
+Callback = Callable[[T], Awaitable[None]]
 
-class PubSubDriver:
+
+class PubSubDriver(Generic[T]):
     initialized: bool = False
     async_callback: bool = False
 
@@ -47,5 +49,5 @@ class PubSubDriver:
     ):
         raise NotImplementedError()
 
-    def parse(self, data: Any):
+    def parse(self, data: T) -> bytes:
         raise NotImplementedError()

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
 
@@ -139,8 +143,8 @@ class TransactionUtility:
     ) -> Optional[Event]:
         action_type = self._get_notification_action_type()
 
-        def received(waiting_for: WaitFor, event: Event, raw_data: bytes):
-            data = self.pubsub.parse(raw_data)
+        def received(waiting_for: WaitFor, event: Event, msg: Any):
+            data = self.pubsub.parse(msg)
             pb = Notification()
             pb.ParseFromString(data)
             if pb.uuid == waiting_for.uuid and pb.action == action_type:
@@ -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.post4161.dist-info → nucliadb_utils-6.4.0.post4171.dist-info}/METADATA

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

{nucliadb_utils-6.4.0.post4161.dist-info → nucliadb_utils-6.4.0.post4171.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=l3ZvrITYMnAs_fv1OOC-1nDZxWPG5qmbBhzvuC3DUzQ,8039
 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,11 +23,11 @@ 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/pubsub.py,sha256=l8i_RwRf7OPzfmPy-gyn66xgYFs5aHidCIjEaU9VOHE,1654
+nucliadb_utils/cache/nats.py,sha256=3aV25374fsA_oOr9ZqLDxhNn700O7CIagJ-0Y-ojwqE,7294
+nucliadb_utils/cache/pubsub.py,sha256=YvJHiEhDaMIPofgQRpGY1belN--pOPyfYO1wFTUclt4,1701
 nucliadb_utils/cache/settings.py,sha256=WVeHOE6Re5i4k2hUHdFKfkoL4n83v_Z6UPBK6GHYb8g,1059
 nucliadb_utils/encryption/__init__.py,sha256=oav6jFOTGgmIen88sdmy-bCK-uj1tyDt2hr7lB2YKik,2690
 nucliadb_utils/encryption/settings.py,sha256=yF2AW6c_mu0yWbBn1grAERDAqDvKIpXqQnwD0OFmROI,1467
@@ -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.post4161.dist-info/METADATA,sha256=9cW6zAHR0b3P2YQHwcwyL7kUdoR3QllRFUznTCFI3YE,2205
-nucliadb_utils-6.4.0.post4161.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
-nucliadb_utils-6.4.0.post4161.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.4.0.post4161.dist-info/RECORD,,
+nucliadb_utils-6.4.0.post4171.dist-info/METADATA,sha256=kcPzFKSAR8qvmsclZQgjIb33iOFL9WcO1k_zaYJ7m_s,2205
+nucliadb_utils-6.4.0.post4171.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
+nucliadb_utils-6.4.0.post4171.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.4.0.post4171.dist-info/RECORD,,

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