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

nucliadb/ingest/orm/knowledgebox.py

@@ -18,55 +18,71 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from datetime import datetime
-from typing import AsyncGenerator, AsyncIterator, Optional, Sequence
+from functools import partial
+from typing import Any, AsyncGenerator, Callable, Coroutine, Optional, Sequence
 from uuid import uuid4
 
 from grpc import StatusCode
 from grpc.aio import AioRpcError
-from nucliadb_protos.knowledgebox_pb2 import (
-    KnowledgeBoxConfig,
-    Labels,
-    LabelSet,
-    SemanticModelMetadata,
-)
-from nucliadb_protos.knowledgebox_pb2 import Synonyms as PBSynonyms
-from nucliadb_protos.knowledgebox_pb2 import VectorSet, VectorSets
-from nucliadb_protos.resources_pb2 import Basic
-from nucliadb_protos.utils_pb2 import ReleaseChannel
 
 from nucliadb.common import datamanagers
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.common.cluster.exceptions import ShardNotFound, ShardsNotFound
+from nucliadb.common.cluster.exceptions import ShardNotFound
 from nucliadb.common.cluster.manager import get_index_node
 from nucliadb.common.cluster.utils import get_shard_manager
-from nucliadb.common.maindb.driver import Driver, Transaction
-from nucliadb.ingest import SERVICE_NAME, logger
-from nucliadb.ingest.orm.exceptions import KnowledgeBoxConflict
-from nucliadb.ingest.orm.resource import (
+
+# XXX: this keys shouldn't be exposed outside datamanagers
+from nucliadb.common.datamanagers.resources import (
     KB_RESOURCE_SLUG,
     KB_RESOURCE_SLUG_BASE,
-    Resource,
 )
-from nucliadb.ingest.orm.synonyms import Synonyms
-from nucliadb.ingest.orm.utils import compute_paragraph_key, get_basic, set_basic
+from nucliadb.common.external_index_providers.base import VectorsetExternalIndex
+from nucliadb.common.external_index_providers.pinecone import PineconeIndexManager
+from nucliadb.common.maindb.driver import Driver, Transaction
+from nucliadb.common.maindb.pg import PGTransaction
+from nucliadb.common.nidx import get_nidx_api_client
+from nucliadb.ingest import SERVICE_NAME, logger
+from nucliadb.ingest.orm.exceptions import (
+    KnowledgeBoxConflict,
+    KnowledgeBoxCreationError,
+    VectorSetConflict,
+)
+from nucliadb.ingest.orm.metrics import processor_observer
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.ingest.orm.utils import choose_matryoshka_dimension, compute_paragraph_key
 from nucliadb.migrator.utils import get_latest_version
-from nucliadb_protos import writer_pb2
+from nucliadb_protos import knowledgebox_pb2, noderesources_pb2, nodewriter_pb2, writer_pb2
+from nucliadb_protos.knowledgebox_pb2 import (
+    CreateExternalIndexProviderMetadata,
+    ExternalIndexProviderType,
+    KnowledgeBoxConfig,
+    SemanticModelMetadata,
+    StoredExternalIndexProviderMetadata,
+)
+from nucliadb_protos.resources_pb2 import Basic
+from nucliadb_utils.settings import is_onprem_nucliadb
 from nucliadb_utils.storages.storage import Storage
-from nucliadb_utils.utilities import get_audit, get_storage
+from nucliadb_utils.utilities import (
+    get_audit,
+    get_storage,
+)
 
 # XXX Eventually all these keys should be moved to datamanagers.kb
 KB_RESOURCE = "/kbs/{kbid}/r/{uuid}"
 
 KB_KEYS = "/kbs/{kbid}/"
 
-KB_VECTORSET = "/kbs/{kbid}/vectorsets"
-
 KB_TO_DELETE_BASE = "/kbtodelete/"
 KB_TO_DELETE_STORAGE_BASE = "/storagetodelete/"
 
+RESOURCE_TO_DELETE_STORAGE_BASE = "/resourcestoragetodelete"
+RESOURCE_TO_DELETE_STORAGE = f"{RESOURCE_TO_DELETE_STORAGE_BASE}/{{kbid}}/{{uuid}}"
+
 KB_TO_DELETE = f"{KB_TO_DELETE_BASE}{{kbid}}"
 KB_TO_DELETE_STORAGE = f"{KB_TO_DELETE_STORAGE_BASE}{{kbid}}"
 
+KB_VECTORSET_TO_DELETE_BASE = "/vectorsettodelete"
+KB_VECTORSET_TO_DELETE = f"{KB_VECTORSET_TO_DELETE_BASE}/{{kbid}}/{{vectorset}}"
+
 
 class KnowledgeBox:
     def __init__(self, txn: Transaction, storage: Storage, kbid: str):
@@ -74,115 +90,154 @@ class KnowledgeBox:
         self.storage = storage
         self.kbid = kbid
         self._config: Optional[KnowledgeBoxConfig] = None
-        self.synonyms = Synonyms(self.txn, self.kbid)
-
-    async def get_config(self) -> Optional[KnowledgeBoxConfig]:
-        if self._config is None:
-            async with datamanagers.with_transaction() as txn:
-                config = await datamanagers.kb.get_config(txn, kbid=self.kbid)
-            if config is not None:
-                self._config = config
-                return config
-            else:
-                return None
-        else:
-            return self._config
-
-    @classmethod
-    async def delete_kb(cls, txn: Transaction, slug: str = "", kbid: str = ""):
-        # Mark storage to be deleted
-        # Mark keys to be deleted
-        logger.info(f"Deleting KB kbid={kbid} slug={slug}")
-        if not kbid and not slug:
-            raise AttributeError()
-
-        if slug and not kbid:
-            kbid_bytes = await txn.get(datamanagers.kb.KB_SLUGS.format(slug=slug))
-            if kbid_bytes is None:
-                raise datamanagers.exceptions.KnowledgeBoxNotFound()
-            kbid = kbid_bytes.decode()
-
-        if kbid and not slug:
-            kbconfig_bytes = await txn.get(datamanagers.kb.KB_UUID.format(kbid=kbid))
-            if kbconfig_bytes is None:
-                raise datamanagers.exceptions.KnowledgeBoxNotFound()
-            pbconfig = KnowledgeBoxConfig()
-            pbconfig.ParseFromString(kbconfig_bytes)
-            slug = pbconfig.slug
-
-        # Delete main anchor
-        async with txn.driver.transaction() as subtxn:
-            key_match = datamanagers.kb.KB_SLUGS.format(slug=slug)
-            logger.info(f"Deleting KB with slug: {slug}")
-            await subtxn.delete(key_match)
 
-            when = datetime.now().isoformat()
-            await subtxn.set(KB_TO_DELETE.format(kbid=kbid), when.encode())
-            await subtxn.commit()
-
-        audit_util = get_audit()
-        if audit_util is not None:
-            await audit_util.delete_kb(kbid)
-        return kbid
+    @staticmethod
+    def new_unique_kbid() -> str:
+        return str(uuid4())
 
     @classmethod
+    @processor_observer.wrap({"type": "create_kb"})
     async def create(
         cls,
-        txn: Transaction,
+        driver: Driver,
+        *,
+        kbid: str,
         slug: str,
-        semantic_model: SemanticModelMetadata,
-        uuid: Optional[str] = None,
-        config: Optional[KnowledgeBoxConfig] = None,
-        release_channel: ReleaseChannel.ValueType = ReleaseChannel.STABLE,
-    ) -> tuple[str, bool]:
-        failed = False
-        exist = await datamanagers.kb.get_kb_uuid(txn, slug=slug)
-        if exist:
-            raise KnowledgeBoxConflict()
-        if uuid is None or uuid == "":
-            uuid = str(uuid4())
-
-        if slug == "":
-            slug = uuid
-
-        await txn.set(
-            datamanagers.kb.KB_SLUGS.format(slug=slug),
-            uuid.encode(),
-        )
-        if config is None:
-            config = KnowledgeBoxConfig()
-
-        config.migration_version = get_latest_version()
-        config.slug = slug
-        await txn.set(
-            datamanagers.kb.KB_UUID.format(kbid=uuid),
-            config.SerializeToString(),
-        )
-        # Create Storage
-        storage = await get_storage(service_name=SERVICE_NAME)
+        title: str = "",
+        description: str = "",
+        semantic_models: Optional[dict[str, SemanticModelMetadata]] = None,
+        external_index_provider: CreateExternalIndexProviderMetadata = CreateExternalIndexProviderMetadata(),
+        hidden_resources_enabled: bool = False,
+        hidden_resources_hide_on_creation: bool = False,
+    ) -> tuple[str, str]:
+        """Creates a new knowledge box and return its id and slug."""
+
+        if not kbid:
+            raise KnowledgeBoxCreationError("A kbid must be provided to create a new KB")
+        if not slug:
+            raise KnowledgeBoxCreationError("A slug must be provided to create a new KB")
+        if hidden_resources_hide_on_creation and not hidden_resources_enabled:
+            raise KnowledgeBoxCreationError(
+                "Cannot hide new resources if the hidden resources feature is disabled"
+            )
+        if semantic_models is None or len(semantic_models) == 0:
+            raise KnowledgeBoxCreationError("KB must define at least one semantic model")
 
