nucliadb-utils 6.1.0.post2473__py3-none-any.whl → 6.1.0.post2476__py3-none-any.whl

nucliadb_utils/const.py

@@ -30,7 +30,7 @@ class Streams:
 
         name = "nucliadb"
         subject = "ndb.consumer.{partition}"
-        group = "nucliadb-{partition}"
+        group = "nucliadb-pull-{partition}"
 
     class INGEST_PROCESSED:
         """
@@ -40,7 +40,7 @@ class Streams:
 
         name = "nucliadb"
         subject = "ndb.consumer.processed"
-        group = "nucliadb-processed"
+        group = "nucliadb-pull-processed"
 
     class INDEX:
         """
@@ -80,3 +80,4 @@ class Features:
     LOG_REQUEST_PAYLOADS = "nucliadb_log_request_payloads"
     IGNORE_EXTRACTED_IN_SEARCH = "nucliadb_ignore_extracted_in_search"
     NIDX_READS = "nucliadb_nidx_reads"
+    PULL_PROCESSED_CONSUMERS_DEPLOYED = "nucliadb_pull_processed_consumers_deployed"

nucliadb_utils/featureflagging.py

@@ -69,6 +69,10 @@ DEFAULT_FLAG_DATA: dict[str, Any] = {
         "rollout": 0,
         "variants": {"environment": ["local"]},
     },
+    const.Features.PULL_PROCESSED_CONSUMERS_DEPLOYED: {
+        "rollout": 0,
+        "variants": {"environment": ["local"]},
+    },
 }
 
 

nucliadb_utils/nats.py

@@ -104,6 +104,11 @@ class MessageProgressUpdater:
 class NatsConnectionManager:
     _nc: NATSClient
     _subscriptions: list[tuple[Subscription, Callable[[], Awaitable[None]]]]
