nucliadb 6.7.2.post4868__py3-none-any.whl → 6.7.2.post4874__py3-none-any.whl

migrations/0016_upgrade_to_paragraphs_v2.py

@@ -55,7 +55,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
 
 
 # async def has_old_paragraphs_index(context: ExecutionContext, kbid: str) -> bool:
-#     async with context.kv_driver.transaction(read_only=True) as txn:
+#     async with context.kv_driver.ro_transaction() as txn:
 #         shards_object = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=False)
 #         if not shards_object:
 #             raise ShardsObjectNotFound()

migrations/0017_multiple_writable_shards.py

@@ -44,7 +44,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
 
     # No longer relevant with nidx
 
-    # async with context.kv_driver.transaction() as txn:
+    # async with context.kv_driver.rw_transaction() as txn:
     #     shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=True)
     #     if shards is None:
     #         logger.error("KB without shards", extra={"kbid": kbid})

migrations/0018_purge_orphan_kbslugs.py

@@ -35,7 +35,7 @@ logger = logging.getLogger(__name__)
 
 
 async def migrate(context: ExecutionContext) -> None:
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         async for key in txn.keys(KB_SLUGS_BASE):
             slug = key.replace(KB_SLUGS_BASE, "")
             value = await txn.get(key, for_update=False)

migrations/0019_upgrade_to_paragraphs_v3.py

@@ -55,7 +55,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
 
 
 # async def has_old_paragraphs_index(context: ExecutionContext, kbid: str) -> bool:
-#     async with context.kv_driver.transaction(read_only=True) as txn:
+#     async with context.kv_driver.ro_transaction() as txn:
 #         shards_object = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
 #         if not shards_object:
 #             raise ShardsObjectNotFound()

migrations/0021_overwrite_vectorsets_key.py

@@ -38,7 +38,7 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         logger.info(f"Overwriting vectorsets key", extra={"kbid": kbid})
         await datamanagers.vectorsets.initialize(txn, kbid=kbid)
         await txn.commit()

migrations/0023_backfill_pg_catalog.py

@@ -44,7 +44,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
         return
 
     BATCH_SIZE = 100
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         txn = cast(PGTransaction, txn)
         continue_sql = ""
         while True:

migrations/0025_assign_models_to_kbs_v2.py

@@ -47,7 +47,7 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    async with context.kv_driver.transaction(read_only=True) as txn:
+    async with context.kv_driver.ro_transaction() as txn:
         vectorsets_count = len([vs async for vs in datamanagers.vectorsets.iter(txn, kbid=kbid)])
     if vectorsets_count > 0:
         logger.info("Skipping KB with vectorsets already populated", extra={"kbid": kbid})
@@ -65,7 +65,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
     learning_matryoshka_dimensions = learning_model_metadata.matryoshka_dimensions
     learning_normalize_vectors = len(learning_matryoshka_dimensions) > 0
 
-    async with context.kv_driver.transaction(read_only=True) as txn:
+    async with context.kv_driver.ro_transaction() as txn:
         semantic_model = await datamanagers.kb.get_model_metadata(txn, kbid=kbid)
 
         maindb_similarity = semantic_model.similarity_function
@@ -103,7 +103,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
         matryoshka_dimensions=maindb_matryoshka_dimensions,
     )
 
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         # Populate KB vectorsets with data from learning. We are skipping KBs
         # with this key already set, so we can set here safely
         await datamanagers.vectorsets.set(txn, kbid=kbid, config=default_vectorset)

migrations/0027_rollover_texts3.py

@@ -49,7 +49,7 @@ async def maybe_fix_vector_dimensions(context: ExecutionContext, kbid: str) -> N
         logger.warning(f"KB has no learning config", extra={"kbid": kbid})
         return
 
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         vectorsets = [vs async for vs in datamanagers.vectorsets.iter(txn, kbid=kbid)]
         if len(vectorsets) != 1:
             # If multiple vectorsets, they are new shards created correctly, we can safely skip it

