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

nucliadb/common/datamanagers/atomic.py

@@ -0,0 +1,102 @@
+# 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/>.
+#
+
+"""
+Atomic datamanagers
+
+This module aims to provide a simple way to call a datamanager function in a
+single transaction, avoiding the need of encapsulating like in this example:
+
+```
+async def <function>(...):
+    async with datamanagers.with_transaction() as txn:
+        await datamanagers.<module>.<function>(...)
+```
+
+Or simply a more handy way to call an datamanager operation without caring about
+it's transaction
+
+"""
+
+import sys
+from functools import wraps
+
+from . import kb as kb_dm
+from . import labels as labels_dm
+from . import resources as resources_dm
+from . import synonyms as synonyms_dm
+from .utils import with_ro_transaction, with_transaction
+
+# XXX: we are using the not exported _ParamSpec to support 3.9. Whenever we
+# upgrade to >= 3.10 we'll be able to use ParamSpecKwargs and improve the
+# typing. We are abusing of ParamSpec anywat to better support text editors, so
+# we also need to ignore some mypy complains
+
+__python_version = (sys.version_info.major, sys.version_info.minor)
+if __python_version == (3, 9):
+    from typing_extensions import ParamSpec
+else:
+    from typing import ParamSpec  # type: ignore
+
+P = ParamSpec("P")
+
+
+def ro_txn_wrap(fun: P) -> P:  # type: ignore
+    @wraps(fun)
+    async def wrapper(**kwargs: P.kwargs):
+        async with with_ro_transaction() as txn:
+            return await fun(txn, **kwargs)
+
+    return wrapper
+
+
+def rw_txn_wrap(fun: P) -> P:  # type: ignore
+    @wraps(fun)
+    async def wrapper(**kwargs: P.kwargs):
+        async with with_transaction() as txn:
+            result = await fun(txn, **kwargs)
+            await txn.commit()
+            return result
+
+    return wrapper
+
+
+class kb:
+    exists_kb = ro_txn_wrap(kb_dm.exists_kb)
+    get_config = ro_txn_wrap(kb_dm.get_config)
+    get_external_index_provider_metadata = ro_txn_wrap(kb_dm.get_external_index_provider_metadata)
+
+
+class resources:
+    get_resource_uuid_from_slug = ro_txn_wrap(resources_dm.get_resource_uuid_from_slug)
+    resource_exists = ro_txn_wrap(resources_dm.resource_exists)
+    slug_exists = ro_txn_wrap(resources_dm.slug_exists)
+
+
+class labelset:
+    get = ro_txn_wrap(labels_dm.get_labelset)
+    set = rw_txn_wrap(labels_dm.set_labelset)
+    delete = rw_txn_wrap(labels_dm.delete_labelset)
+    get_all = ro_txn_wrap(labels_dm.get_labels)
+
+
+class synonyms:
+    get = ro_txn_wrap(synonyms_dm.get)
+    set = rw_txn_wrap(synonyms_dm.set)

nucliadb/common/datamanagers/cluster.py

@@ -31,13 +31,13 @@ logger = logging.getLogger(__name__)
 KB_SHARDS = "/kbs/{kbid}/shards"
 
 
-async def get_kb_shards(txn: Transaction, *, kbid: str) -> Optional[writer_pb2.Shards]:
+async def get_kb_shards(
+    txn: Transaction, *, kbid: str, for_update: bool = False
+) -> Optional[writer_pb2.Shards]:
     key = KB_SHARDS.format(kbid=kbid)
-    return await get_kv_pb(txn, key, writer_pb2.Shards)
+    return await get_kv_pb(txn, key, writer_pb2.Shards, for_update=for_update)
 
 
-async def update_kb_shards(
-    txn: Transaction, *, kbid: str, shards: writer_pb2.Shards
-) -> None:
+async def update_kb_shards(txn: Transaction, *, kbid: str, shards: writer_pb2.Shards) -> None:
     key = KB_SHARDS.format(kbid=kbid)
     await txn.set(key, shards.SerializeToString())

nucliadb/common/datamanagers/entities.py

