nucliadb-utils 6.4.0.post4276__py3-none-any.whl → 6.4.0.post4283__py3-none-any.whl

nucliadb_utils/helpers.py

@@ -17,7 +17,13 @@
 # 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 AsyncGenerator
+import asyncio
+import logging
+from typing import AsyncGenerator, Awaitable, Callable
+
+from nucliadb_telemetry.errors import capture_exception
+
+logger = logging.getLogger(__name__)
 
 
 async def async_gen_lookahead(
@@ -44,3 +50,54 @@ async def async_gen_lookahead(
     # Yield the last chunk if there is one
     if buffered_chunk is not None:
         yield buffered_chunk, True
+
+
+class MessageProgressUpdater:
+    """
+    Context manager to send progress updates to NATS.
+
+    This should allow lower ack_wait time settings without causing
+    messages to be redelivered.
+    """
+
+    _task: asyncio.Task
+
+    def __init__(self, seqid: str, cb: Callable[[], Awaitable[bool]], timeout: float):
+        self.seqid = seqid
+        self.cb = cb
+        self.timeout = timeout
+
+    def start(self):
+        task_name = f"MessageProgressUpdater: {id(self)} (seqid={self.seqid})"
+        self._task = asyncio.create_task(self._progress(), name=task_name)
+
+    async def end(self):
+        self._task.cancel()
+        try:
+            await self._task
+        except asyncio.CancelledError:  # pragma: no cover
+            logger.info("MessageProgressUpdater cancelled")
+            pass
+        except Exception as exc:  # pragma: no cover
+            capture_exception(exc)
+            logger.exception("Error in MessageProgressUpdater")
+            pass
+
+    async def __aenter__(self):
+        self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.end()
+
+    async def _progress(self):
+        while True:
+            try:
+                await asyncio.sleep(self.timeout)
+                done = await self.cb()
+                if done:  # all done, do not mark with in_progress
+                    return
+            except (RuntimeError, asyncio.CancelledError):
+                return
+            except Exception:  # pragma: no cover
+                logger.exception("Error sending task progress to NATS")

nucliadb_utils/nats.py

@@ -28,17 +28,17 @@ import nats
 import nats.errors
 import nats.js.api
 from nats.aio.client import Client as NATSClient
-from nats.aio.client import Msg
+from nats.aio.msg import Msg
 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.metrics import Counter
+from nucliadb_utils.helpers import MessageProgressUpdater
 
 logger = logging.getLogger(__name__)
 
@@ -47,7 +47,7 @@ logger = logging.getLogger(__name__)
 from nucliadb_telemetry.jetstream import get_traced_jetstream  # noqa
 
 
-class MessageProgressUpdater:
+class NatsMessageProgressUpdater(MessageProgressUpdater):
     """
     Context manager to send progress updates to NATS.
 
@@ -55,47 +55,15 @@ class MessageProgressUpdater:
     messages to be redelivered.
     """
 
-    _task: asyncio.Task
-
     def __init__(self, msg: Msg, timeout: float):
-        self.msg = msg
-        self.timeout = timeout
-
-    def start(self):
-        seqid = self.msg.reply.split(".")[5]
-        task_name = f"MessageProgressUpdater: {id(self)} (seqid={seqid})"
-        self._task = asyncio.create_task(self._progress(), name=task_name)
-
-    async def end(self):
-        self._task.cancel()
-        try:
-            await self._task
-        except asyncio.CancelledError:  # pragma: no cover
-            logger.info("MessageProgressUpdater cancelled")
-            pass
-        except Exception as exc:  # pragma: no cover
-            capture_exception(exc)
-            logger.exception("Error in MessageProgressUpdater")
-            pass
-
-    async def __aenter__(self):
-        self.start()
-        return self
-
-    async def __aexit__(self, exc_type, exc_value, traceback):
-        await self.end()
-
-    async def _progress(self):
-        while True:
-            try:
-                await asyncio.sleep(self.timeout)
-                if self.msg._ackd:  # all done, do not mark with in_progress
-                    return
-                await self.msg.in_progress()
-            except (RuntimeError, asyncio.CancelledError):
-                return
-            except Exception:  # pragma: no cover
-                logger.exception("Error sending task progress to NATS")
+        async def update_msg() -> bool:
+            if msg._ackd:  # all done, do not mark with in_progress
+                return True
+            await msg.in_progress()
+            return False
+
+        seqid = msg.reply.split(".")[5]
+        super().__init__(seqid, update_msg, timeout)
 
 
 class NatsConnectionManager:

{nucliadb_utils-6.4.0.post4276.dist-info → nucliadb_utils-6.4.0.post4283.dist-info}/METADATA

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

{nucliadb_utils-6.4.0.post4276.dist-info → nucliadb_utils-6.4.0.post4283.dist-info}/RECORD

@@ -6,8 +6,8 @@ nucliadb_utils/debug.py,sha256=Q56Nx9Dp7V2ae3CU2H0ztaZcHTJXdlflPLKLeOPZ170,2436
 nucliadb_utils/exceptions.py,sha256=y_3wk77WLVUtdo-5FtbBsdSkCtK_DsJkdWb5BoPn3qo,1094
 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=uMs5dTaAplBuj0D0NFFmXGmobrnl4N3FzhCMa-YsnFc,16204
+nucliadb_utils/helpers.py,sha256=eed7_E1MKh9eW3CpqOXka3OvLw5C9eJGC_R-1MPYdfY,3336
+nucliadb_utils/nats.py,sha256=On-2TQRD7U7SVZjhoJWONPH18CUyrWB2Pz-Ax5ji3MA,15191
 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
@@ -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.post4276.dist-info/METADATA,sha256=EHKOe5MxYezmQuO_4PEgq3lzIGxLX3AfZs1LpllDTQ4,2156
-nucliadb_utils-6.4.0.post4276.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-nucliadb_utils-6.4.0.post4276.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.4.0.post4276.dist-info/RECORD,,
+nucliadb_utils-6.4.0.post4283.dist-info/METADATA,sha256=SbycFCOY8JRYEUVzhwyDTQxNYEkVTh3s8Z-mHyQabZU,2156
+nucliadb_utils-6.4.0.post4283.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+nucliadb_utils-6.4.0.post4283.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.4.0.post4283.dist-info/RECORD,,