+    _pull_subscriptions: list[
+        tuple[
+            JetStreamContext.PullSubscription, asyncio.Task, Callable[[], Awaitable[None]], asyncio.Event
+        ]
+    ]
     _unhealthy_timeout = 10  # needs to be unhealth for 10 seconds to be unhealthy and force exit
 
     def __init__(
@@ -117,6 +122,7 @@ class NatsConnectionManager:
         self._nats_servers = nats_servers
         self._nats_creds = nats_creds
         self._subscriptions = []
+        self._pull_subscriptions = []
         self._lock = asyncio.Lock()
         self._healthy = True
         self._last_unhealthy: Optional[float] = None
@@ -155,11 +161,26 @@ class NatsConnectionManager:
 
     async def finalize(self):
         async with self._lock:
+            # Finalize push subscriptions
             for sub, _ in self._subscriptions:
                 try:
                     await sub.drain()
                 except nats.errors.ConnectionClosedError:  # pragma: no cover
                     pass
+            self._subscriptions = []
+
+            # Finalize pull subscriptions
+            for pull_sub, task, _, cancelled in self._pull_subscriptions:
+                cancelled.set()
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+                await pull_sub.unsubscribe()
+            self._pull_subscriptions = []
+
+            # Close the connection
             try:
                 await asyncio.wait_for(self._nc.drain(), timeout=1)
             except (
@@ -168,7 +189,6 @@ class NatsConnectionManager:
             ):  # pragma: no cover
                 pass
             await self._nc.close()
-            self._subscriptions = []
 
     async def disconnected_cb(self) -> None:
         logger.info("Disconnected from NATS!")
@@ -193,6 +213,26 @@ class NatsConnectionManager:
                     # should force exit here to restart the service
                     self._healthy = False
                     raise
+
+            existing_pull_subs = self._pull_subscriptions
+            self._pull_subscriptions = []
+            for pull_sub, task, recon_callback, cancelled in existing_pull_subs:
+                cancelled.set()
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+                try:
+                    await pull_sub.unsubscribe()
+                    await recon_callback()
+                except Exception:
+                    logger.exception(
+                        f"Error resubscribing to pull subscription on {self._nc.connected_url.netloc}"
+                    )
+                    # should force exit here to restart the service
+                    self._healthy = False
+                    raise
         self._healthy = True
         self._last_unhealthy = None  # reset the last unhealthy time
 
@@ -236,17 +276,67 @@ class NatsConnectionManager:
 
         return sub
 
-    async def _remove_subscription(self, subscription: Subscription):
+    async def pull_subscribe(
+        self,
+        *,
+        subject: str,
+        stream: str,
+        cb: Callable[[Msg], Awaitable[None]],
+        subscription_lost_cb: Callable[[], Awaitable[None]],
+        durable: Optional[str] = None,
+        config: Optional[nats.js.api.ConsumerConfig] = None,
+    ) -> JetStreamContext.PullSubscription:
+        psub = await self.js.pull_subscribe(
+            subject,
+            durable=durable,  # type: ignore
+            stream=stream,
+            config=config,  # type: ignore
+        )
+
+        cancelled = asyncio.Event()
+
+        async def consume(psub: JetStreamContext.PullSubscription, subject: str):
+            while True:
+                if cancelled.is_set():
+                    break
+                try:
+                    messages = await psub.fetch(batch=1)
+                    for message in messages:
+                        await cb(message)
+                except asyncio.CancelledError:
+                    # Handle task cancellation
+                    logger.info("Pull subscription consume task cancelled", extra={"subject": subject})
+                    break
+                except TimeoutError:
+                    pass
+                except Exception:
+                    logger.exception("Error in pull_subscribe task", extra={"subject": subject})
+
+        task = asyncio.create_task(consume(psub, subject), name=f"pull_subscribe_{subject}")
+        self._pull_subscriptions.append((psub, task, subscription_lost_cb, cancelled))
+        return psub
+
+    async def _remove_subscription(
+        self, subscription: Union[Subscription, JetStreamContext.PullSubscription]
+    ):
         async with self._lock:
-            sub_index = None
             for index, (sub, _) in enumerate(self._subscriptions):
                 if sub is not subscription:
                     continue
-                sub_index = index
-                break
-            if sub_index is not None:
-                self._subscriptions.pop(sub_index)
+                self._subscriptions.pop(index)
+                return
+            for index, (psub, task, _, cancelled) in enumerate(self._pull_subscriptions):
+                if psub is not subscription:
+                    continue
+                self._pull_subscriptions.pop(index)
+                cancelled.set()
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+                return
 
-    async def unsubscribe(self, subscription: Subscription):
+    async def unsubscribe(self, subscription: Union[Subscription, JetStreamContext.PullSubscription]):
         await subscription.unsubscribe()
         await self._remove_subscription(subscription)

nucliadb_utils/tests/gcs.py

@@ -75,7 +75,7 @@ class GCS(BaseImage):
 @pytest.fixture(scope="session")
 def gcs():
     container = GCS()
-    host, port = container.run()
+    _, port = container.run()
     public_api_url = f"http://172.17.0.1:{port}"
     yield public_api_url
     container.stop()

{nucliadb_utils-6.1.0.post2473.dist-info → nucliadb_utils-6.1.0.post2476.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb_utils
-Version: 6.1.0.post2473
+Version: 6.1.0.post2476
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -24,8 +24,8 @@ Requires-Dist: PyNaCl
 Requires-Dist: pyjwt>=2.4.0
 Requires-Dist: memorylru>=1.1.2
 Requires-Dist: mrflagly>=0.2.9
-Requires-Dist: nucliadb-protos>=6.1.0.post2473
-Requires-Dist: nucliadb-telemetry>=6.1.0.post2473
+Requires-Dist: nucliadb-protos>=6.1.0.post2476
+Requires-Dist: nucliadb-telemetry>=6.1.0.post2476
 Provides-Extra: cache
 Requires-Dist: redis>=4.3.4; extra == "cache"
 Requires-Dist: orjson>=3.6.7; extra == "cache"

{nucliadb_utils-6.1.0.post2473.dist-info → nucliadb_utils-6.1.0.post2476.dist-info}/RECORD

@@ -1,14 +1,14 @@
 nucliadb_utils/__init__.py,sha256=EvBCH1iTODe-AgXm48aj4kVUt_Std3PeL8QnwimR5wI,895
 nucliadb_utils/asyncio_utils.py,sha256=h8Y-xpcFFRgNzaiIW0eidz7griAQa7ggbNk34-tAt2c,2888
 nucliadb_utils/authentication.py,sha256=N__d2Ez3JHJv5asYK5TgUcIkKqcAC8ZTLlnfLhfSneM,5837
-nucliadb_utils/const.py,sha256=_UhW7A3VC0EG-OQDse4maZdP8hfsI_NQrLfBkonEFtA,2477
+nucliadb_utils/const.py,sha256=0zsA5Ys2NQoJECenjX8aSgQGqizqmsUZCcIp87E-61A,2572
 nucliadb_utils/debug.py,sha256=Q56Nx9Dp7V2ae3CU2H0ztaZcHTJXdlflPLKLeOPZ170,2436
 nucliadb_utils/exceptions.py,sha256=y_3wk77WLVUtdo-5FtbBsdSkCtK_DsJkdWb5BoPn3qo,1094
-nucliadb_utils/featureflagging.py,sha256=dKp8u4g52KTcupKFJkdGYk3Z2pGQdi_q7IaoY_GQ0nE,2953
+nucliadb_utils/featureflagging.py,sha256=5J5Je_cZ05f0xnKLvWDMLi3eqKGKAOfR_MItDnAHJj0,3086
 nucliadb_utils/grpc.py,sha256=apu0uePnkGHCAT7GRQ9YZfRYyFj26kJ440i8jitbM3U,3314
 nucliadb_utils/helpers.py,sha256=nPw8yod3hP-pxq80VF8QC36s7ygSg0dBUdfI-LatvCs,1600
 nucliadb_utils/indexing.py,sha256=Luaqcar3CySpdYOFp6Q9Fyr8ZYwhYhaKRHQ_VGL78f8,2318
-nucliadb_utils/nats.py,sha256=prs98dIrGh2-WchQcRBloc5wYhpjEnB0VtCERPnoReg,8430
+nucliadb_utils/nats.py,sha256=SGBWtP0h3_9Ne6VEfZ2OLnWlEvzfAj63U252FrTITjY,11870
 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
@@ -54,13 +54,13 @@ nucliadb_utils/tests/__init__.py,sha256=Oo9CAE7B0eW5VHn8sHd6o30SQzOWUhktLPRXdlDO
 nucliadb_utils/tests/asyncbenchmark.py,sha256=x4be2IwCawle9zWgMOJkmwoUwk5p1tv7cLQGmybkEOg,10587
 nucliadb_utils/tests/azure.py,sha256=Dg-Eb4KVScG-O6P9y-bVQZTAKTNUMQ0i-CKEd9IdrWw,4474
 nucliadb_utils/tests/fixtures.py,sha256=i0sqPqe5a5JlKGFdaIvOlHYkZ3pHZ2hTIgTsaIB3vSM,3472
-nucliadb_utils/tests/gcs.py,sha256=weTGkuPyAbOzUwlXC4VHGQMJG1iKlsny842yJ8wslhk,4502
+nucliadb_utils/tests/gcs.py,sha256=wO4xRhl-_wrddxlGyEktqa_bXHeQarQN-l_A72Tzr9A,4499
 nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQFxk,1513
 nucliadb_utils/tests/local.py,sha256=fXIBasrvdaFJM-sw2wk1_oiFzBcm9O10iCyC-OiXwY8,1914
 nucliadb_utils/tests/nats.py,sha256=xqpww4jZjTKY9oPGlJdDJG67L3FIBQsa9qDHxILR8r8,7687
 nucliadb_utils/tests/s3.py,sha256=pl-RJFjA4MH6iXkqhsh5g8gDuEhrYu1nPZ-laxlrMlE,3704
-nucliadb_utils-6.1.0.post2473.dist-info/METADATA,sha256=A2nR8tUEDZEnLS67mG611J17y1csF5DOHz9dHkJmymw,2055
-nucliadb_utils-6.1.0.post2473.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-nucliadb_utils-6.1.0.post2473.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.1.0.post2473.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-6.1.0.post2473.dist-info/RECORD,,
+nucliadb_utils-6.1.0.post2476.dist-info/METADATA,sha256=9xT2in9NaLwz4F16XWog7Kjcx3Udzqz3UlCHKN0dH1U,2055
+nucliadb_utils-6.1.0.post2476.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+nucliadb_utils-6.1.0.post2476.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.1.0.post2476.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-6.1.0.post2476.dist-info/RECORD,,