migrations/0029_backfill_field_status.py

@@ -45,7 +45,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None: ...
 
 async def do_batch(context: ExecutionContext, start: str) -> Optional[str]:
     logger.info(f"Running batch from {start}")
-    async with context.kv_driver.transaction(read_only=False) as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         async with txn.connection.cursor() as cur:  # type: ignore
             # Retrieve a batch of fields
             await cur.execute(

migrations/0032_remove_old_relations.py

@@ -47,7 +47,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
 
 async def do_batch(context: ExecutionContext, start: str) -> Optional[str]:
     logger.info(f"Running batch from {start}")
-    async with context.kv_driver.transaction(read_only=False) as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         async with txn.connection.cursor() as cur:  # type: ignore
             # Retrieve a batch of fields
             await cur.execute(

migrations/0036_backfill_catalog_slug.py

@@ -37,7 +37,7 @@ async def migrate(context: ExecutionContext) -> None:
     driver = cast(PGDriver, context.kv_driver)
 
     BATCH_SIZE = 10_000
-    async with driver.transaction() as txn:
+    async with driver.rw_transaction() as txn:
         txn = cast(PGTransaction, txn)
         start_key = ""
         while True:

migrations/0037_backfill_catalog_facets.py

@@ -37,7 +37,7 @@ async def migrate(context: ExecutionContext) -> None:
     driver = cast(PGDriver, context.kv_driver)
 
     BATCH_SIZE = 1_000
-    async with driver.transaction() as txn:
+    async with driver.rw_transaction() as txn:
         txn = cast(PGTransaction, txn)
         start_kbid = "00000000000000000000000000000000"
         start_rid = "00000000000000000000000000000000"

migrations/0038_backfill_catalog_field_labels.py

@@ -45,7 +45,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
         return
 
     BATCH_SIZE = 100
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         txn = cast(PGTransaction, txn)
         start = ""
         while True:

nucliadb/backups/create.py

@@ -269,7 +269,7 @@ async def backup_search_configurations(context: ApplicationContext, kbid: str, b
 async def get_metadata(
     context: ApplicationContext, kbid: str, backup_id: str
 ) -> Optional[BackupMetadata]:
-    async with context.kv_driver.transaction(read_only=True) as txn:
+    async with context.kv_driver.ro_transaction() as txn:
         metadata_raw = await txn.get(MaindbKeys.METADATA.format(kbid=kbid, backup_id=backup_id))
         if metadata_raw is None:
             return None
@@ -277,7 +277,7 @@ async def get_metadata(
 
 
 async def set_metadata(context: ApplicationContext, kbid: str, backup_id: str, metadata: BackupMetadata):
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         await txn.set(
             MaindbKeys.METADATA.format(kbid=kbid, backup_id=backup_id),
             metadata.model_dump_json().encode(),
@@ -286,7 +286,7 @@ async def set_metadata(context: ApplicationContext, kbid: str, backup_id: str, m
 
 
 async def delete_metadata(context: ApplicationContext, kbid: str, backup_id: str):
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.rw_transaction() as txn:
         await txn.delete(MaindbKeys.METADATA.format(kbid=kbid, backup_id=backup_id))
         await txn.commit()
 

nucliadb/backups/restore.py

@@ -103,7 +103,7 @@ async def restore_resources(context: ApplicationContext, kbid: str, backup_id: s
 
 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:
+    async with context.kv_driver.ro_transaction() as txn:
         raw = await txn.get(key)
         if raw is None:
             return None
@@ -112,14 +112,14 @@ async def get_last_restored(context: ApplicationContext, kbid: str, backup_id: s
 
 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:
+    async with context.kv_driver.rw_transaction() as txn:
         await txn.set(key, resource_id.encode())
         await txn.commit()
 
 
 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:
+    async with context.kv_driver.rw_transaction() as txn:
         await txn.delete(key)
         await txn.commit()
 

nucliadb/common/cache.py

@@ -90,7 +90,7 @@ class ResourceCache(Cache[[str, str], ResourceORM]):
         @alru_cache(maxsize=cache_size)
         async def _get_resource(kbid: str, rid: str) -> Optional[ResourceORM]:
             storage = await get_storage()
-            async with get_driver().transaction(read_only=True) as txn:
+            async with get_driver().ro_transaction() as txn:
                 kb = KnowledgeBoxORM(txn, storage, kbid)
                 return await kb.get(rid)
 

nucliadb/common/datamanagers/utils.py

@@ -42,7 +42,7 @@ async def get_kv_pb(
 @contextlib.asynccontextmanager
 async def with_rw_transaction():
     driver = get_driver()
-    async with driver.transaction(read_only=False) as txn:
+    async with driver.rw_transaction() as txn:
         yield txn
 
 
@@ -53,5 +53,5 @@ with_transaction = with_rw_transaction
 @contextlib.asynccontextmanager
 async def with_ro_transaction():
     driver = get_driver()
-    async with driver.transaction(read_only=True) as ro_txn:
+    async with driver.ro_transaction() as ro_txn:
         yield ro_txn

nucliadb/common/locking.py

@@ -75,7 +75,7 @@ class _Lock:
         start = time.time()
         while True:
             try:
-                async with self.driver.transaction() as txn:
+                async with self.driver.rw_transaction() as txn:
                     lock_data = await self.get_lock_data(txn)
                     if lock_data is None:
                         await self._set_lock_value(txn)
@@ -128,7 +128,7 @@ class _Lock:
         while True:
             try:
                 await asyncio.sleep(self.refresh_timeout)
-                async with self.driver.transaction() as txn:
+                async with self.driver.rw_transaction() as txn:
                     await self._update_lock_value(txn)
                     await txn.commit()
             except (asyncio.CancelledError, RuntimeError):
@@ -138,12 +138,12 @@ class _Lock:
 
     async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
         self.task.cancel()
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             await txn.delete(self.key)
             await txn.commit()
 
     async def is_locked(self) -> bool:
-        async with get_driver().transaction(read_only=True) as txn:
+        async with get_driver().ro_transaction() as txn:
             lock_data = await self.get_lock_data(txn)
         return lock_data is not None and time.time() < lock_data.expires_at
 

nucliadb/common/maindb/driver.py

@@ -81,5 +81,15 @@ class Driver:
                     pass
 
     @asynccontextmanager
-    async def transaction(self, read_only: bool = False) -> AsyncGenerator[Transaction, None]:
+    async def _transaction(self, *, read_only: bool) -> AsyncGenerator[Transaction, None]:
         yield Transaction()
+
+    @asynccontextmanager
+    async def ro_transaction(self) -> AsyncGenerator[Transaction, None]:
+        async with self._transaction(read_only=True) as txn:
+            yield txn
+
+    @asynccontextmanager
+    async def rw_transaction(self) -> AsyncGenerator[Transaction, None]:
+        async with self._transaction(read_only=False) as txn:
+            yield txn

nucliadb/common/maindb/local.py

@@ -222,7 +222,7 @@ class LocalDriver(Driver):
         pass
 
     @asynccontextmanager
-    async def transaction(self, read_only: bool = False) -> AsyncGenerator[Transaction, None]:
+    async def _transaction(self, *, read_only: bool) -> AsyncGenerator[Transaction, None]:
         if self.url is None:
             raise AttributeError("Invalid url")
         txn = LocalTransaction(self.url, self)

nucliadb/common/maindb/pg.py

@@ -330,7 +330,7 @@ class PGDriver(Driver):
             metric.set(value)
 
     @asynccontextmanager
-    async def transaction(self, read_only: bool = False) -> AsyncGenerator[Transaction, None]:
+    async def _transaction(self, *, read_only: bool) -> AsyncGenerator[Transaction, None]:
         if read_only:
             yield ReadOnlyPGTransaction(self)
         else:

nucliadb/export_import/datamanager.py

@@ -54,7 +54,7 @@ class ExportImportDataManager:
 
     async def get_metadata(self, type: str, kbid: str, id: str) -> Metadata:
         key = self._get_maindb_metadata_key(type, kbid, id)
-        async with self.driver.transaction(read_only=True) as txn:
+        async with self.driver.ro_transaction() as txn:
             data = await txn.get(key)
         if data is None or data == b"":
             raise MetadataNotFound()
@@ -89,13 +89,13 @@ class ExportImportDataManager:
         metadata.modified = datetime.now(timezone.utc)
         key = self._get_maindb_metadata_key(type, metadata.kbid, metadata.id)
         data = metadata.model_dump_json().encode("utf-8")
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             await txn.set(key, data)
             await txn.commit()
 
     async def delete_metadata(self, type: str, metadata: Metadata):
         key = self._get_maindb_metadata_key(type, metadata.kbid, metadata.id)
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             await txn.delete(key)
             await txn.commit()
 

nucliadb/ingest/consumer/shard_creator.py

@@ -91,7 +91,7 @@ class ShardCreatorHandler:
     @metrics.handler_histo.wrap({"type": "shard_creator"})
     async def process_kb(self, kbid: str) -> None:
         logger.info({"message": "Processing notification for kbid", "kbid": kbid})
-        async with self.driver.transaction(read_only=True) as txn:
+        async with self.driver.ro_transaction() as txn:
             current_shard = await self.shard_manager.get_current_active_shard(txn, kbid)
 
         if current_shard is None:

nucliadb/ingest/orm/knowledgebox.py

@@ -127,7 +127,7 @@ class KnowledgeBox:
         rollback_ops: list[Callable[[], Coroutine[Any, Any, Any]]] = []
 
         try:
-            async with driver.transaction() as txn:
+            async with driver.rw_transaction() as txn:
                 exists = await datamanagers.kb.get_kb_uuid(
                     txn, slug=slug
                 ) or await datamanagers.kb.exists_kb(txn, kbid=kbid)
@@ -281,7 +281,7 @@ class KnowledgeBox:
 
     @classmethod
     async def delete(cls, driver: Driver, kbid: str):
-        async with driver.transaction() as txn:
+        async with driver.rw_transaction() as txn:
             exists = await datamanagers.kb.exists_kb(txn, kbid=kbid)
             if not exists:
                 return
@@ -347,7 +347,7 @@ class KnowledgeBox:
 
         nidx_api = get_nidx_api_client()
 
-        async with driver.transaction() as txn:
+        async with driver.rw_transaction() as txn:
             storage_to_delete = KB_TO_DELETE_STORAGE.format(kbid=kbid)
             await txn.set(storage_to_delete, b"")
 
@@ -377,7 +377,7 @@ class KnowledgeBox:
     @classmethod
     async def delete_all_kb_keys(cls, driver: Driver, kbid: str, chunk_size: int = 1_000):
         prefix = KB_KEYS.format(kbid=kbid)
-        async with driver.transaction() as txn:
+        async with driver.rw_transaction() as txn:
             await txn.delete_by_prefix(prefix)
             await txn.commit()
 

nucliadb/ingest/orm/processor/auditing.py

@@ -34,7 +34,7 @@ async def collect_audit_fields(
         return []
 
     audit_storage_fields: list[audit_pb2.AuditField] = []
-    async with driver.transaction(read_only=True) as txn:
+    async with driver.ro_transaction() as txn:
         kb = KnowledgeBox(txn, storage, message.kbid)
         resource = Resource(txn, storage, kb, message.uuid)
         field_keys = await resource.get_fields_ids()

nucliadb/ingest/orm/processor/processor.py

@@ -208,7 +208,7 @@ class Processor:
         partition: str,
         transaction_check: bool = True,
     ) -> None:
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             try:
                 kb = KnowledgeBox(txn, self.storage, message.kbid)
 
@@ -267,7 +267,7 @@ class Processor:
         # so we commit it in a different transaction to make it as short as possible
         prev_txn = resource.txn
         try:
-            async with self.driver.transaction() as txn:
+            async with self.driver.rw_transaction() as txn:
                 resource.txn = txn
                 await resource.set_slug()
                 await txn.commit()
@@ -294,7 +294,7 @@ class Processor:
                     await txn.commit()
             return None
 
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             try:
                 multi = messages[0].multiid
                 kb = KnowledgeBox(txn, self.storage, kbid)
@@ -731,7 +731,7 @@ class Processor:
             logger.info(f"Skip when resource does not even have basic metadata: {resource}")
             return
         try:
-            async with self.driver.transaction() as txn:
+            async with self.driver.rw_transaction() as txn:
                 kb.txn = resource.txn = txn
                 resource.basic.metadata.status = resources_pb2.Metadata.Status.ERROR
                 await resource.set_basic(resource.basic)

nucliadb/ingest/orm/processor/sequence_manager.py

@@ -31,7 +31,7 @@ async def get_last_seqid(driver: Driver, worker: str) -> Optional[int]:
     This is oriented towards the ingest consumer and processor,
     which is the only one that should be writing to this key.
     """
-    async with driver.transaction(read_only=True) as txn:
+    async with driver.ro_transaction() as txn:
         key = TXNID.format(worker=worker)
         last_seq = await txn.get(key)
         if not last_seq:

nucliadb/ingest/serialize.py

@@ -135,7 +135,7 @@ async def serialize(
     slug: Optional[str] = None,
 ) -> Optional[Resource]:
     driver = get_driver()
-    async with driver.transaction(read_only=True) as txn:
+    async with driver.ro_transaction() as txn:
         return await managed_serialize(
             txn,
             kbid,
@@ -392,6 +392,6 @@ async def get_resource_uuid_by_slug(
 ) -> Optional[str]:
     storage = await get_storage(service_name=service_name)
     driver = get_driver()
-    async with driver.transaction(read_only=True) as txn:
+    async with driver.ro_transaction() as txn:
         kb = KnowledgeBox(txn, storage, kbid)
         return await kb.get_resource_uuid_by_slug(slug)

nucliadb/ingest/service/writer.py

@@ -167,7 +167,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
             )
 
         try:
-            async with self.driver.transaction() as txn:
+            async with self.driver.rw_transaction() as txn:
                 kbid = await KnowledgeBoxORM.update(
                     txn, uuid=request.uuid, slug=request.slug, config=request.config
                 )
@@ -221,13 +221,13 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
         self, request: NewEntitiesGroupRequest, context=None
     ) -> NewEntitiesGroupResponse:
         response = NewEntitiesGroupResponse()
-        async with self.driver.transaction(read_only=True) as ro_txn:
+        async with self.driver.ro_transaction() as ro_txn:
             kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
             if kbobj is None:
                 response.status = NewEntitiesGroupResponse.Status.KB_NOT_FOUND
                 return response
 
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             kbobj.txn = txn
             entities_manager = EntitiesManager(kbobj, txn)
             try:
@@ -244,7 +244,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
         self, request: GetEntitiesRequest, context=None
     ) -> GetEntitiesResponse:
         response = GetEntitiesResponse()
-        async with self.driver.transaction(read_only=True) as txn:
+        async with self.driver.ro_transaction() as txn:
             kbobj = await self.proc.get_kb_obj(txn, request.kb)
             if kbobj is None:
                 response.status = GetEntitiesResponse.Status.NOTFOUND
@@ -266,7 +266,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
         self, request: ListEntitiesGroupsRequest, context=None
     ) -> ListEntitiesGroupsResponse:
         response = ListEntitiesGroupsResponse()
-        async with self.driver.transaction(read_only=True) as txn:
+        async with self.driver.ro_transaction() as txn:
             kbobj = await self.proc.get_kb_obj(txn, request.kb)
             if kbobj is None:
                 response.status = ListEntitiesGroupsResponse.Status.NOTFOUND
@@ -290,7 +290,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
         self, request: GetEntitiesGroupRequest, context=None
     ) -> GetEntitiesGroupResponse:
         response = GetEntitiesGroupResponse()
-        async with self.driver.transaction(read_only=True) as txn:
+        async with self.driver.ro_transaction() as txn:
             kbobj = await self.proc.get_kb_obj(txn, request.kb)
             if kbobj is None:
                 response.status = GetEntitiesGroupResponse.Status.KB_NOT_FOUND
@@ -315,13 +315,13 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
 
     async def SetEntities(self, request: SetEntitiesRequest, context=None) -> OpStatusWriter:  # type: ignore
         response = OpStatusWriter()
-        async with self.driver.transaction(read_only=True) as ro_txn:
+        async with self.driver.ro_transaction() as ro_txn:
             kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
             if kbobj is None:
                 response.status = OpStatusWriter.Status.NOTFOUND
                 return response
 
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             kbobj.txn = txn
             entities_manager = EntitiesManager(kbobj, txn)
             try:
@@ -339,13 +339,13 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
         self, request: UpdateEntitiesGroupRequest, context=None
     ) -> UpdateEntitiesGroupResponse:
         response = UpdateEntitiesGroupResponse()
-        async with self.driver.transaction(read_only=True) as ro_txn:
+        async with self.driver.ro_transaction() as ro_txn:
             kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
             if kbobj is None:
                 response.status = UpdateEntitiesGroupResponse.Status.KB_NOT_FOUND
                 return response
 
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             kbobj.txn = txn
             entities_manager = EntitiesManager(kbobj, txn)
             try:
@@ -368,13 +368,13 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
     async def DelEntities(self, request: DelEntitiesRequest, context=None) -> OpStatusWriter:  # type: ignore
         response = OpStatusWriter()
 
-        async with self.driver.transaction(read_only=True) as ro_txn:
+        async with self.driver.ro_transaction() as ro_txn:
             kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
             if kbobj is None:
                 response.status = OpStatusWriter.Status.NOTFOUND
                 return response
 
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             kbobj.txn = txn
             entities_manager = EntitiesManager(kbobj, txn)
             try:
@@ -393,7 +393,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
     ) -> WriterStatusResponse:
         logger.info("Status Call")
         response = WriterStatusResponse()
-        async with self.driver.transaction(read_only=True) as txn:
+        async with self.driver.ro_transaction() as txn:
             async for _, slug in datamanagers.kb.get_kbs(txn):
                 response.knowledgeboxes.append(slug)
 
@@ -405,7 +405,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
             return response
 
     async def Index(self, request: IndexResource, context=None) -> IndexStatus:  # type: ignore
-        async with self.driver.transaction() as txn:
+        async with self.driver.ro_transaction() as txn:
             kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
             resobj = ResourceORM(txn, self.storage, kbobj, request.rid)
             bm = await generate_broker_message(resobj)
@@ -419,7 +419,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
 
     async def ReIndex(self, request: IndexResource, context=None) -> IndexStatus:  # type: ignore
         try:
-            async with self.driver.transaction() as txn:
+            async with self.driver.rw_transaction() as txn:
                 kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
                 resobj = ResourceORM(txn, self.storage, kbobj, request.rid)
                 resobj.disable_vectors = not request.reindex_vectors

nucliadb/metrics_exporter.py

@@ -42,7 +42,7 @@ async def iter_kbids(context: ApplicationContext) -> AsyncGenerator[str, None]:
     """
     Return a list of all KB ids.
     """
-    async with context.kv_driver.transaction(read_only=True) as txn:
+    async with context.kv_driver.ro_transaction() as txn:
         async for kbid, _ in datamanagers.kb.get_kbs(txn):
             yield kbid