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

nucliadb/writer/api/v1/services.py

@@ -19,30 +19,17 @@
 #
 from fastapi import HTTPException, Response
 from fastapi_versioning import version
-from nucliadb_protos.knowledgebox_pb2 import KnowledgeBoxID
-from nucliadb_protos.knowledgebox_pb2 import Label as LabelPB
-from nucliadb_protos.knowledgebox_pb2 import LabelSet as LabelSetPB
-from nucliadb_protos.writer_pb2 import (
-    DelEntitiesRequest,
-    DelLabelsRequest,
-    DelVectorSetRequest,
-    NewEntitiesGroupRequest,
-    NewEntitiesGroupResponse,
-    OpStatusWriter,
-    SetLabelsRequest,
-    SetSynonymsRequest,
-    UpdateEntitiesGroupRequest,
-    UpdateEntitiesGroupResponse,
-)
 from starlette.requests import Request
 
+from nucliadb.common import datamanagers
+from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
+from nucliadb.common.models_utils import to_proto
 from nucliadb.models.responses import (
     HTTPConflict,
     HTTPInternalServerError,
     HTTPNotFound,
 )
 from nucliadb.writer.api.v1.router import KB_PREFIX, api
-from nucliadb.writer.resource.vectors import create_vectorset  # type: ignore
 from nucliadb_models.entities import (
     CreateEntitiesGroupPayload,
     UpdateEntitiesGroupPayload,
@@ -50,7 +37,17 @@ from nucliadb_models.entities import (
 from nucliadb_models.labels import LabelSet
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_models.synonyms import KnowledgeBoxSynonyms
-from nucliadb_models.vectors import VectorSet
+from nucliadb_protos import writer_pb2
+from nucliadb_protos.knowledgebox_pb2 import Label as LabelPB
+from nucliadb_protos.knowledgebox_pb2 import LabelSet as LabelSetPB
+from nucliadb_protos.writer_pb2 import (
+    DelEntitiesRequest,
+    NewEntitiesGroupRequest,
+    NewEntitiesGroupResponse,
+    OpStatusWriter,
+    UpdateEntitiesGroupRequest,
+    UpdateEntitiesGroupResponse,
+)
 from nucliadb_utils.authentication import requires
 from nucliadb_utils.utilities import get_ingest
 
@@ -58,15 +55,13 @@ from nucliadb_utils.utilities import get_ingest
 @api.post(
     f"/{KB_PREFIX}/{{kbid}}/entitiesgroups",
     status_code=200,
-    name="Create Knowledge Box Entities Group",
+    summary="Create Knowledge Box Entities Group",
     tags=["Knowledge Box Services"],
     openapi_extra={"x-operation_order": 1},
 )
 @requires(NucliaDBRoles.WRITER)
 @version(1)
-async def create_entities_group(
-    request: Request, kbid: str, item: CreateEntitiesGroupPayload
-):
+async def create_entities_group(request: Request, kbid: str, item: CreateEntitiesGroupPayload):
     ingest = get_ingest()
 
     pbrequest: NewEntitiesGroupRequest = NewEntitiesGroupRequest()
@@ -95,15 +90,13 @@ async def create_entities_group(
             detail=f"Entities group {item.group} already exists in this Knowledge box",
         )
     elif status.status == NewEntitiesGroupResponse.Status.ERROR:
-        return HTTPInternalServerError(
-            detail="Error on settings entities on a Knowledge box"
-        )
+        return HTTPInternalServerError(detail="Error on settings entities on a Knowledge box")
 
 
 @api.patch(
     f"/{KB_PREFIX}/{{kbid}}/entitiesgroup/{{group}}",
     status_code=200,
-    name="Update Knowledge Box Entities Group",
+    summary="Update Knowledge Box Entities Group",
     tags=["Knowledge Box Services"],
     openapi_extra={"x-operation_order": 2},
 )
@@ -142,15 +135,13 @@ async def update_entities_group(
     elif status.status == UpdateEntitiesGroupResponse.Status.ENTITIES_GROUP_NOT_FOUND:
         return HTTPNotFound(detail="Entities group does not exist")
     elif status.status == UpdateEntitiesGroupResponse.Status.ERROR:
-        return HTTPInternalServerError(
-            detail="Error on settings entities on a Knowledge box"
-        )
+        return HTTPInternalServerError(detail="Error on settings entities on a Knowledge box")
 
 
 @api.delete(
     f"/{KB_PREFIX}/{{kbid}}/entitiesgroup/{{group}}",
     status_code=200,
-    name="Delete Knowledge Box Entities",
+    summary="Delete Knowledge Box Entities",
     tags=["Knowledge Box Services"],
     openapi_extra={"x-operation_order": 3},
 )
@@ -168,9 +159,7 @@ async def delete_entities(request: Request, kbid: str, group: str):
     elif status.status == OpStatusWriter.Status.NOTFOUND:
         raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
     elif status.status == OpStatusWriter.Status.ERROR:
-        raise HTTPException(
-            status_code=500, detail="Error on deleting entities from a Knowledge box"
-        )
+        raise HTTPException(status_code=500, detail="Error on deleting entities from a Knowledge box")
 
     return Response(status_code=204)
 
@@ -178,27 +167,31 @@ async def delete_entities(request: Request, kbid: str, group: str):
 @api.post(
     f"/{KB_PREFIX}/{{kbid}}/labelset/{{labelset}}",
     status_code=200,
-    name="Set Knowledge Box Labels",
+    summary="Set Knowledge Box Labels",
     tags=["Knowledge Box Services"],
     openapi_extra={"x-operation_order": 1},
 )
 @requires(NucliaDBRoles.WRITER)
 @version(1)
-async def set_labels(request: Request, kbid: str, labelset: str, item: LabelSet):
-    ingest = get_ingest()
-    pbrequest: SetLabelsRequest = SetLabelsRequest(id=labelset)
-    pbrequest.kb.uuid = kbid
-
-    if item.title:
-        pbrequest.labelset.title = item.title
+async def set_labelset_endpoint(request: Request, kbid: str, labelset: str, item: LabelSet):
+    try:
+        await set_labelset(kbid, labelset, item)
+    except KnowledgeBoxNotFound:
+        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
 
-    if item.color:
-        pbrequest.labelset.color = item.color
 
-    pbrequest.labelset.multiple = item.multiple
+async def set_labelset(kbid: str, labelset_id: str, item: LabelSet):
+    kb_exists = await datamanagers.atomic.kb.exists_kb(kbid=kbid)
+    if not kb_exists:
+        raise KnowledgeBoxNotFound()
+    labelset = writer_pb2.LabelSet()
+    if item.title is not None:
+        labelset.title = item.title
+    if item.color is not None:
+        labelset.color = item.color
+    labelset.multiple = item.multiple
     for kind in item.kind:
-        pbrequest.labelset.kind.append(LabelSetPB.LabelSetKind.Value(kind))
-
+        labelset.kind.append(LabelSetPB.LabelSetKind.Value(kind))
     for label_input in item.labels:
         lbl = LabelPB()
         if label_input.uri:
@@ -209,126 +202,65 @@ async def set_labels(request: Request, kbid: str, labelset: str, item: LabelSet)
             lbl.related = label_input.related
         if label_input.title:
             lbl.title = label_input.title
-        pbrequest.labelset.labels.append(lbl)
-    status: OpStatusWriter = await ingest.SetLabels(pbrequest)  # type: ignore
-    if status.status == OpStatusWriter.Status.OK:
-        return None
-    elif status.status == OpStatusWriter.Status.NOTFOUND:
-        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    elif status.status == OpStatusWriter.Status.ERROR:
-        raise HTTPException(
-            status_code=500, detail="Error on settings labels on a Knowledge box"
-        )
+        labelset.labels.append(lbl)
+    await datamanagers.atomic.labelset.set(kbid=kbid, labelset_id=labelset_id, labelset=labelset)
 
 
 @api.delete(
     f"/{KB_PREFIX}/{{kbid}}/labelset/{{labelset}}",
     status_code=200,
-    name="Delete Knowledge Box Label",
+    summary="Delete Knowledge Box Label",
     tags=["Knowledge Box Services"],
     openapi_extra={"x-operation_order": 3},
 )
 @requires(NucliaDBRoles.WRITER)
 @version(1)
-async def delete_labels(request: Request, kbid: str, labelset: str):
-    ingest = get_ingest()
-    pbrequest: DelLabelsRequest = DelLabelsRequest()
-    pbrequest.kb.uuid = kbid
-    pbrequest.id = labelset
-    status: OpStatusWriter = await ingest.DelLabels(pbrequest)  # type: ignore
-    if status.status == OpStatusWriter.Status.OK:
-        return None
-    elif status.status == OpStatusWriter.Status.NOTFOUND:
+async def delete_labelset_endpoint(request: Request, kbid: str, labelset: str):
+    try:
+        await delete_labelset(kbid, labelset)
+    except KnowledgeBoxNotFound:
         raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    elif status.status == OpStatusWriter.Status.ERROR:
-        raise HTTPException(
-            status_code=500, detail="Error on deleting labels from a Knowledge box"
-        )
-    return Response(status_code=204)
-
-
-@api.post(
-    f"/{KB_PREFIX}/{{kbid}}/vectorset/{{vectorset}}",
-    status_code=200,
-    name="Set Knowledge Box VectorSet",
-    tags=["Knowledge Box Services"],
-    openapi_extra={"x-operation_order": 1},
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def set_vectorset(request: Request, kbid: str, vectorset: str, item: VectorSet):
-    await create_vectorset(kbid, vectorset, item.dimension, similarity=item.similarity)
 
 
-@api.delete(
-    f"/{KB_PREFIX}/{{kbid}}/vectorset/{{vectorset}}",
-    status_code=200,
-    name="Delete Knowledge Box VectorSet",
-    tags=["Knowledge Box Services"],
-    openapi_extra={"x-operation_order": 3},
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def delete_vectorset(request: Request, kbid: str, vectorset: str):
-    ingest = get_ingest()
-    pbrequest: DelVectorSetRequest = DelVectorSetRequest()
-    pbrequest.kb.uuid = kbid
-    pbrequest.vectorset = vectorset
-    status: OpStatusWriter = await ingest.DelVectorSet(pbrequest)  # type: ignore
-    if status.status == OpStatusWriter.Status.OK:
-        return None
-    elif status.status == OpStatusWriter.Status.NOTFOUND:
-        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    elif status.status == OpStatusWriter.Status.ERROR:
-        raise HTTPException(
-            status_code=500, detail="Error on deleting vectorset from a Knowledge box"
-        )
-    return Response(status_code=204)
+async def delete_labelset(kbid: str, labelset_id: str):
+    kb_exists = await datamanagers.atomic.kb.exists_kb(kbid=kbid)
+    if not kb_exists:
+        raise KnowledgeBoxNotFound()
+    await datamanagers.atomic.labelset.delete(kbid=kbid, labelset_id=labelset_id)
 
 
 @api.put(
     f"/{KB_PREFIX}/{{kbid}}/custom-synonyms",
     status_code=204,
-    name="Set Knowledge Box Custom Synonyms",
+    summary="Set Knowledge Box Custom Synonyms",
     tags=["Knowledge Box Services"],
     openapi_extra={"x-operation_order": 1},
 )
 @requires(NucliaDBRoles.WRITER)
 @version(1)
 async def set_custom_synonyms(request: Request, kbid: str, item: KnowledgeBoxSynonyms):
-    ingest = get_ingest()
-    pbrequest = SetSynonymsRequest()
-    pbrequest.kbid.uuid = kbid
-    pbrequest.synonyms.CopyFrom(item.to_message())
-    status: OpStatusWriter = await ingest.SetSynonyms(pbrequest)  # type: ignore
-    if status.status == OpStatusWriter.Status.OK:
-        return Response(status_code=204)
-    elif status.status == OpStatusWriter.Status.NOTFOUND:
+    if not await datamanagers.atomic.kb.exists_kb(kbid=kbid):
         raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    elif status.status == OpStatusWriter.Status.ERROR:
-        raise HTTPException(
-            status_code=500, detail="Error setting synonyms of a Knowledge box"
-        )
+    synonyms = to_proto.kb_synonyms(item)
+    await datamanagers.atomic.synonyms.set(kbid=kbid, synonyms=synonyms)
+    return Response(status_code=204)
 
 
 @api.delete(
     f"/{KB_PREFIX}/{{kbid}}/custom-synonyms",
     status_code=204,
-    name="Delete Knowledge Box Custom Synonyms",
+    summary="Delete Knowledge Box Custom Synonyms",
     tags=["Knowledge Box Services"],
     openapi_extra={"x-operation_order": 3},
 )
 @requires(NucliaDBRoles.WRITER)
 @version(1)
 async def delete_custom_synonyms(request: Request, kbid: str):
-    ingest = get_ingest()
-    pbrequest = KnowledgeBoxID(uuid=kbid)
-    status: OpStatusWriter = await ingest.DelSynonyms(pbrequest)  # type: ignore
-    if status.status == OpStatusWriter.Status.OK:
-        return Response(status_code=204)
-    elif status.status == OpStatusWriter.Status.NOTFOUND:
-        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    elif status.status == OpStatusWriter.Status.ERROR:
-        raise HTTPException(
-            status_code=500, detail="Error deleting synonyms of a Knowledge box"
-        )
+    async with datamanagers.with_transaction() as txn:
+        if not await datamanagers.kb.exists_kb(txn, kbid=kbid):
+            raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
+
+        await datamanagers.synonyms.delete(txn, kbid=kbid)
+        await txn.commit()
+
+    return Response(status_code=204)

nucliadb/writer/api/v1/slug.py

@@ -0,0 +1,61 @@
+# 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/>.
+#
+import contextlib
+
+from fastapi import HTTPException
+
+from nucliadb.common import datamanagers, locking
+
+
+@contextlib.asynccontextmanager
+async def noop_context_manager():
+    """
+    This is used for requests where slug is not set by the user and we don't need to
+    care about uniqueness
+    """
+    yield
+
+
+@contextlib.asynccontextmanager
+async def ensure_slug_uniqueness(kbid: str, slug: str):
+    """
+    Make sure slug is unique by:
+    - First check if the slug is already taken by another existing resource
+    - Otherwise, use the transaction lock to prevent from multiple concurrent
+      create resource requests having the same slug.
+    """
+    try:
+        async with locking.distributed_lock(
+            key=locking.RESOURCE_CREATION_SLUG_LOCK.format(kbid=kbid, resource_slug=slug),
+            # We don't want to wait here. If the lock exists, just raise exception
+            lock_timeout=0.0,
+            # Matches aprox the max amount of time that the ingest can take
+            # to ingest a broker message from the writer.
+            expire_timeout=60.0,
+            # We don't want to refresh it here
+            refresh_timeout=120.0,
+        ):
+            if await datamanagers.atomic.resources.slug_exists(kbid=kbid, slug=slug):
+                raise HTTPException(status_code=409, detail=f"Resource slug {slug} already exists")
+            yield
+    except locking.ResourceLocked:
+        raise HTTPException(
+            status_code=409, detail=f"Another resource with the same {slug} is already being ingested"
+        )

nucliadb/writer/api/v1/transaction.py

@@ -0,0 +1,67 @@
+# 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/>.
+#
+
+import backoff
+from fastapi import HTTPException
+
+from nucliadb_protos.writer_pb2 import BrokerMessage
+from nucliadb_utils import const
+from nucliadb_utils.transaction import (
+    MaxTransactionSizeExceededError,
+    StreamingServerError,
+    TransactionCommitTimeoutError,
+)
+from nucliadb_utils.utilities import get_transaction_utility
+
+
+async def commit(writer: BrokerMessage, partition: int, wait: bool = True) -> None:
+    try:
+        await transaction_commit(writer, partition, wait)
+    except TransactionCommitTimeoutError:
+        raise HTTPException(
+            status_code=501,
+            detail="Inconsistent write. This resource will not be processed and may not be stored.",
+        )
+    except MaxTransactionSizeExceededError:
+        raise HTTPException(
+            status_code=413,
+            detail="Transaction size exceeded. The resource is too large to be stored. Consider using file fields or split into multiple requests.",
+        )
+    except StreamingServerError:
+        raise HTTPException(
+            status_code=504,
+            detail="Timeout waiting for the streaming server to respond. Please back off and retry.",
+        )
+
+
+@backoff.on_exception(
+    backoff.expo,
+    (StreamingServerError,),
+    jitter=backoff.random_jitter,
+    max_tries=3,
+)
+async def transaction_commit(writer: BrokerMessage, partition: int, wait: bool = True):
+    transaction = get_transaction_utility()
+    await transaction.commit(
+        writer,
+        partition,
+        wait=wait,
+        target_subject=const.Streams.INGEST.subject.format(partition=partition),
+    )