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

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,10 +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 get_traced_nats_client
+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):
@@ -47,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]
@@ -69,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:
@@ -201,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

@@ -21,7 +21,7 @@ import asyncio
 import logging
 import sys
 import time
-from functools import cached_property
+from functools import cached_property, partial
 from typing import Any, Awaitable, Callable, Optional, Union
 
 import nats
@@ -348,13 +348,18 @@ class NatsConnectionManager:
         durable: Optional[str] = None,
         config: Optional[nats.js.api.ConsumerConfig] = None,
     ) -> JetStreamContext.PullSubscription:
+        wrapped_cb: Callable[[Msg], Awaitable[None]]
+        if isinstance(self.js, JetStreamContextTelemetry):
+            wrapped_cb = partial(self.js.trace_pull_subscriber_message, cb)
+        else:
+            wrapped_cb = cb
+
         psub = await self.js.pull_subscribe(
             subject,
-            durable=durable,  # type: ignore
+            durable=durable,
             stream=stream,
-            config=config,  # type: ignore
+            config=config,
         )
-
         cancelled = asyncio.Event()
 
         async def consume(psub: JetStreamContext.PullSubscription, subject: str):
@@ -372,7 +377,7 @@ class NatsConnectionManager:
                         self.pull_utilization_metrics.inc({"status": "waiting"}, received - start_wait)
 
                     for message in messages:
-                        await cb(message)
+                        await wrapped_cb(message)
 
                     if self.pull_utilization_metrics:
                         processed = time.monotonic()

nucliadb_utils/transaction.py

@@ -143,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:

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

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

{nucliadb_utils-6.4.0.post4167.dist-info → nucliadb_utils-6.4.0.post4176.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=T5A_pH0koGPNH9IPNPv1LsYGnAS5hnJgN8VuNCqO_p8,15989
+nucliadb_utils/nats.py,sha256=uMs5dTaAplBuj0D0NFFmXGmobrnl4N3FzhCMa-YsnFc,16204
 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=K_eXwbbC_exc1DKqce-eIYkMEwbLISmR4dqVOb8geAo,8051
+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
@@ -26,8 +26,8 @@ nucliadb_utils/audit/basic.py,sha256=4NX3IkZtTM_F2jD9TGE8-1eoS-hc-_zWwFGYVahGNSs
 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=zCZNJ2kxT6L0y-KPwqB7IIRyssVQBTtpKM2-jQslln4,7159
-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.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.post4176.dist-info/METADATA,sha256=xc81R9dPXXlDQ26yTGi9u_XxsvpD9J4S8xavW7Snfs0,2205
+nucliadb_utils-6.4.0.post4176.dist-info/WHEEL,sha256=ooBFpIzZCPdw3uqIQsOo4qqbA4ZRPxHnOH7peeONza0,91
+nucliadb_utils-6.4.0.post4176.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.4.0.post4176.dist-info/RECORD,,