nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/common/datamanagers/resources.py

@@ -17,7 +17,7 @@
 # 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
 
@@ -26,13 +26,15 @@ 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.settings import settings as ingest_settings
-from nucliadb_protos import noderesources_pb2, resources_pb2, writer_pb2
+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
@@ -41,11 +43,16 @@ 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_SHARD = "/kbs/{kbid}/r/{uuid}/shard"
+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"
+
 
 async def resource_exists(txn: Transaction, *, kbid: str, rid: str) -> bool:
     basic = await get_basic_raw(txn, kbid=kbid, rid=rid)
@@ -55,10 +62,8 @@ async def resource_exists(txn: Transaction, *, kbid: str, rid: str) -> bool:
 # 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))
+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()
@@ -95,13 +100,11 @@ async def modify_slug(txn: Transaction, *, kbid: str, rid: str, new_slug: str) -
 # resource-shard
 
 
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3
-)
+@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
+    txn: Transaction, *, kbid: str, rid: str, for_update: bool = False
 ) -> Optional[str]:
-    shard = await txn.get(KB_RESOURCE_SHARD.format(kbid=kbid, uuid=rid))
+    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:
@@ -115,9 +118,7 @@ async def set_resource_shard_id(txn: Transaction, *, kbid: str, rid: str, shard:
 # Basic
 
 
-async def get_basic(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[resources_pb2.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
@@ -134,9 +135,7 @@ async def get_basic_raw(txn: Transaction, *, kbid: str, rid: str) -> Optional[by
     return raw_basic
 
 
-async def set_basic(
-    txn: Transaction, *, kbid: str, rid: str, basic: resources_pb2.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),
@@ -152,16 +151,12 @@ async def set_basic(
 # Origin
 
 
-async def get_origin(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[resources_pb2.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
-):
+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())
 
@@ -169,16 +164,12 @@ async def set_origin(
 # Extra
 
 
-async def get_extra(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[resources_pb2.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
-):
+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())
 
@@ -186,16 +177,12 @@ async def set_extra(
 # Security
 
 
-async def get_security(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[resources_pb2.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
-):
+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())
 
@@ -203,16 +190,12 @@ async def set_security(
 # Relations
 
 
-async def get_relations(
-    txn: Transaction, *, kbid: str, rid: str
-) -> Optional[resources_pb2.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
-):
+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())
 
@@ -241,25 +224,17 @@ 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
-)
+@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
-        ):
+    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]
 
 
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3
-)
+@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]
-        )
+    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]
 
 
@@ -288,26 +263,24 @@ 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())
 
 
 # Fields (materialized key with all field ids)
 
 
 async def get_all_field_ids(
-    txn: Transaction, *, kbid: str, rid: str
+    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)
+    return await get_kv_pb(txn, key, resources_pb2.AllFieldIDs, for_update=for_update)
 
 
 async def set_all_field_ids(
@@ -317,9 +290,7 @@ async def set_all_field_ids(
     await txn.set(key, allfields.SerializeToString())
 
 
-async def has_field(
-    txn: Transaction, *, kbid: str, rid: str, field_id: resources_pb2.FieldID
-) -> bool:
+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
@@ -332,25 +303,8 @@ async def has_field(
 # ORM mix (this functions shouldn't belong here)
 
 
-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
-
-    resource.disable_vectors = False
-    resource.txn = txn
-    bm = await resource.generate_broker_message()
-    return bm
-
-
-@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]:
+@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.
@@ -362,19 +316,3 @@ async def get_resource(
 
     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_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

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)
 
@@ -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/common/datamanagers/synonyms.py

@@ -29,7 +29,7 @@ KB_SYNONYMS = "/kbs/{kbid}/synonyms"
 
 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)
+    return await get_kv_pb(txn, key, knowledgebox_pb2.Synonyms, for_update=False)
 
 
 async def set(txn: Transaction, *, kbid: str, synonyms: knowledgebox_pb2.Synonyms):

nucliadb/common/datamanagers/utils.py

@@ -29,9 +29,9 @@ 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]:
-    serialized: Optional[bytes] = await txn.get(key)
+    serialized: Optional[bytes] = await txn.get(key, for_update=for_update)
     if serialized is None:
         return None
     pb = pb_type()
