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

nucliadb/writer/api/v1/services.py

@@ -19,21 +19,11 @@
 #
 from fastapi import HTTPException, Response
 from fastapi_versioning import version
-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,
-    NewEntitiesGroupRequest,
-    NewEntitiesGroupResponse,
-    OpStatusWriter,
-    SetLabelsRequest,
-    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,
@@ -47,6 +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_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
 
@@ -60,9 +61,7 @@ from nucliadb_utils.utilities import get_ingest
 )
 @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()
@@ -91,9 +90,7 @@ 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(
@@ -138,9 +135,7 @@ 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(
@@ -164,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)
 
@@ -180,21 +173,25 @@ async def delete_entities(request: Request, kbid: str, group: str):
 )
 @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:
@@ -205,16 +202,8 @@ 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(
@@ -226,21 +215,18 @@ async def set_labels(request: Request, kbid: str, labelset: str, item: LabelSet)
 )
 @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)
+
+
+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(
@@ -253,15 +239,10 @@ async def delete_labels(request: Request, kbid: str, labelset: str):
 @requires(NucliaDBRoles.WRITER)
 @version(1)
 async def set_custom_synonyms(request: Request, kbid: str, item: KnowledgeBoxSynonyms):
-    synonyms = item.to_message()
-
-    async with datamanagers.with_transaction() as txn:
-        if not datamanagers.kb.exists_kb(txn, kbid=kbid):
-            raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-
-        await datamanagers.synonyms.set(txn, kbid=kbid, synonyms=synonyms)
-        await txn.commit()
-
+    if not await datamanagers.atomic.kb.exists_kb(kbid=kbid):
+        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
+    synonyms = to_proto.kb_synonyms(item)
+    await datamanagers.atomic.synonyms.set(kbid=kbid, synonyms=synonyms)
     return Response(status_code=204)
 
 
@@ -276,7 +257,7 @@ async def set_custom_synonyms(request: Request, kbid: str, item: KnowledgeBoxSyn
 @version(1)
 async def delete_custom_synonyms(request: Request, kbid: str):
     async with datamanagers.with_transaction() as txn:
-        if not datamanagers.kb.exists_kb(txn, kbid=kbid):
+        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)

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),
+    )