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

nucliadb/common/datamanagers/fields.py

@@ -0,0 +1,84 @@
+# 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 Optional
+
+from google.protobuf.message import Message
+
+from nucliadb.common.datamanagers.utils import get_kv_pb
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb_protos import writer_pb2
+
+KB_RESOURCE_FIELD = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}"
+KB_RESOURCE_FIELD_ERROR = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}/error"
+
+
+async def get_raw(
+    txn: Transaction, *, kbid: str, rid: str, field_type: str, field_id: str
+) -> Optional[bytes]:
+    key = KB_RESOURCE_FIELD.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
+    return await txn.get(key)
+
+
+async def set(
+    txn: Transaction,
+    *,
+    kbid: str,
+    rid: str,
+    field_type: str,
+    field_id: str,
+    value: Message,
+):
+    key = KB_RESOURCE_FIELD.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
+    await txn.set(key, value.SerializeToString())
+
+
+async def delete(txn: Transaction, *, kbid: str, rid: str, field_type: str, field_id: str):
+    base_key = KB_RESOURCE_FIELD.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
+    # Make sure we explicitly delete the field and any nested key
+    keys_to_delete = []
+    async for key in txn.keys(base_key):
+        keys_to_delete.append(key)
+
+    for key in keys_to_delete:
+        await txn.delete(key)
+
+
+# Error
+
+
+async def get_error(
+    txn: Transaction, *, kbid: str, rid: str, field_type: str, field_id: str
+) -> Optional[writer_pb2.Error]:
+    key = KB_RESOURCE_FIELD_ERROR.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
+    return await get_kv_pb(txn, key, writer_pb2.Error)
+
+
+async def set_error(
+    txn: Transaction,
+    *,
+    kbid: str,
+    rid: str,
+    field_type: str,
+    field_id: str,
+    error: writer_pb2.Error,
+):
+    key = KB_RESOURCE_FIELD_ERROR.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
+    await txn.set(key, error.SerializeToString())

nucliadb/common/datamanagers/kb.py

@@ -33,10 +33,8 @@ KB_SLUGS = KB_SLUGS_BASE + "{slug}"
 logger = logging.getLogger(__name__)
 
 
-async def get_kbs(
-    txn: Transaction, *, prefix: str = ""
-) -> AsyncIterator[tuple[str, str]]:
-    async for key in txn.keys(KB_SLUGS.format(slug=prefix), count=-1):
+async def get_kbs(txn: Transaction, *, prefix: str = "") -> AsyncIterator[tuple[str, str]]:
+    async for key in txn.keys(KB_SLUGS.format(slug=prefix)):
         slug = key.replace(KB_SLUGS_BASE, "")
         uuid = await get_kb_uuid(txn, slug=slug)
         if uuid is None:
@@ -46,22 +44,32 @@ async def get_kbs(
 
 
 async def exists_kb(txn: Transaction, *, kbid: str) -> bool:
-    return await get_config(txn, kbid=kbid) is not None
+    return await get_config(txn, kbid=kbid, for_update=False) is not None
 
 
 async def get_kb_uuid(txn: Transaction, *, slug: str) -> Optional[str]:
-    uuid = await txn.get(KB_SLUGS.format(slug=slug))
+    uuid = await txn.get(KB_SLUGS.format(slug=slug), for_update=False)
     if uuid is not None:
         return uuid.decode()
     else:
         return None
 
 
+async def set_kbid_for_slug(txn: Transaction, *, slug: str, kbid: str):
+    key = KB_SLUGS.format(slug=slug)
+    await txn.set(key, kbid.encode())
+
+
+async def delete_kb_slug(txn: Transaction, *, slug: str):
+    key = KB_SLUGS.format(slug=slug)
+    await txn.delete(key)
+
+
 async def get_config(
-    txn: Transaction, *, kbid: str
+    txn: Transaction, *, kbid: str, for_update: bool = False
 ) -> Optional[knowledgebox_pb2.KnowledgeBoxConfig]:
     key = KB_UUID.format(kbid=kbid)
-    payload = await txn.get(key)
+    payload = await txn.get(key, for_update=for_update)
     if payload is None:
         return None
     response = knowledgebox_pb2.KnowledgeBoxConfig()
@@ -69,17 +77,18 @@ async def get_config(
     return response
 
 
-async def set_config(
-    txn: Transaction, *, kbid: str, config: knowledgebox_pb2.KnowledgeBoxConfig
-):
+async def set_config(txn: Transaction, *, kbid: str, config: knowledgebox_pb2.KnowledgeBoxConfig):
     key = KB_UUID.format(kbid=kbid)
     await txn.set(key, config.SerializeToString())
 
 
-async def get_model_metadata(
-    txn: Transaction, *, kbid: str
-) -> knowledgebox_pb2.SemanticModelMetadata:
-    shards_obj = await cluster.get_kb_shards(txn, kbid=kbid)
+async def delete_config(txn: Transaction, *, kbid: str) -> None:
+    key = KB_UUID.format(kbid=kbid)
+    await txn.delete(key)
+
+
+async def get_model_metadata(txn: Transaction, *, kbid: str) -> knowledgebox_pb2.SemanticModelMetadata:
+    shards_obj = await cluster.get_kb_shards(txn, kbid=kbid, for_update=False)
     if shards_obj is None:
         raise KnowledgeBoxNotFound(kbid)
     if shards_obj.HasField("model"):
@@ -87,30 +96,67 @@ async def get_model_metadata(
     else:
         # B/c code for old KBs that do not have the `model` attribute set in the Shards object.
         # Cleanup this code after a migration is done unifying all fields under `model` (on-prem and cloud).
-        return knowledgebox_pb2.SemanticModelMetadata(
-            similarity_function=shards_obj.similarity
-        )
+        return knowledgebox_pb2.SemanticModelMetadata(similarity_function=shards_obj.similarity)
 
 
+# DEPRECATED: this function should be removed once the "default" vectorset
+# concept is removed and processing sends us all messages with a vectorset_id
 async def get_matryoshka_vector_dimension(
-    txn: Transaction, *, kbid: str
+    txn: Transaction,
+    *,
+    kbid: str,
+    vectorset_id: Optional[str] = None,
 ) -> Optional[int]:
     """Return vector dimension for matryoshka models"""
-    model_metadata = await get_model_metadata(txn, kbid=kbid)
-    dimension = None
-    if (
-        len(model_metadata.matryoshka_dimensions) > 0
-        and model_metadata.vector_dimension
-    ):
-        if model_metadata.vector_dimension in model_metadata.matryoshka_dimensions:
-            dimension = model_metadata.vector_dimension
-        else:
-            logger.error(
-                "KB has an invalid matryoshka dimension!",
-                extra={
-                    "kbid": kbid,
-                    "vector_dimension": model_metadata.vector_dimension,
-                    "matryoshka_dimensions": model_metadata.matryoshka_dimensions,
-                },
-            )
-    return dimension
+    from . import vectorsets
+
+    async for _, vs in vectorsets.iter(txn, kbid=kbid):
+        if len(vs.matryoshka_dimensions) > 0 and vs.vectorset_index_config.vector_dimension:
+            if vs.vectorset_index_config.vector_dimension in vs.matryoshka_dimensions:
+                return vs.vectorset_index_config.vector_dimension
+            else:
+                logger.error(
+                    "KB has an invalid matryoshka dimension!",
+                    extra={
+                        "kbid": kbid,
+                        "vector_dimension": vs.vectorset_index_config.vector_dimension,
+                        "matryoshka_dimensions": vs.matryoshka_dimensions,
+                    },
+                )
+        return None
+    else:
+        # fallback for KBs that don't have vectorset
+        model_metadata = await get_model_metadata(txn, kbid=kbid)
+        dimension = None
+        if len(model_metadata.matryoshka_dimensions) > 0 and model_metadata.vector_dimension:
+            if model_metadata.vector_dimension in model_metadata.matryoshka_dimensions:
+                dimension = model_metadata.vector_dimension
+            else:
+                logger.error(
+                    "KB has an invalid matryoshka dimension!",
+                    extra={
+                        "kbid": kbid,
+                        "vector_dimension": model_metadata.vector_dimension,
+                        "matryoshka_dimensions": model_metadata.matryoshka_dimensions,
+                    },
+                )
+        return dimension
+
+
+async def get_external_index_provider_metadata(
+    txn: Transaction, *, kbid: str
+) -> Optional[knowledgebox_pb2.StoredExternalIndexProviderMetadata]:
+    kb_config = await get_config(txn, kbid=kbid)
+    if kb_config is None:
+        return None
+    return kb_config.external_index_provider
+
+
+async def set_external_index_provider_metadata(
+    txn: Transaction, *, kbid: str, metadata: knowledgebox_pb2.StoredExternalIndexProviderMetadata
+):
+    kb_config = await get_config(txn, kbid=kbid)
+    if kb_config is None:
+        raise KnowledgeBoxNotFound(kbid)
+    kb_config.external_index_provider.CopyFrom(metadata)
+    await set_config(txn, kbid=kbid, config=kb_config)

nucliadb/common/datamanagers/labels.py

@@ -17,6 +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/>.
 #
+import logging
 from typing import Optional
 
 import orjson
@@ -24,6 +25,8 @@ import orjson
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb_protos import knowledgebox_pb2 as kb_pb2
 
+logger = logging.getLogger(__name__)
+
 KB_LABELS = "/kbs/{kbid}/labels"
 KB_LABELSET = "/kbs/{kbid}/labels/{id}"
 KB_LABELSET_IDS = "/kbs/{kbid}/ids-labels"
@@ -34,7 +37,9 @@ async def get_labels(txn: Transaction, *, kbid: str) -> kb_pb2.Labels:
     Get all labels for a knowledge box (from multiple labelsets)
     """
     labels = kb_pb2.Labels()
-    labelset_ids = await _get_labelset_ids_bw_compat(txn, kbid=kbid)
+    labelset_ids = await _get_labelset_ids(txn, kbid=kbid)
+    if labelset_ids is None:
+        return labels
     for labelset_id in labelset_ids:
         labelset = await txn.get(KB_LABELSET.format(kbid=kbid, id=labelset_id))
         if not labelset:
@@ -45,76 +50,41 @@ async def get_labels(txn: Transaction, *, kbid: str) -> kb_pb2.Labels:
     return labels
 
 
-async def _get_labelset_ids_bw_compat(txn: Transaction, *, kbid: str) -> list[str]:
-    labelsets = await _get_labelset_ids(txn, kbid=kbid)
-    if labelsets is not None:
-        return labelsets
-    # TODO: Remove this after migration #11
-    return await _deprecated_scan_labelset_ids(txn, kbid=kbid)
-
-
-async def _deprecated_scan_labelset_ids(txn: Transaction, *, kbid: str) -> list[str]:
-    labelsets = []
-    labels_key = KB_LABELS.format(kbid=kbid)
-    async for key in txn.keys(labels_key, count=-1, include_start=False):
-        lsid = key.split("/")[-1]
-        labelsets.append(lsid)
-    return labelsets
-
-
 async def _get_labelset_ids(txn: Transaction, *, kbid: str) -> Optional[list[str]]:
     key = KB_LABELSET_IDS.format(kbid=kbid)
-    data = await txn.get(key)
+    data = await txn.get(key, for_update=True)
     if not data:
         return None
     return orjson.loads(data)
 
 
-async def _add_to_labelset_ids(
-    txn: Transaction, *, kbid: str, labelsets: list[str]
-) -> None:
+async def _add_to_labelset_ids(txn: Transaction, *, kbid: str, labelsets: list[str]) -> None:
+    updated = set(labelsets)
     previous = await _get_labelset_ids(txn, kbid=kbid)
-    needs_set = False
-    if previous is None:
-        # TODO: Remove this after migration #11
-        needs_set = True
-        previous = await _deprecated_scan_labelset_ids(txn, kbid=kbid)
-    for labelset in labelsets:
-        if labelset not in previous:
-            needs_set = True
-            previous.append(labelset)
-    if needs_set:
-        await _set_labelset_ids(txn, kbid=kbid, labelsets=previous)
-
-
-async def _delete_from_labelset_ids(
-    txn: Transaction, *, kbid: str, labelsets: list[str]
-) -> None:
-    needs_set = False
+    if previous is not None:
+        updated.update(previous)
+    if previous is None or previous != updated:
+        await _set_labelset_ids(txn, kbid=kbid, labelsets=list(updated))
+
+
+async def _delete_from_labelset_ids(txn: Transaction, *, kbid: str, labelsets: list[str]) -> None:
     previous = await _get_labelset_ids(txn, kbid=kbid)
     if previous is None:
-        # TODO: Remove this after migration #11
-        needs_set = True
-        previous = await _deprecated_scan_labelset_ids(txn, kbid=kbid)
-    for labelset in labelsets:
-        if labelset in previous:
-            needs_set = True
-            previous.remove(labelset)
-    if needs_set:
-        await _set_labelset_ids(txn, kbid=kbid, labelsets=previous)
-
-
-async def _set_labelset_ids(
-    txn: Transaction, *, kbid: str, labelsets: list[str]
-) -> None:
+        # Nothing to delete
+        return
+    previous_set = set(previous)
+    updated = previous_set - set(labelsets)
+    if previous_set != updated:
+        await _set_labelset_ids(txn, kbid=kbid, labelsets=list(updated))
+
+
+async def _set_labelset_ids(txn: Transaction, *, kbid: str, labelsets: list[str]) -> None:
     key = KB_LABELSET_IDS.format(kbid=kbid)
     data = orjson.dumps(labelsets)
     await txn.set(key, data)
 
 
-async def get_labelset(
-    txn: Transaction, *, kbid: str, labelset_id: str
-) -> Optional[kb_pb2.LabelSet]:
+async def get_labelset(txn: Transaction, *, kbid: str, labelset_id: str) -> Optional[kb_pb2.LabelSet]:
     labelset_key = KB_LABELSET.format(kbid=kbid, id=labelset_id)
     payload = await txn.get(labelset_key)
     if payload:

nucliadb/common/datamanagers/processing.py

@@ -28,9 +28,7 @@ logger = logging.getLogger(__name__)
 PULL_PARTITION_OFFSET = "/processing/pull-offset/{pull_type_id}/{partition}"
 
 
-async def get_pull_offset(
-    txn: Transaction, *, pull_type_id: str, partition: str
-) -> Optional[int]:
+async def get_pull_offset(txn: Transaction, *, pull_type_id: str, partition: str) -> Optional[int]:
     key = PULL_PARTITION_OFFSET.format(pull_type_id=pull_type_id, partition=partition)
     val: Optional[bytes] = await txn.get(key)
     if val is not None:
@@ -38,8 +36,6 @@ async def get_pull_offset(
     return None
 
 
-async def set_pull_offset(
-    txn: Transaction, *, pull_type_id: str, partition: str, offset: int
-) -> None:
+async def set_pull_offset(txn: Transaction, *, pull_type_id: str, partition: str, offset: int) -> None:
     key = PULL_PARTITION_OFFSET.format(pull_type_id=pull_type_id, partition=partition)
     await txn.set(key, str(offset).encode())