nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/common/datamanagers/resources.py

@@ -17,47 +17,191 @@
 # 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, Optional
+from typing import TYPE_CHECKING, AsyncGenerator, Optional
 
 import backoff
-from nucliadb_protos.resources_pb2 import Basic
 
+from nucliadb.common.datamanagers.utils import get_kv_pb
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb.common.maindb.exceptions import ConflictError, NotFoundError
 
 # These should be refactored
-from nucliadb.ingest.orm.resource import KB_RESOURCE_SLUG, KB_RESOURCE_SLUG_BASE
-from nucliadb.ingest.orm.resource import Resource as ResourceORM
-from nucliadb.ingest.orm.utils import get_basic, set_basic
-from nucliadb_protos import noderesources_pb2, writer_pb2
+from nucliadb.ingest.settings import settings as ingest_settings
+from nucliadb_protos import resources_pb2
 from nucliadb_utils.utilities import get_storage
 
-from .utils import with_transaction
+from .utils import with_ro_transaction
 
+if TYPE_CHECKING:
+    from nucliadb.ingest.orm.resource import Resource as ResourceORM
+
+
+KB_RESOURCE_BASIC = "/kbs/{kbid}/r/{uuid}"
+KB_RESOURCE_BASIC_FS = "/kbs/{kbid}/r/{uuid}/basic"  # Only used on FS driver
+KB_RESOURCE_ORIGIN = "/kbs/{kbid}/r/{uuid}/origin"
+KB_RESOURCE_EXTRA = "/kbs/{kbid}/r/{uuid}/extra"
+KB_RESOURCE_SECURITY = "/kbs/{kbid}/r/{uuid}/security"
+KB_RESOURCE_RELATIONS = "/kbs/{kbid}/r/{uuid}/relations"
+
+KB_RESOURCE_SLUG_BASE = "/kbs/{kbid}/s/"
+KB_RESOURCE_SLUG = f"{KB_RESOURCE_SLUG_BASE}{{slug}}"
+
+KB_RESOURCE_FIELDS = "/kbs/{kbid}/r/{uuid}/f/"
+
+KB_RESOURCE_ALL_FIELDS = "/kbs/{kbid}/r/{uuid}/allfields"
 KB_MATERIALIZED_RESOURCES_COUNT = "/kbs/{kbid}/materialized/resources/count"