@@ -40,9 +40,18 @@ async def get_kv_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

nucliadb/common/datamanagers/vectorsets.py

@@ -0,0 +1,110 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 AsyncIterator, Optional
+
+from nucliadb.common.datamanagers.utils import get_kv_pb
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb_protos import knowledgebox_pb2
+
+KB_VECTORSETS = "/kbs/{kbid}/vectorsets"
+
+
+class BrokenInvariant(Exception):
+    pass
+
+
+async def initialize(txn: Transaction, *, kbid: str):
+    key = KB_VECTORSETS.format(kbid=kbid)
+    await txn.set(key, knowledgebox_pb2.KnowledgeBoxVectorSetsConfig().SerializeToString())
+
+
+async def get(
+    txn: Transaction, *, kbid: str, vectorset_id: str
+) -> Optional[knowledgebox_pb2.VectorSetConfig]:
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
+    index = _find_vectorset(kb_vectorsets, vectorset_id)
+    if index is None:
+        return None
+    return kb_vectorsets.vectorsets[index]
+
+
+async def exists(txn, *, kbid: str, vectorset_id: str) -> bool:
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
+    return _find_vectorset(kb_vectorsets, vectorset_id) is not None
+
+
+async def iter(
+    txn: Transaction, *, kbid: str
+) -> AsyncIterator[tuple[str, knowledgebox_pb2.VectorSetConfig]]:
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
+    for config in kb_vectorsets.vectorsets:
+        yield config.vectorset_id, config
+
+
+async def set(txn: Transaction, *, kbid: str, config: knowledgebox_pb2.VectorSetConfig):
+    """Create or update a vectorset configuration"""
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=True)
+    index = _find_vectorset(kb_vectorsets, config.vectorset_id)
+    if index is None:
+        # adding a new vectorset
+        kb_vectorsets.vectorsets.append(config)
+    else:
+        # updating a vectorset
+        kb_vectorsets.vectorsets[index].CopyFrom(config)
+
+    key = KB_VECTORSETS.format(kbid=kbid)
+    await txn.set(key, kb_vectorsets.SerializeToString())
+
+
+async def delete(txn: Transaction, *, kbid: str, vectorset_id: str):
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=True)
+    index = _find_vectorset(kb_vectorsets, vectorset_id)
+    if index is None:
+        # already deleted
+        return
+
+    del kb_vectorsets.vectorsets[index]
+    key = KB_VECTORSETS.format(kbid=kbid)
+    await txn.set(key, kb_vectorsets.SerializeToString())
+
+
+# XXX At some point in the vectorset epic, we should make this key mandatory and
+# fail instead of providing a default
+async def _get_or_default(
+    txn: Transaction,
+    *,
+    kbid: str,
+    for_update: bool = True,
+) -> knowledgebox_pb2.KnowledgeBoxVectorSetsConfig:
+    key = KB_VECTORSETS.format(kbid=kbid)
+    stored = await get_kv_pb(
+        txn, key, knowledgebox_pb2.KnowledgeBoxVectorSetsConfig, for_update=for_update
+    )
+    return stored or knowledgebox_pb2.KnowledgeBoxVectorSetsConfig()
+
+
+def _find_vectorset(
+    kb_vectorsets: knowledgebox_pb2.KnowledgeBoxVectorSetsConfig, vectorset_id: str
+) -> Optional[int]:
+    """Return the position of the vectorset in `vectorsets` or `None` if not found."""
+    for idx, vectorset in enumerate(kb_vectorsets.vectorsets):
+        if vectorset.vectorset_id == vectorset_id:
+            return idx
+    return None