nucliadb 6.3.1.post3590__py3-none-any.whl → 6.3.1.post3598__py3-none-any.whl

nucliadb/backups/create.py

@@ -170,7 +170,7 @@ async def backup_resource_with_binaries(
         """
         nonlocal total_size
 
-        for cloud_file in get_cloud_files(bm):
+        for index, cloud_file in enumerate(get_cloud_files(bm)):
             if not await exists_cf(context, cloud_file):
                 logger.warning(
                     "Cloud file not found in storage, skipping",
@@ -184,13 +184,13 @@ async def backup_resource_with_binaries(
                 yield serialized_cf
 
             async for chunk in to_tar(
-                name=f"cloud-files/{cloud_file.uri}", size=len(serialized_cf), chunks=cf_iterator()
+                name=f"cloud-files/{index}", size=len(serialized_cf), chunks=cf_iterator()
             ):
                 yield chunk
                 total_size += len(chunk)
 
             async for chunk in to_tar(
-                name=f"binaries/{cloud_file.uri}",
+                name=f"binaries/{index}",
                 size=cloud_file.size,
                 chunks=download_binary(context, cloud_file),
             ):

nucliadb/backups/restore.py

@@ -31,7 +31,7 @@ from nucliadb.backups.settings import settings
 from nucliadb.common.context import ApplicationContext
 from nucliadb.export_import.utils import (
     import_binary,
-    import_broker_message,
+    restore_broker_message,
     set_entities_groups,
     set_labels,
 )
@@ -69,11 +69,11 @@ async def restore_kb(context: ApplicationContext, kbid: str, backup_id: str):
     await restore_resources(context, kbid, backup_id)
     await restore_labels(context, kbid, backup_id)
     await restore_entities(context, kbid, backup_id)
-    await delete_last_restored_resource_key(context, kbid, backup_id)
+    await delete_last_restored(context, kbid, backup_id)
 
 
 async def restore_resources(context: ApplicationContext, kbid: str, backup_id: str):
-    last_restored = await get_last_restored_resource_key(context, kbid, backup_id)
+    last_restored = await get_last_restored(context, kbid, backup_id)
     tasks = []
     async for object_info in context.blob_storage.iterate_objects(
         bucket=settings.backups_bucket,
@@ -81,21 +81,19 @@ async def restore_resources(context: ApplicationContext, kbid: str, backup_id: s
         start=last_restored,
     ):
         key = object_info.name
-        resource_id = key.split("/")[-1].rstrip(".tar")
+        resource_id = key.split("/")[-1].split(".tar")[0]
         tasks.append(asyncio.create_task(restore_resource(context, kbid, backup_id, resource_id)))
         if len(tasks) > settings.restore_resources_concurrency:
             await asyncio.gather(*tasks)
             tasks = []
-            await set_last_restored_resource_key(context, kbid, backup_id, key)
+            await set_last_restored(context, kbid, backup_id, key)
     if len(tasks) > 0:
         await asyncio.gather(*tasks)
         tasks = []
-        await set_last_restored_resource_key(context, kbid, backup_id, key)
+        await set_last_restored(context, kbid, backup_id, key)
 
 
-async def get_last_restored_resource_key(
-    context: ApplicationContext, kbid: str, backup_id: str
-) -> Optional[str]:
+async def get_last_restored(context: ApplicationContext, kbid: str, backup_id: str) -> Optional[str]:
     key = MaindbKeys.LAST_RESTORED.format(kbid=kbid, backup_id=backup_id)
     async with context.kv_driver.transaction(read_only=True) as txn:
         raw = await txn.get(key)
@@ -104,16 +102,14 @@ async def get_last_restored_resource_key(
         return raw.decode()
 
 
-async def set_last_restored_resource_key(
-    context: ApplicationContext, kbid: str, backup_id: str, resource_id: str
-):
+async def set_last_restored(context: ApplicationContext, kbid: str, backup_id: str, resource_id: str):
     key = MaindbKeys.LAST_RESTORED.format(kbid=kbid, backup_id=backup_id)
     async with context.kv_driver.transaction() as txn:
         await txn.set(key, resource_id.encode())
         await txn.commit()
 
 
-async def delete_last_restored_resource_key(context: ApplicationContext, kbid: str, backup_id: str):
+async def delete_last_restored(context: ApplicationContext, kbid: str, backup_id: str):
     key = MaindbKeys.LAST_RESTORED.format(kbid=kbid, backup_id=backup_id)
     async with context.kv_driver.transaction() as txn:
         await txn.delete(key)
@@ -134,6 +130,7 @@ class ResourceBackupReader:
     def __init__(self, download_stream: AsyncIterator[bytes]):
         self.download_stream = download_stream
         self.buffer = b""
+        self.cloud_files: dict[int, CloudFile] = {}
 
     async def read(self, size: int) -> bytes:
         while len(self.buffer) < size:
@@ -194,15 +191,18 @@ class ResourceBackupReader:
             bm.ParseFromString(raw_bm)
             return bm
         elif tarinfo.name.startswith("cloud-files"):
+            cf_index = int(tarinfo.name.split("cloud-files/")[-1])
             raw_cf = await self.read_data(tarinfo)
             cf = CloudFile()
             cf.ParseFromString(raw_cf)
+            self.cloud_files[cf_index] = cf
             return cf
         elif tarinfo.name.startswith("binaries"):
-            uri = tarinfo.name.lstrip("binaries/")
+            bin_index = int(tarinfo.name.split("binaries/")[-1])
             size = tarinfo.size
             download_stream = functools.partial(self.iter_data, size)
-            return CloudFileBinary(uri, download_stream)
+            cf = self.cloud_files[bin_index]
+            return CloudFileBinary(cf.uri, download_stream)
         else:  # pragma: no cover
             raise ValueError(f"Unknown tar entry: {tarinfo.name}")
 
@@ -235,8 +235,8 @@ async def restore_resource(context: ApplicationContext, kbid: str, backup_id: st
                 extra={"item_type": type(item), kbid: kbid, resource_id: resource_id},
             )
             continue
-
-    await import_broker_message(context, kbid, bm)
+    if bm is not None:
+        await restore_broker_message(context, kbid, bm)
 
 
 async def restore_labels(context: ApplicationContext, kbid: str, backup_id: str):

nucliadb/export_import/utils.py

@@ -39,6 +39,7 @@ from nucliadb_protos import knowledgebox_pb2 as kb_pb2
 from nucliadb_protos import resources_pb2, writer_pb2
 from nucliadb_utils.const import Streams
 from nucliadb_utils.transaction import MaxTransactionSizeExceededError
+from nucliadb_utils.utilities import get_ingest
 
 BinaryStream = AsyncIterator[bytes]
 BinaryStreamGenerator = Callable[[int], BinaryStream]
@@ -72,6 +73,24 @@ async def import_broker_message(
         await transaction_commit(context, pb, partition)
 
 
+async def restore_broker_message(
+    context: ApplicationContext, kbid: str, bm: writer_pb2.BrokerMessage
+) -> None:
+    bm.kbid = kbid
+
+    # First, ingest the broker message writer part synchronously
+    async def ingest_request_stream() -> AsyncIterator[writer_pb2.BrokerMessage]:
+        yield get_writer_bm(bm)
+
+    response = await get_ingest().ProcessMessage(ingest_request_stream())  # type: ignore
+    assert response.status == writer_pb2.OpStatusWriter.Status.OK, "Failed to process broker message"
+
+    # Then enqueue the processor part asynchronously
+    processor_bm = get_processor_bm(bm)
+    partition = context.partitioning.generate_partition(kbid, bm.uuid)
+    await transaction_commit(context, processor_bm, partition)
+
+
 async def transaction_commit(
     context: ApplicationContext, bm: writer_pb2.BrokerMessage, partition: int
 ) -> None:

nucliadb/ingest/service/writer.py

@@ -202,9 +202,9 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
             return DeleteKnowledgeBoxResponse(status=KnowledgeBoxResponseStatus.ERROR)
         return DeleteKnowledgeBoxResponse(status=KnowledgeBoxResponseStatus.OK)
 
-    async def ProcessMessage(  # type: ignore
+    async def ProcessMessage(
         self, request_stream: AsyncIterator[BrokerMessage], context=None
-    ):
+    ) -> OpStatusWriter:
         response = OpStatusWriter()
         async for message in request_stream:
             try:

nucliadb/search/lifecycle.py

@@ -22,6 +22,7 @@ from contextlib import asynccontextmanager
 from fastapi import FastAPI
 
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
+from nucliadb.common.context.fastapi import inject_app_context
 from nucliadb.common.maindb.utils import setup_driver
 from nucliadb.common.nidx import start_nidx_utility
 from nucliadb.ingest.utils import start_ingest, stop_ingest
@@ -51,7 +52,8 @@ async def lifespan(app: FastAPI):
 
     await start_audit_utility(SERVICE_NAME)
 
-    yield
+    async with inject_app_context(app):
+        yield
 
     await stop_ingest()
     if get_utility(Utility.PARTITION):

nucliadb/tasks/consumer.py

@@ -105,7 +105,7 @@ class NatsTaskConsumer(Generic[MsgType]):
 
     async def _subscription_worker_as_task(self, msg: Msg):
         seqid = int(msg.reply.split(".")[5])
-        task_name = f"NatsTaskConsumer({self.name}, msg={seqid})"
+        task_name = f"NatsTaskConsumer({self.name}, stream={self.stream.name}, subject={self.consumer.subject}, seqid={seqid})"
         task = asyncio.create_task(self.subscription_worker(msg), name=task_name)
         task.add_done_callback(self._running_tasks_remove)
         self.running_tasks.append(task)

nucliadb/train/lifecycle.py

@@ -22,6 +22,7 @@ from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
+from nucliadb.common.context.fastapi import inject_app_context
 from nucliadb.common.nidx import start_nidx_utility, stop_nidx_utility
 from nucliadb.train import SERVICE_NAME
 from nucliadb.train.utils import (
@@ -42,7 +43,8 @@ async def lifespan(app: FastAPI):
     await start_train_grpc(SERVICE_NAME)
     await start_audit_utility(SERVICE_NAME)
 
-    yield
+    async with inject_app_context(app):
+        yield
 
     await stop_audit_utility()
     await stop_train_grpc()

{nucliadb-6.3.1.post3590.dist-info → nucliadb-6.3.1.post3598.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nucliadb
-Version: 6.3.1.post3590
+Version: 6.3.1.post3598
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -20,11 +20,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.3.1.post3590
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.1.post3590
-Requires-Dist: nucliadb-protos>=6.3.1.post3590
-Requires-Dist: nucliadb-models>=6.3.1.post3590
-Requires-Dist: nidx-protos>=6.3.1.post3590
+Requires-Dist: nucliadb-telemetry[all]>=6.3.1.post3598
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.1.post3598
+Requires-Dist: nucliadb-protos>=6.3.1.post3598
+Requires-Dist: nucliadb-models>=6.3.1.post3598
+Requires-Dist: nidx-protos>=6.3.1.post3598
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn

{nucliadb-6.3.1.post3590.dist-info → nucliadb-6.3.1.post3598.dist-info}/RECORD

@@ -41,10 +41,10 @@ nucliadb/openapi.py,sha256=wDiw0dVEvTpJvbatkJ0JZLkKm9RItZT5PWRHjqRfqTA,2272
 nucliadb/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/backups/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/backups/const.py,sha256=9vPAhLxQO_gNAjSdPxWuv3V66s9WcdpjOQ89CZlfmuk,1894
-nucliadb/backups/create.py,sha256=w6KpjkTOpIoPVrtOZrxB99oMdWp6jYkyIJ1-qo_k4Mw,11230
+nucliadb/backups/create.py,sha256=x5Ezd1p8qqlcn-srbE_iFxTwEEjycKWhTezUIkXL-oo,11230
 nucliadb/backups/delete.py,sha256=1rnBhVUGYYZJXSZUrrgYMDZ5NyswEWkIA-G-crRCyHk,2404
 nucliadb/backups/models.py,sha256=-hITU4Mv6AxePu12toBu_fjpEv6vVGcwNVxV22O9jQA,1273
-nucliadb/backups/restore.py,sha256=wepEgv4vBN5yeiZU-f17PbuFV4xT4_SVKplNr8xSJrE,10001
+nucliadb/backups/restore.py,sha256=YD3Bbo9ry4YLMM6imB-DXbOAMXfGxVzJtTAAUFDvB0I,10153
 nucliadb/backups/settings.py,sha256=SyzsInj1BRbBI0atg5IXWbMbOZ_eVg4eSQ3IcnUhCxQ,1357
 nucliadb/backups/tasks.py,sha256=4_kOVJ2yCwMvDEpzJgTuTt75TNlpq5woyw9sTAcaSkw,4194
 nucliadb/backups/utils.py,sha256=_Vogjqcru5oqNZM-bZ0q7Ju79Bv1PD-LVFEa7Z-Q13I,1261
@@ -109,7 +109,7 @@ nucliadb/export_import/exporter.py,sha256=k2QVx1EjqFlDYiggriWiEJzwtMXzHbldsqWdpG
 nucliadb/export_import/importer.py,sha256=v5cq9Nn8c2zrY_K_00mydR52f8mdFxR7tLdtNLQ0qvk,4229
 nucliadb/export_import/models.py,sha256=dbjScNkiMRv4X3Ktudy1JRliD25bfoDTy3JmEZgQSCc,2121
 nucliadb/export_import/tasks.py,sha256=DWbdqY97ffoyfipelGXz3Jqz1iam6JCjQSh367Fc3NA,2947
-nucliadb/export_import/utils.py,sha256=iAQAjYuNx0dhM2b5-1A0NEs8tSRsznuT-izysUrTwS0,19986
+nucliadb/export_import/utils.py,sha256=DlGUHaqT43b3jG9U-vZ48GpC4O2OgD2WSP_0-hrYW9k,20774
 nucliadb/ingest/__init__.py,sha256=fsw3C38VP50km3R-nHL775LNGPpJ4JxqXJ2Ib1f5SqE,1011
 nucliadb/ingest/app.py,sha256=rX1KE5vsAzG9hlArBk8WE2SOlvdYylcb-jNkMQNPJdQ,7407
 nucliadb/ingest/partitions.py,sha256=2NIhMYbNT0TNBL6bX1UMSi7vxFGICstCKEqsB0TXHOE,2410
@@ -152,7 +152,7 @@ nucliadb/ingest/orm/processor/processor.py,sha256=wjZmfJumypdGoKVB0BEJ51ha3xNRs1
 nucliadb/ingest/orm/processor/sequence_manager.py,sha256=uqEphtI1Ir_yk9jRl2gPf7BlzzXWovbARY5MNZSBI_8,1704
 nucliadb/ingest/service/__init__.py,sha256=MME_G_ERxzJR6JW_hfE2qcfXpmpH1kdG-S0a-M0qRm8,2043
 nucliadb/ingest/service/exceptions.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/ingest/service/writer.py,sha256=8oHJru1Yc8e05MXKmdQMbcNOzLHFFbWmHlJXIMjJcnY,22315
+nucliadb/ingest/service/writer.py,sha256=duFp2EOA52Mkz8kJcV1apCDgsUmkwpwGuubVra5Aa7o,22317
 nucliadb/middleware/__init__.py,sha256=A8NBlBuEkunCFMKpR9gnfNELsVn0Plc55BIQMbWDM8Q,2202
 nucliadb/migrator/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/migrator/command.py,sha256=dKbJ1tAmP6X4lMVRSSlz351euaqs2wBPpOczLjATUes,2089
@@ -188,7 +188,7 @@ nucliadb/reader/reader/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXM
 nucliadb/reader/reader/notifications.py,sha256=HVZNUlfbSuoZ9BsSs8wmzPeYurl0U0O2ooVlR9KSM3U,7792
 nucliadb/search/__init__.py,sha256=tnypbqcH4nBHbGpkINudhKgdLKpwXQCvDtPchUlsyY4,1511
 nucliadb/search/app.py,sha256=-WEX1AZRA8R_9aeOo9ovOTwjXW_7VfwWN7N2ccSoqXg,3387
-nucliadb/search/lifecycle.py,sha256=DW8v4WUi4rZqc7xTOi3rE67W7877WG7fH9oTZbolHdE,2099
+nucliadb/search/lifecycle.py,sha256=V_Pj5PRP0yyDY8d5LytO4X8p9HhN7UomqRG6Ri0UaFA,2206
 nucliadb/search/openapi.py,sha256=t3Wo_4baTrfPftg2BHsyLWNZ1MYn7ZRdW7ht-wFOgRs,1016
 nucliadb/search/predict.py,sha256=z2-RkhMkH-5T6PtFkfESxNof07XiS5FxicLHPRyCUXc,22284
 nucliadb/search/predict_models.py,sha256=ZAe0dneUsPmV9uBar57cCFADCGOrYDsJHuqKlA5zWag,5937
@@ -272,7 +272,7 @@ nucliadb/standalone/static/favicon.ico,sha256=96pKGp6Sx457JkTfjy1dtApMhkitixfU6i
 nucliadb/standalone/static/index.html,sha256=PEZfuEQFYnYACAL1ceN8xC0im8lBrUx838RkE8tbvgA,3833
 nucliadb/standalone/static/logo.svg,sha256=-wQqSvPGTdlKjUP6pHE6kiq005pgYjDzp9nPl0X71Mk,2639
 nucliadb/tasks/__init__.py,sha256=oFJ3A8HD7w11mBu-IixYE_KxA7juMGlYQb7YD_y6WPM,975
-nucliadb/tasks/consumer.py,sha256=xc0Ql3N1Iq52dJ3t4YYGJFj1NCQAly0J5W_brfLa_F8,6894
+nucliadb/tasks/consumer.py,sha256=4CWfBdXVr2a25n7seldbQ0PaK0FcxJZuWgosU6aODS8,6956
 nucliadb/tasks/logger.py,sha256=C7keOEO_mjLVp5VbqAZ2QXfqVB2Hot7NgBlUP_SDSMw,924
 nucliadb/tasks/models.py,sha256=qrZKi5DNDQ07waMsp5L4_Fi7WRs57YiO-kmXlrBzEAA,1168
 nucliadb/tasks/producer.py,sha256=UnpJAzhj_GElsCoO5G6T4m6MshsgOaqR2tVzJmEta64,2625
@@ -284,7 +284,7 @@ nucliadb/tests/vectors.py,sha256=CcNKx-E8LPpyvRyljbmb-Tn_wST9Juw2CBoogWrKiTk,628
 nucliadb/train/__init__.py,sha256=NVwe5yULoHXb80itIJT8YJYEz2xbiOPQ7_OMys6XJw8,1301
 nucliadb/train/app.py,sha256=TiRttTvekLuZdIvi46E4HyuumDTkR4G4Luqq3fEdjes,2824
 nucliadb/train/generator.py,sha256=0_zqWsLUHmJZl0lXhGorO5CWSkl42-k78dqb1slZ5h0,3904
-nucliadb/train/lifecycle.py,sha256=7rW2Be9cMUg7-fY9JhC5krXBB7MDOM4BpfNRRO5IAec,1713
+nucliadb/train/lifecycle.py,sha256=a96KuAVZ0sf9TVVW6v6szVXn2eGBaboszwnv_HdWiZk,1820
 nucliadb/train/models.py,sha256=BmgmMjDsu_1Ih5JDAqo6whhume90q0ASJcDP9dkMQm8,1198
 nucliadb/train/nodes.py,sha256=HROQMRw2g5sJTnuBagh3B0id3iWonRJ68tg3skOme9k,5748
 nucliadb/train/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -351,8 +351,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.3.1.post3590.dist-info/METADATA,sha256=cnyGLiXObJPzQkmY-I7Ilyq2TvPG9l48K1TrnPhqreQ,4291
-nucliadb-6.3.1.post3590.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-nucliadb-6.3.1.post3590.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.3.1.post3590.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.3.1.post3590.dist-info/RECORD,,
+nucliadb-6.3.1.post3598.dist-info/METADATA,sha256=66igSpA2Xq5tC8ibQhQSNAKLryF2gO_ZZxpOXHbET3I,4291
+nucliadb-6.3.1.post3598.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
+nucliadb-6.3.1.post3598.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.3.1.post3598.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.3.1.post3598.dist-info/RECORD,,