+
 KB_RESOURCE_SHARD = "/kbs/{kbid}/r/{uuid}/shard"
 
 
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3
-)
-async def _iter_resource_slugs(*, kbid: str) -> AsyncGenerator[str, None]:
-    async with with_transaction() as txn:
-        async for key in txn.keys(
-            match=KB_RESOURCE_SLUG_BASE.format(kbid=kbid), count=-1
-        ):
-            yield key.split("/")[-1]
+async def resource_exists(txn: Transaction, *, kbid: str, rid: str) -> bool:
+    basic = await get_basic_raw(txn, kbid=kbid, rid=rid)
+    return basic is not None
 
 
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3
-)
-async def _get_resource_ids_from_slugs(kbid: str, slugs: list[str]) -> list[str]:
-    async with with_transaction() as txn:
-        rids = await txn.batch_get(
-            [KB_RESOURCE_SLUG.format(kbid=kbid, slug=slug) for slug in slugs]
+# id and slug
+
+
+async def get_resource_uuid_from_slug(txn: Transaction, *, kbid: str, slug: str) -> Optional[str]:
+    encoded_uuid = await txn.get(KB_RESOURCE_SLUG.format(kbid=kbid, slug=slug, for_update=False))
+    if not encoded_uuid:
+        return None
+    return encoded_uuid.decode()
+
+
+async def slug_exists(txn: Transaction, *, kbid: str, slug: str) -> bool:
+    key = KB_RESOURCE_SLUG.format(kbid=kbid, slug=slug)
+    encoded_slug: Optional[bytes] = await txn.get(key)
+    return encoded_slug not in (None, b"")
+
+
+async def modify_slug(txn: Transaction, *, kbid: str, rid: str, new_slug: str) -> str:
+    basic = await get_basic(txn, kbid=kbid, rid=rid)
+    if basic is None:
+        raise NotFoundError()
+    old_slug = basic.slug
+
+    uuid_for_new_slug = await get_resource_uuid_from_slug(txn, kbid=kbid, slug=new_slug)
+    if uuid_for_new_slug is not None:
+        if uuid_for_new_slug == rid:
+            # Nothing to change
+            return old_slug
+        else:
+            raise ConflictError(f"Slug {new_slug} already exists")
+    key = KB_RESOURCE_SLUG.format(kbid=kbid, slug=old_slug)
+    await txn.delete(key)
+    key = KB_RESOURCE_SLUG.format(kbid=kbid, slug=new_slug)
+    await txn.set(key, rid.encode())
+    basic.slug = new_slug
+    await set_basic(txn, kbid=kbid, rid=rid, basic=basic)
+    return old_slug
+
+
+# resource-shard
+
+
+@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
+async def get_resource_shard_id(
+    txn: Transaction, *, kbid: str, rid: str, for_update: bool = False
+) -> Optional[str]:
+    shard = await txn.get(KB_RESOURCE_SHARD.format(kbid=kbid, uuid=rid, for_update=for_update))
+    if shard is not None:
+        return shard.decode()
+    else:
+        return None
+
+
+async def set_resource_shard_id(txn: Transaction, *, kbid: str, rid: str, shard: str):
+    await txn.set(KB_RESOURCE_SHARD.format(kbid=kbid, uuid=rid), shard.encode())
+
+
+# Basic
+
+
+async def get_basic(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Basic]:
+    raw = await get_basic_raw(txn, kbid=kbid, rid=rid)
+    if raw is None:
+        return None
+    basic = resources_pb2.Basic()
+    basic.ParseFromString(raw)
+    return basic
+
+
+async def get_basic_raw(txn: Transaction, *, kbid: str, rid: str) -> Optional[bytes]:
+    if ingest_settings.driver == "local":
+        raw_basic = await txn.get(KB_RESOURCE_BASIC_FS.format(kbid=kbid, uuid=rid))
+    else:
+        raw_basic = await txn.get(KB_RESOURCE_BASIC.format(kbid=kbid, uuid=rid))
+    return raw_basic
+
+
+async def set_basic(txn: Transaction, *, kbid: str, rid: str, basic: resources_pb2.Basic):
+    if ingest_settings.driver == "local":
+        await txn.set(
+            KB_RESOURCE_BASIC_FS.format(kbid=kbid, uuid=rid),
+            basic.SerializeToString(),
         )
-    return [rid.decode() for rid in rids if rid is not None]
+    else:
+        await txn.set(
+            KB_RESOURCE_BASIC.format(kbid=kbid, uuid=rid),
+            basic.SerializeToString(),
+        )
+
+
+# Origin
+
+
+async def get_origin(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Origin]:
+    key = KB_RESOURCE_ORIGIN.format(kbid=kbid, uuid=rid)
+    return await get_kv_pb(txn, key, resources_pb2.Origin)
+
+
+async def set_origin(txn: Transaction, *, kbid: str, rid: str, origin: resources_pb2.Origin):
+    key = KB_RESOURCE_ORIGIN.format(kbid=kbid, uuid=rid)
+    await txn.set(key, origin.SerializeToString())
+
+
+# Extra
+
+
+async def get_extra(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Extra]:
+    key = KB_RESOURCE_EXTRA.format(kbid=kbid, uuid=rid)
+    return await get_kv_pb(txn, key, resources_pb2.Extra)
+
+
+async def set_extra(txn: Transaction, *, kbid: str, rid: str, extra: resources_pb2.Extra):
+    key = KB_RESOURCE_EXTRA.format(kbid=kbid, uuid=rid)
+    await txn.set(key, extra.SerializeToString())
+
+
+# Security
+
+
+async def get_security(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Security]:
+    key = KB_RESOURCE_SECURITY.format(kbid=kbid, uuid=rid)
+    return await get_kv_pb(txn, key, resources_pb2.Security)
+
+
+async def set_security(txn: Transaction, *, kbid: str, rid: str, security: resources_pb2.Security):
+    key = KB_RESOURCE_SECURITY.format(kbid=kbid, uuid=rid)
+    await txn.set(key, security.SerializeToString())
+
+
+# Relations
+
+
+async def get_relations(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Relations]:
+    key = KB_RESOURCE_RELATIONS.format(kbid=kbid, uuid=rid)
+    return await get_kv_pb(txn, key, resources_pb2.Relations)
+
+
+async def set_relations(txn: Transaction, *, kbid: str, rid: str, relations: resources_pb2.Relations):
+    key = KB_RESOURCE_RELATIONS.format(kbid=kbid, uuid=rid)
+    await txn.set(key, relations.SerializeToString())
+
+
+# KB resource ids (this functions use internal transactions, breaking the
+# datamanager contract. We should rethink them at some point)
 
 
 async def iterate_resource_ids(*, kbid: str) -> AsyncGenerator[str, None]:
@@ -80,52 +224,21 @@ async def iterate_resource_ids(*, kbid: str) -> AsyncGenerator[str, None]:
             yield rid
 
 
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3
-)
-async def get_resource_shard_id(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[str]:
-    shard = await txn.get(KB_RESOURCE_SHARD.format(kbid=kbid, uuid=rid))
-    if shard is not None:
-        return shard.decode()
-    else:
-        return None
-
-
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3
-)
-async def get_resource(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[ResourceORM]:
-    """
-    Not ideal to return Resource type here but refactoring would
-    require a lot of changes.
-
-    At least this isolated that dependency here.
-    """
-    # prevent circulat imports -- this is not ideal that we have the ORM mix here.
-    from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
+@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
+async def _iter_resource_slugs(*, kbid: str) -> AsyncGenerator[str, None]:
+    async with with_ro_transaction() as txn:
+        async for key in txn.keys(match=KB_RESOURCE_SLUG_BASE.format(kbid=kbid)):
+            yield key.split("/")[-1]
 
-    kb_orm = KnowledgeBoxORM(txn, await get_storage(), kbid)
-    return await kb_orm.get(rid)
 
+@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
+async def _get_resource_ids_from_slugs(kbid: str, slugs: list[str]) -> list[str]:
+    async with with_ro_transaction() as txn:
+        rids = await txn.batch_get([KB_RESOURCE_SLUG.format(kbid=kbid, slug=slug) for slug in slugs])
+    return [rid.decode() for rid in rids if rid is not None]
 
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3
-)
-async def get_resource_index_message(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[noderesources_pb2.Resource]:
-    # prevent circulat imports -- this is not ideal that we have the ORM mix here.
-    from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
 
-    kb_orm = KnowledgeBoxORM(txn, await get_storage(), kbid)
-    res = await kb_orm.get(rid)
-    if res is None:
-        return None
-    return (await res.generate_index_message()).brain
+# KB resource count (materialized key)
 
 
 async def calculate_number_of_resources(txn: Transaction, *, kbid: str) -> int:
@@ -150,72 +263,56 @@ async def get_number_of_resources(txn: Transaction, *, kbid: str) -> int:
     """
     Return cached number of resources in a knowledgebox.
     """
-    raw_value = await txn.get(KB_MATERIALIZED_RESOURCES_COUNT.format(kbid=kbid))
+    raw_value = await txn.get(KB_MATERIALIZED_RESOURCES_COUNT.format(kbid=kbid), for_update=False)
     if raw_value is None:
         return -1
     return int(raw_value)
 
 
 async def set_number_of_resources(txn: Transaction, kbid: str, value: int) -> None:
-    await txn.set(
-        KB_MATERIALIZED_RESOURCES_COUNT.format(kbid=kbid), str(value).encode()
-    )
+    await txn.set(KB_MATERIALIZED_RESOURCES_COUNT.format(kbid=kbid), str(value).encode())
 
 
-async def get_broker_message(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[writer_pb2.BrokerMessage]:
-    resource = await get_resource(txn, kbid=kbid, rid=rid)
-    if resource is None:
-        return None
+# Fields (materialized key with all field ids)
 
-    resource.disable_vectors = False
-    resource.txn = txn
-    bm = await resource.generate_broker_message()
-    return bm
 
+async def get_all_field_ids(
+    txn: Transaction, *, kbid: str, rid: str, for_update: bool = False
+) -> Optional[resources_pb2.AllFieldIDs]:
+    key = KB_RESOURCE_ALL_FIELDS.format(kbid=kbid, uuid=rid)
+    return await get_kv_pb(txn, key, resources_pb2.AllFieldIDs, for_update=for_update)
 
-async def get_resource_basic(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[Basic]:
-    raw_basic = await get_basic(txn, kbid, rid)
-    if not raw_basic:
-        return None
-    basic = Basic()
-    basic.ParseFromString(raw_basic)
-    return basic
 
+async def set_all_field_ids(
+    txn: Transaction, *, kbid: str, rid: str, allfields: resources_pb2.AllFieldIDs
+):
+    key = KB_RESOURCE_ALL_FIELDS.format(kbid=kbid, uuid=rid)
+    await txn.set(key, allfields.SerializeToString())
 
-async def get_resource_uuid_from_slug(
-    txn: Transaction, *, kbid: str, slug: str
-) -> Optional[str]:
-    encoded_uuid = await txn.get(KB_RESOURCE_SLUG.format(kbid=kbid, slug=slug))
-    if not encoded_uuid:
-        return None
-    return encoded_uuid.decode()
 
+async def has_field(txn: Transaction, *, kbid: str, rid: str, field_id: resources_pb2.FieldID) -> bool:
+    fields = await get_all_field_ids(txn, kbid=kbid, rid=rid)
+    if fields is None:
+        return False
+    for resource_field_id in fields.fields:
+        if field_id == resource_field_id:
+            return True
+    return False
 
-async def modify_slug(txn: Transaction, *, kbid: str, rid: str, new_slug: str) -> str:
-    basic = await get_resource_basic(txn, kbid=kbid, rid=rid)
-    if basic is None:
-        raise NotFoundError()
-    old_slug = basic.slug
 
-    uuid_for_new_slug = await get_resource_uuid_from_slug(txn, kbid=kbid, slug=new_slug)
-    if uuid_for_new_slug is not None:
-        if uuid_for_new_slug == rid:
-            # Nothing to change
-            return old_slug
-        else:
-            raise ConflictError(f"Slug {new_slug} already exists")
-    key = KB_RESOURCE_SLUG.format(kbid=kbid, slug=old_slug)
-    await txn.delete(key)
-    key = KB_RESOURCE_SLUG.format(kbid=kbid, slug=new_slug)
-    await txn.set(key, rid.encode())
-    basic.slug = new_slug
-    await set_basic(txn, kbid, rid, basic)
-    return old_slug
+# ORM mix (this functions shouldn't belong here)
 
 
-async def set_resource_shard_id(txn: Transaction, *, kbid: str, rid: str, shard: str):
-    await txn.set(KB_RESOURCE_SHARD.format(kbid=kbid, uuid=rid), shard.encode())
+@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
+async def get_resource(txn: Transaction, *, kbid: str, rid: str) -> Optional["ResourceORM"]:
+    """
+    Not ideal to return Resource type here but refactoring would
+    require a lot of changes.
+
+    At least this isolated that dependency here.
+    """
+    # prevent circulat imports -- this is not ideal that we have the ORM mix here.
+    from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
+
+    kb_orm = KnowledgeBoxORM(txn, await get_storage(), kbid)
+    return await kb_orm.get(rid)

nucliadb/common/datamanagers/rollover.py

@@ -21,22 +21,42 @@ import logging
 from typing import AsyncGenerator, Optional
 
 import orjson
+from pydantic import BaseModel
 
 from nucliadb.common.maindb.driver import Transaction
+from nucliadb_protos import knowledgebox_pb2 as kb_pb2
 from nucliadb_protos import writer_pb2
 
-from .utils import get_kv_pb, with_transaction
+from .utils import get_kv_pb, with_ro_transaction
 
 logger = logging.getLogger(__name__)
 
+KB_ROLLOVER_STATE = "/kbs/{kbid}/rollover/state"
 KB_ROLLOVER_SHARDS = "/kbs/{kbid}/rollover/shards"
+KB_ROLLOVER_EXTERNAL_INDEX_METADATA = "/kbs/{kbid}/rollover/external_index_metadata"
 KB_ROLLOVER_RESOURCES_TO_INDEX = "/kbs/{kbid}/rollover/to-index/{resource}"
 KB_ROLLOVER_RESOURCES_INDEXED = "/kbs/{kbid}/rollover/indexed/{resource}"
 
 
-async def get_kb_rollover_shards(
-    txn: Transaction, *, kbid: str
-) -> Optional[writer_pb2.Shards]:
+class RolloverState(BaseModel):
+    rollover_shards_created: bool = False
+    external_index_created: bool = False
+    resources_scheduled: bool = False
+    resources_indexed: bool = False
+    cutover_shards: bool = False
+    cutover_external_index: bool = False
+    resources_validated: bool = False
+
+
+class RolloverStateNotFoundError(Exception):
+    """
+    Raised when the rollover state is not found.
+    """
+
+    ...
+
+
+async def get_kb_rollover_shards(txn: Transaction, *, kbid: str) -> Optional[writer_pb2.Shards]:
     key = KB_ROLLOVER_SHARDS.format(kbid=kbid)
     return await get_kv_pb(txn, key, writer_pb2.Shards)
 
@@ -90,7 +110,7 @@ async def add_indexed(
     kbid: str,
     resource_id: str,
     shard_id: str,
-    modification_time: int
+    modification_time: int,
 ) -> None:
     to_index = KB_ROLLOVER_RESOURCES_TO_INDEX.format(kbid=kbid, resource=resource_id)
     indexed = KB_ROLLOVER_RESOURCES_INDEXED.format(kbid=kbid, resource=resource_id)
@@ -106,7 +126,9 @@ async def get_indexed_data(
     val = await txn.get(key)
     if val is not None:
         data = orjson.loads(val)
-        return tuple(data)  # type: ignore
+        shard_id: str = data[0]
+        modification_time: int = data[1]
+        return shard_id, modification_time
     return None
 
 
@@ -122,15 +144,13 @@ async def iter_indexed_keys(*, kbid: str) -> AsyncGenerator[str, None]:
     internally managed
     """
     start_key = KB_ROLLOVER_RESOURCES_INDEXED.format(kbid=kbid, resource="")
-    async with with_transaction() as txn:
-        async for key in txn.keys(match=start_key, count=-1):
+    async with with_ro_transaction() as txn:
+        async for key in txn.keys(match=start_key):
             yield key.split("/")[-1]
 
 
-async def _get_batch_indexed_data(
-    *, kbid, batch: list[str]
-) -> list[tuple[str, tuple[str, int]]]:
-    async with with_transaction() as txn:
+async def _get_batch_indexed_data(*, kbid, batch: list[str]) -> list[tuple[str, tuple[str, int]]]:
+    async with with_ro_transaction() as txn:
         values = await txn.batch_get(
             [
                 KB_ROLLOVER_RESOURCES_INDEXED.format(kbid=kbid, resource=resource_id)
@@ -140,14 +160,15 @@ async def _get_batch_indexed_data(
     results: list[tuple[str, tuple[str, int]]] = []
     for key, val in zip(batch, values):
         if val is not None:
-            data: tuple[str, int] = tuple(orjson.loads(val))  # type: ignore
+            shard_id: str
+            modification_time: int
+            shard_id, modification_time = orjson.loads(val)
+            data = (shard_id, modification_time)
             results.append((key.split("/")[-1], data))
     return results
 
 
-async def iterate_indexed_data(
-    *, kbid: str
-) -> AsyncGenerator[tuple[str, tuple[str, int]], None]:
+async def iterate_indexed_data(*, kbid: str) -> AsyncGenerator[tuple[str, tuple[str, int]], None]:
     """
     This function is optimized for reducing the time a transaction is open.
 
@@ -164,3 +185,43 @@ async def iterate_indexed_data(
     if len(batch) > 0:
         for key, val in await _get_batch_indexed_data(kbid=kbid, batch=batch):
             yield key, val
+
+
+async def get_rollover_state(txn: Transaction, kbid: str) -> RolloverState:
+    key = KB_ROLLOVER_STATE.format(kbid=kbid)
+    val = await txn.get(key)
+    if not val:
+        raise RolloverStateNotFoundError(kbid)
+    return RolloverState.model_validate_json(val)
+
+
+async def set_rollover_state(txn: Transaction, kbid: str, state: RolloverState) -> None:
+    key = KB_ROLLOVER_STATE.format(kbid=kbid)
+    await txn.set(key, state.model_dump_json().encode())
+
+
+async def clear_rollover_state(txn: Transaction, kbid: str) -> None:
+    key = KB_ROLLOVER_STATE.format(kbid=kbid)
+    await txn.delete(key)
+
+
+async def update_kb_rollover_external_index_metadata(
+    txn: Transaction, *, kbid: str, metadata: kb_pb2.StoredExternalIndexProviderMetadata
+) -> None:
+    key = KB_ROLLOVER_EXTERNAL_INDEX_METADATA.format(kbid=kbid)
+    await txn.set(key, metadata.SerializeToString())
+
+
+async def get_kb_rollover_external_index_metadata(
+    txn: Transaction, *, kbid: str
+) -> Optional[kb_pb2.StoredExternalIndexProviderMetadata]:
+    key = KB_ROLLOVER_EXTERNAL_INDEX_METADATA.format(kbid=kbid)
+    val = await txn.get(key)
+    if not val:
+        return None
+    return kb_pb2.StoredExternalIndexProviderMetadata.FromString(val)
+
+
+async def delete_kb_rollover_external_index_metadata(txn: Transaction, *, kbid: str) -> None:
+    key = KB_ROLLOVER_EXTERNAL_INDEX_METADATA.format(kbid=kbid)
+    await txn.delete(key)

nucliadb/{ingest/orm → common/datamanagers}/synonyms.py

@@ -17,38 +17,26 @@
 # 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 Optional
 
-from nucliadb_protos.knowledgebox_pb2 import Synonyms as PBSynonyms
+from typing import Optional
 
+from nucliadb.common.datamanagers.utils import get_kv_pb
 from nucliadb.common.maindb.driver import Transaction
+from nucliadb_protos import knowledgebox_pb2
 
 KB_SYNONYMS = "/kbs/{kbid}/synonyms"
 
 
-class Synonyms:
-    def __init__(self, txn: Transaction, kbid: str):
-        self.txn = txn
-        self.kbid = kbid
-
-    @property
-    def key(self) -> str:
-        return KB_SYNONYMS.format(kbid=self.kbid)
-
-    async def set(self, synonyms: PBSynonyms):
-        body = synonyms.SerializeToString()
-        await self.txn.set(self.key, body)
-
-    async def get(self) -> Optional[PBSynonyms]:
-        try:
-            payload = await self.txn.get(self.key)
-        except KeyError:
-            return None
-        if payload is None:
-            return None
-        body = PBSynonyms()
-        body.ParseFromString(payload)
-        return body
-
-    async def clear(self):
-        await self.txn.delete(self.key)
+async def get(txn: Transaction, *, kbid: str) -> Optional[knowledgebox_pb2.Synonyms]:
+    key = KB_SYNONYMS.format(kbid=kbid)
+    return await get_kv_pb(txn, key, knowledgebox_pb2.Synonyms, for_update=False)
+
+
+async def set(txn: Transaction, *, kbid: str, synonyms: knowledgebox_pb2.Synonyms):
+    key = KB_SYNONYMS.format(kbid=kbid)
+    await txn.set(key, synonyms.SerializeToString())
+
+
+async def delete(txn: Transaction, *, kbid: str):
+    key = KB_SYNONYMS.format(kbid=kbid)
+    await txn.delete(key)

nucliadb/common/datamanagers/utils.py

@@ -29,21 +29,29 @@ PB_TYPE = TypeVar("PB_TYPE", bound=Message)
 
 
 async def get_kv_pb(
-    txn: Transaction, key: str, pb_type: Type[PB_TYPE]
+    txn: Transaction, key: str, pb_type: Type[PB_TYPE], for_update: bool = True
 ) -> Optional[PB_TYPE]:
-    kb_shards_bytes: Optional[bytes] = await txn.get(key)
-    if kb_shards_bytes:
-        kb_shards = pb_type()
-        kb_shards.ParseFromString(kb_shards_bytes)
-        return kb_shards
-    else:
+    serialized: Optional[bytes] = await txn.get(key, for_update=for_update)
+    if serialized is None:
         return None
+    pb = pb_type()
+    pb.ParseFromString(serialized)
+    return pb
 
 
 @contextlib.asynccontextmanager
-async def with_transaction(read_only: bool = False, wait_for_abort: bool = True):
+async def with_rw_transaction():
     driver = get_driver()
-    async with driver.transaction(
-        read_only=read_only, wait_for_abort=wait_for_abort
-    ) as txn:
+    async with driver.transaction(read_only=False) as txn:
         yield txn
+
+
+# For backwards compatibility
+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:
+        yield ro_txn