-        created = await storage.create_kb(uuid)
-        if created is False:
-            logger.error(f"{uuid} KB could not be created")
-            failed = True
+        rollback_ops: list[Callable[[], Coroutine[Any, Any, Any]]] = []
 
-        if failed is False:
-            shard_manager = get_shard_manager()
-            try:
-                await shard_manager.create_shard_by_kbid(
-                    txn,
-                    uuid,
-                    semantic_model=semantic_model,
-                    release_channel=release_channel,
+        try:
+            async with driver.transaction() as txn:
+                exists = await datamanagers.kb.get_kb_uuid(
+                    txn, slug=slug
+                ) or await datamanagers.kb.exists_kb(txn, kbid=kbid)
+                if exists:
+                    raise KnowledgeBoxConflict()
+
+                # Create in maindb
+                await datamanagers.kb.set_kbid_for_slug(txn, slug=slug, kbid=kbid)
+
+                # all KBs have the vectorset key initialized, although (for
+                # now), not every KB will store vectorsets there
+                await datamanagers.vectorsets.initialize(txn, kbid=kbid)
+
+                kb_shards = writer_pb2.Shards()
+                kb_shards.kbid = kbid
+                # B/c with Shards.actual
+                kb_shards.actual = -1
+
+                vs_external_indexes = []
+                for vectorset_id, semantic_model in semantic_models.items():  # type: ignore
+                    # if this KB uses a matryoshka model, we can choose a different
+                    # dimension
+                    if len(semantic_model.matryoshka_dimensions) > 0:
+                        dimension = choose_matryoshka_dimension(semantic_model.matryoshka_dimensions)
+                    else:
+                        dimension = semantic_model.vector_dimension
+
+                    vs_external_indexes.append(
+                        VectorsetExternalIndex(
+                            vectorset_id=vectorset_id,
+                            dimension=dimension,
+                            similarity=semantic_model.similarity_function,
+                        )
+                    )
+
+                    vectorset_config = knowledgebox_pb2.VectorSetConfig(
+                        vectorset_id=vectorset_id,
+                        vectorset_index_config=nodewriter_pb2.VectorIndexConfig(
+                            similarity=semantic_model.similarity_function,
+                            # XXX: hardcoded value
+                            vector_type=nodewriter_pb2.VectorType.DENSE_F32,
+                            normalize_vectors=len(semantic_model.matryoshka_dimensions) > 0,
+                            vector_dimension=dimension,
+                        ),
+                        matryoshka_dimensions=semantic_model.matryoshka_dimensions,
+                    )
+                    await datamanagers.vectorsets.set(txn, kbid=kbid, config=vectorset_config)
+
+                stored_external_index_provider = await cls._maybe_create_external_indexes(
+                    kbid, request=external_index_provider, indexes=vs_external_indexes
+                )
+                rollback_ops.append(
+                    partial(
+                        cls._maybe_delete_external_indexes,
+                        kbid,
+                        stored_external_index_provider,
+                    )
                 )
-            except Exception as e:
-                await storage.delete_kb(uuid)
-                raise e
-
-        if failed:
-            await storage.delete_kb(uuid)
 
-        return uuid, failed
+                config = KnowledgeBoxConfig(
+                    title=title,
+                    description=description,
+                    slug=slug,
+                    migration_version=get_latest_version(),
+                    hidden_resources_enabled=hidden_resources_enabled,
+                    hidden_resources_hide_on_creation=hidden_resources_hide_on_creation,
+                )
+                config.external_index_provider.CopyFrom(stored_external_index_provider)
+                await datamanagers.kb.set_config(txn, kbid=kbid, config=config)
+                await datamanagers.cluster.update_kb_shards(txn, kbid=kbid, shards=kb_shards)
+
+                # shard creation will alter this value on maindb, make sure nobody
+                # uses this variable anymore
+                del kb_shards
+
+                # Create in storage
+
+                storage = await get_storage(service_name=SERVICE_NAME)
+
+                created = await storage.create_kb(kbid)
+                if not created:
+                    logger.error(f"KB {kbid} could not be created")
+                    raise KnowledgeBoxCreationError(
+                        f"KB blob storage could not be created (slug={slug})"
+                    )
+                rollback_ops.append(partial(storage.delete_kb, kbid))
+
+                # Create shards in index nodes
+
+                shard_manager = get_shard_manager()
+                # XXX creating a shard is a slow IO operation that requires a write
+                # txn to be open!
+                await shard_manager.create_shard_by_kbid(txn, kbid)
+                # shards don't need a rollback as they will be eventually purged
+
+                await txn.commit()
+
+        except Exception as exc:
+            # rollback all changes on the db and raise the exception
+            for op in reversed(rollback_ops):
+                try:
+                    await op()
+                except Exception:
+                    if isinstance(op, partial):
+                        name: str = op.func.__name__
+                    else:
+                        getattr(op, "__name__", "unknown?")
+                    logger.exception(f"Unexpected error rolling back {name}. Keep rolling back")
+            raise exc
+
+        return (kbid, slug)
 
     @classmethod
     async def update(
@@ -192,7 +247,7 @@ class KnowledgeBox:
         slug: Optional[str] = None,
         config: Optional[KnowledgeBoxConfig] = None,
     ) -> str:
-        exist = await datamanagers.kb.get_config(txn, kbid=uuid)
+        exist = await datamanagers.kb.get_config(txn, kbid=uuid, for_update=True)
         if not exist:
             raise datamanagers.exceptions.KnowledgeBoxNotFound()
 
@@ -209,94 +264,61 @@ class KnowledgeBox:
 
         if config and exist != config:
             exist.MergeFrom(config)
+            exist.hidden_resources_enabled = config.hidden_resources_enabled
+            exist.hidden_resources_hide_on_creation = config.hidden_resources_hide_on_creation
 
-        await txn.set(
-            datamanagers.kb.KB_UUID.format(kbid=uuid),
-            exist.SerializeToString(),
-        )
+        if exist.hidden_resources_hide_on_creation and not exist.hidden_resources_enabled:
+            raise KnowledgeBoxCreationError(
+                "Cannot hide new resources if the hidden resources feature is disabled"
+            )
+
+        await datamanagers.kb.set_config(txn, kbid=uuid, config=exist)
 
         return uuid
 
-    async def iterate_kb_nodes(self) -> AsyncIterator[tuple[AbstractIndexNode, str]]:
-        async with datamanagers.with_transaction() as txn:
-            shards_obj = await datamanagers.cluster.get_kb_shards(txn, kbid=self.kbid)
-            if shards_obj is None:
-                raise ShardsNotFound(self.kbid)
-
-        for shard in shards_obj.shards:
-            for replica in shard.replicas:
-                node = get_index_node(replica.node)
-                if node is not None:
-                    yield node, replica.shard.id
-
-    # Vectorset
-    async def get_vectorsets(self, response: writer_pb2.GetVectorSetsResponse):
-        vectorset_key = KB_VECTORSET.format(kbid=self.kbid)
-        payload = await self.txn.get(vectorset_key)
-        if payload is not None:
-            response.vectorsets.ParseFromString(payload)
-
-    async def del_vectorset(self, id: str):
-        vectorset_key = KB_VECTORSET.format(kbid=self.kbid)
-        payload = await self.txn.get(vectorset_key)
-        vts = VectorSets()
-        if payload is not None:
-            vts.ParseFromString(payload)
-        del vts.vectorsets[id]
-        # For each Node on the KB delete the vectorset
-        async for node, shard in self.iterate_kb_nodes():
-            await node.del_vectorset(shard, id)
-        payload = vts.SerializeToString()
-        await self.txn.set(vectorset_key, payload)
-
-    async def set_vectorset(self, id: str, vs: VectorSet):
-        vectorset_key = KB_VECTORSET.format(kbid=self.kbid)
-        payload = await self.txn.get(vectorset_key)
-        vts = VectorSets()
-        if payload is not None:
-            vts.ParseFromString(payload)
-        vts.vectorsets[id].CopyFrom(vs)
-        # For each Node on the KB add the vectorset
-        async for node, shard in self.iterate_kb_nodes():
-            await node.set_vectorset(shard, id, similarity=vs.similarity)
-        payload = vts.SerializeToString()
-        await self.txn.set(vectorset_key, payload)
-
-    # Labels
-    async def set_labelset(self, id: str, labelset: LabelSet):
-        await datamanagers.labels.set_labelset(
-            self.txn, kbid=self.kbid, labelset_id=id, labelset=labelset
-        )
+    @classmethod
+    async def delete(cls, driver: Driver, kbid: str):
+        async with driver.transaction() as txn:
+            exists = await datamanagers.kb.exists_kb(txn, kbid=kbid)
+            if not exists:
+                return
 
-    async def get_labels(self) -> Labels:
-        return await datamanagers.labels.get_labels(self.txn, kbid=self.kbid)
+            # Delete main anchor
+            kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
+            if kb_config is not None:
+                slug = kb_config.slug
+                await datamanagers.kb.delete_kb_slug(txn, slug=slug)
 
-    async def get_labelset(
-        self, labelset: str, labelset_response: writer_pb2.GetLabelSetResponse
-    ):
-        ls = await datamanagers.labels.get_labelset(
-            self.txn,
-            kbid=self.kbid,
-            labelset_id=labelset,
-        )
-        if ls is not None:
-            labelset_response.labelset.CopyFrom(ls)
+            await datamanagers.kb.delete_config(txn, kbid=kbid)
 
-    async def del_labelset(self, id: str):
-        await datamanagers.labels.delete_labelset(
-            self.txn, kbid=self.kbid, labelset_id=id
-        )
+            # Mark KB to purge. This will eventually delete all KB keys, storage
+            # and index data (for the old index nodes)
+            when = datetime.now().isoformat()
+            await txn.set(KB_TO_DELETE.format(kbid=kbid), when.encode())
 
-    async def get_synonyms(self, synonyms: PBSynonyms):
-        pbsyn = await self.synonyms.get()
-        if pbsyn is not None:
-            synonyms.CopyFrom(pbsyn)
+            shards_obj = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
 
-    async def set_synonyms(self, synonyms: PBSynonyms):
-        await self.synonyms.set(synonyms)
+            await txn.commit()
 
-    async def delete_synonyms(self):
-        await self.synonyms.clear()
+        if shards_obj is None:
+            logger.warning(f"Shards not found for KB while deleting it", extra={"kbid": kbid})
+        else:
+            nidx_api = get_nidx_api_client()
+            # Delete shards from nidx. They'll be marked for eventual deletion,
+            # so this call shouldn't be costly
+            if nidx_api is not None:
+                for shard in shards_obj.shards:
+                    if shard.nidx_shard_id:
+                        await nidx_api.DeleteShard(noderesources_pb2.ShardId(id=shard.nidx_shard_id))
+
+        if kb_config is not None:
+            await cls._maybe_delete_external_indexes(kbid, kb_config.external_index_provider)
+
+        audit = get_audit()
+        if audit is not None:
+            audit.delete_kb(kbid=kbid)
+
+        return kbid
 
     @classmethod
     async def purge(cls, driver: Driver, kbid: str):
@@ -307,6 +329,8 @@ class KnowledgeBox:
         need to delete the kb shards and also deletes the related storage
         buckets.
 
+        Removes all catalog entries related to the kb.
+
         As non-empty buckets cannot be deleted, they are scheduled to be
         deleted instead. Actually, this empties the bucket asynchronouysly
         but it doesn't delete it. To do it, we save a marker using the
@@ -322,16 +346,13 @@ class KnowledgeBox:
             storage_to_delete = KB_TO_DELETE_STORAGE.format(kbid=kbid)
             await txn.set(storage_to_delete, b"")
 
-            # Delete KB Shards
-            shards_match = datamanagers.cluster.KB_SHARDS.format(kbid=kbid)
-            payload = await txn.get(shards_match)
+            await catalog_delete_kb(txn, kbid)
 
-            if payload is None:
-                logger.warning(f"Shards not found for kbid={kbid}")
+            # Delete KB Shards
+            shards_obj = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
+            if shards_obj is None:
+                logger.warning(f"Shards not found for KB while purging it", extra={"kbid": kbid})
             else:
-                shards_obj = writer_pb2.Shards()
-                shards_obj.ParseFromString(payload)  # type: ignore
-
                 for shard in shards_obj.shards:
                     # Delete the shard on nodes
                     for replica in shard.replicas:
@@ -357,29 +378,14 @@ class KnowledgeBox:
         await cls.delete_all_kb_keys(driver, kbid)
 
     @classmethod
-    async def delete_all_kb_keys(
-        cls, driver: Driver, kbid: str, chunk_size: int = 1_000
-    ):
+    async def delete_all_kb_keys(cls, driver: Driver, kbid: str, chunk_size: int = 1_000):
         prefix = KB_KEYS.format(kbid=kbid)
-        while True:
-            async with driver.transaction() as txn:
-                all_keys = [key async for key in txn.keys(match=prefix, count=-1)]
-
-            if len(all_keys) == 0:
-                break
-
-            # We commit deletions in chunks because otherwise
-            # tikv complains if there is too much data to commit
-            for chunk_of_keys in chunker(all_keys, chunk_size):
-                async with driver.transaction() as txn:
-                    for key in chunk_of_keys:
-                        await txn.delete(key)
-                    await txn.commit()
-
-    async def get_resource_shard(
-        self, shard_id: str
-    ) -> Optional[writer_pb2.ShardObject]:
-        async with datamanagers.with_transaction() as txn:
+        async with driver.transaction() as txn:
+            await txn.delete_by_prefix(prefix)
+            await txn.commit()
+
+    async def get_resource_shard(self, shard_id: str) -> Optional[writer_pb2.ShardObject]:
+        async with datamanagers.with_ro_transaction() as txn:
             pb = await datamanagers.cluster.get_kb_shards(txn, kbid=self.kbid)
             if pb is None:
                 logger.warning("Shards not found for kbid", extra={"kbid": self.kbid})
@@ -390,52 +396,55 @@ class KnowledgeBox:
         return None
 
     async def get(self, uuid: str) -> Optional[Resource]:
-        raw_basic = await get_basic(self.txn, self.kbid, uuid)
-        if raw_basic:
-            return Resource(
-                txn=self.txn,
-                storage=self.storage,
-                kb=self,
-                uuid=uuid,
-                basic=Resource.parse_basic(raw_basic),
-                disable_vectors=False,
-            )
-        else:
+        basic = await datamanagers.resources.get_basic(self.txn, kbid=self.kbid, rid=uuid)
+        if basic is None:
             return None
+        return Resource(
+            txn=self.txn,
+            storage=self.storage,
+            kb=self,
+            uuid=uuid,
+            basic=basic,
+            disable_vectors=False,
+        )
 
-    async def delete_resource(self, uuid: str):
-        raw_basic = await get_basic(self.txn, self.kbid, uuid)
-        if raw_basic:
-            basic = Resource.parse_basic(raw_basic)
-        else:
-            basic = None
-
-        async for key in self.txn.keys(
-            KB_RESOURCE.format(kbid=self.kbid, uuid=uuid), count=-1
-        ):
-            await self.txn.delete(key)
-
+    async def maindb_delete_resource(self, uuid: str):
+        basic = await datamanagers.resources.get_basic(self.txn, kbid=self.kbid, rid=uuid)
+        await self.txn.delete_by_prefix(KB_RESOURCE.format(kbid=self.kbid, uuid=uuid))
         if basic and basic.slug:
-            slug_key = KB_RESOURCE_SLUG.format(kbid=self.kbid, slug=basic.slug)
             try:
-                await self.txn.delete(slug_key)
+                await self.txn.delete(KB_RESOURCE_SLUG.format(kbid=self.kbid, slug=basic.slug))
             except Exception:
-                pass
+                logger.exception("Error deleting slug")
+
+    async def storage_delete_resource(self, uuid: str):
+        if is_onprem_nucliadb():
+            await self.storage.delete_resource(self.kbid, uuid)
+        else:
+            # Deleting from storage can be slow, so we schedule its deletion and the purge cronjob
+            # will take care of it
+            await self.schedule_delete_resource(self.kbid, uuid)
+
+    async def schedule_delete_resource(self, kbid: str, uuid: str):
+        key = RESOURCE_TO_DELETE_STORAGE.format(kbid=kbid, uuid=uuid)
+        await self.txn.set(key, b"")
 
-        await self.storage.delete_resource(self.kbid, uuid)
+    async def delete_resource(self, uuid: str):
+        with processor_observer({"type": "delete_resource_maindb"}):
+            await self.maindb_delete_resource(uuid)
+        with processor_observer({"type": "delete_resource_storage"}):
+            await self.storage_delete_resource(uuid)
 
     async def get_resource_uuid_by_slug(self, slug: str) -> Optional[str]:
-        uuid = await self.txn.get(KB_RESOURCE_SLUG.format(kbid=self.kbid, slug=slug))
-        if uuid is not None:
-            return uuid.decode()
-        else:
-            return None
+        return await datamanagers.resources.get_resource_uuid_from_slug(
+            self.txn, kbid=self.kbid, slug=slug
+        )
 
     async def get_unique_slug(self, uuid: str, slug: str) -> str:
         key = KB_RESOURCE_SLUG.format(kbid=self.kbid, slug=slug)
         key_ok = False
         while key_ok is False:
-            found = await self.txn.get(key)
+            found = await self.txn.get(key, for_update=False)
             if found and found.decode() != uuid:
                 slug += ".c"
                 key = KB_RESOURCE_SLUG.format(kbid=self.kbid, slug=slug)
@@ -443,17 +452,7 @@ class KnowledgeBox:
                 key_ok = True
         return slug
 
-    @classmethod
-    async def resource_slug_exists(
-        self, 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 add_resource(
-        self, uuid: str, slug: str, basic: Optional[Basic] = None
-    ) -> Resource:
+    async def add_resource(self, uuid: str, slug: str, basic: Optional[Basic] = None) -> Resource:
         if basic is None:
             basic = Basic()
         if slug == "":
@@ -461,7 +460,7 @@ class KnowledgeBox:
         slug = await self.get_unique_slug(uuid, slug)
         basic.slug = slug
         fix_paragraph_annotation_keys(uuid, basic)
-        await set_basic(self.txn, self.kbid, uuid, basic)
+        await datamanagers.resources.set_basic(self.txn, kbid=self.kbid, rid=uuid, basic=basic)
         return Resource(
             storage=self.storage,
             txn=self.txn,
@@ -473,7 +472,7 @@ class KnowledgeBox:
 
     async def iterate_resources(self) -> AsyncGenerator[Resource, None]:
         base = KB_RESOURCE_SLUG_BASE.format(kbid=self.kbid)
-        async for key in self.txn.keys(match=base, count=-1):
+        async for key in self.txn.keys(match=base):
             slug = key.split("/")[-1]
             uuid = await self.get_resource_uuid_by_slug(slug)
             if uuid is not None:
@@ -485,6 +484,55 @@ class KnowledgeBox:
                     disable_vectors=False,
                 )
 
+    async def create_vectorset(self, config: knowledgebox_pb2.VectorSetConfig):
+        if await datamanagers.vectorsets.exists(
+            self.txn, kbid=self.kbid, vectorset_id=config.vectorset_id
+        ):
+            raise VectorSetConflict(f"Vectorset {config.vectorset_id} already exists")
+        await datamanagers.vectorsets.set(self.txn, kbid=self.kbid, config=config)
+
+        # Remove the async deletion mark if it exists, just in case there was a previous deletion
+        deletion_mark_key = KB_VECTORSET_TO_DELETE.format(kbid=self.kbid, vectorset=config.vectorset_id)
+        deletion_mark = await self.txn.get(deletion_mark_key, for_update=True)
+        if deletion_mark is not None:
+            await self.txn.delete(deletion_mark_key)
+
+        shard_manager = get_shard_manager()
+        await shard_manager.create_vectorset(self.kbid, config)
+
+    async def delete_vectorset(self, vectorset_id: str):
+        await datamanagers.vectorsets.delete(self.txn, kbid=self.kbid, vectorset_id=vectorset_id)
+
+        # mark vectorset for async deletion
+        deletion_mark_key = KB_VECTORSET_TO_DELETE.format(kbid=self.kbid, vectorset=vectorset_id)
+        await self.txn.set(deletion_mark_key, b"")
+
+        shard_manager = get_shard_manager()
+        await shard_manager.delete_vectorset(self.kbid, vectorset_id)
+
+    @classmethod
+    async def _maybe_create_external_indexes(
+        cls,
+        kbid: str,
+        request: CreateExternalIndexProviderMetadata,
+        indexes: list[VectorsetExternalIndex],
+    ) -> StoredExternalIndexProviderMetadata:
+        if request.type != ExternalIndexProviderType.PINECONE:
+            return StoredExternalIndexProviderMetadata(type=request.type)
+        # Only pinecone is supported for now
+        return await PineconeIndexManager.create_indexes(kbid, request, indexes)
+
+    @classmethod
+    async def _maybe_delete_external_indexes(
+        cls,
+        kbid: str,
+        stored: StoredExternalIndexProviderMetadata,
+    ) -> None:
+        if stored.type != ExternalIndexProviderType.PINECONE:
+            return
+        # Only pinecone is supported for now
+        await PineconeIndexManager.delete_indexes(kbid, stored)
+
 
 def chunker(seq: Sequence, size: int):
     return (seq[pos : pos + size] for pos in range(0, len(seq), size))
@@ -498,3 +546,11 @@ def fix_paragraph_annotation_keys(uuid: str, basic: Basic) -> None:
         for paragraph_annotation in ufm.paragraphs:
             key = compute_paragraph_key(uuid, paragraph_annotation.key)
             paragraph_annotation.key = key
+
+
+@processor_observer.wrap({"type": "catalog_delete_kb"})
+async def catalog_delete_kb(txn: Transaction, kbid: str):
+    if not isinstance(txn, PGTransaction):
+        return
+    async with txn.connection.cursor() as cur:
+        await cur.execute("DELETE FROM catalog where kbid = %(kbid)s", {"kbid": kbid})