@@ -85,11 +85,9 @@ async def set_entities_group(
     await txn.set(key, entities.SerializeToString())
 
 
-async def iterate_entities_groups(
-    txn: Transaction, *, kbid: str
-) -> AsyncGenerator[str, None]:
+async def iterate_entities_groups(txn: Transaction, *, kbid: str) -> AsyncGenerator[str, None]:
     entities_key = KB_ENTITIES.format(kbid=kbid)
-    async for key in txn.keys(entities_key, count=-1):
+    async for key in txn.keys(entities_key):
         group = key.split("/")[-1]
         yield group
 
@@ -106,9 +104,7 @@ async def get_entities_group(
     return eg
 
 
-async def get_deleted_groups(
-    txn: Transaction, *, kbid: str
-) -> kb_pb2.DeletedEntitiesGroups:
+async def get_deleted_groups(txn: Transaction, *, kbid: str) -> kb_pb2.DeletedEntitiesGroups:
     deleted_groups_key = KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid)
     payload = await txn.get(deleted_groups_key)
     deg = kb_pb2.DeletedEntitiesGroups()
@@ -122,18 +118,14 @@ async def mark_group_as_deleted(txn: Transaction, *, kbid: str, group: str) -> N
     deg = await get_deleted_groups(txn, kbid=kbid)
     if group not in deg.entities_groups:
         deg.entities_groups.append(group)
-        await txn.set(
-            KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid), deg.SerializeToString()
-        )
+        await txn.set(KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid), deg.SerializeToString())
 
 
 async def unmark_group_as_deleted(txn: Transaction, *, kbid: str, group: str) -> None:
     deg = await get_deleted_groups(txn, kbid=kbid)
     if group in deg.entities_groups:
         deg.entities_groups.remove(group)
-        await txn.set(
-            KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid), deg.SerializeToString()
-        )
+        await txn.set(KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid), deg.SerializeToString())
 
 
 async def get_entities_meta_cache(txn: Transaction, *, kbid: str) -> EntitiesMetaCache:
@@ -143,7 +135,5 @@ async def get_entities_meta_cache(txn: Transaction, *, kbid: str) -> EntitiesMet
     return pickle.loads(value)
 
 
-async def set_entities_meta_cache(
-    txn: Transaction, kbid: str, cache: EntitiesMetaCache
-) -> None:
+async def set_entities_meta_cache(txn: Transaction, kbid: str, cache: EntitiesMetaCache) -> None:
     await txn.set(KB_ENTITIES_CACHE.format(kbid=kbid), pickle.dumps(cache, protocol=5))

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,15 +33,43 @@ 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)):
+        slug = key.replace(KB_SLUGS_BASE, "")
+        uuid = await get_kb_uuid(txn, slug=slug)
+        if uuid is None:
+            logger.error(f"KB with slug ({slug}) but without uuid?")
+            continue
+        yield (uuid, slug)
+
+
 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), 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()
@@ -49,10 +77,18 @@ async def get_config(
     return response
 
 
-async def get_model_metadata(
-    txn: Transaction, *, kbid: str
-) -> knowledgebox_pb2.SemanticModelMetadata:
-    shards_obj = await cluster.get_kb_shards(txn, kbid=kbid)
+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 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"):
@@ -60,26 +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)
 
 
-async def get_kb_uuid(txn: Transaction, *, slug: str) -> Optional[str]:
-    uuid = await txn.get(KB_SLUGS.format(slug=slug))
-    if uuid is not None:
-        return uuid.decode()
+# 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,
+    vectorset_id: Optional[str] = None,
+) -> Optional[int]:
+    """Return vector dimension for matryoshka models"""
+    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 get_kbs(
-    txn: Transaction, *, prefix: str = ""
-) -> AsyncIterator[tuple[str, str]]:
-    async for key in txn.keys(KB_SLUGS.format(slug=prefix), count=-1):
-        slug = key.replace(KB_SLUGS_BASE, "")
-        uuid = await get_kb_uuid(txn, slug=slug)
-        if uuid is None:
-            logger.error(f"KB with slug ({slug}) but without uuid?")
-            continue
-        yield (uuid, slug)
+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())