nucliadb 6.7.2.post4874__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/ingest/serialize.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional, Union
 
 import nucliadb_models as models
 from nucliadb.common import datamanagers
@@ -32,15 +31,22 @@ from nucliadb.ingest.fields.link import Link
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.ingest.orm.resource import Resource as ORMResource
 from nucliadb_models.common import FieldTypeName
+from nucliadb_models.metadata import Extra, Origin, Relation
 from nucliadb_models.resource import (
     ConversationFieldData,
     ConversationFieldExtractedData,
     Error,
     ExtractedDataType,
     ExtractedDataTypeName,
+    ExtractedText,
+    FieldComputedMetadata,
+    FieldQuestionAnswers,
+    FileExtractedData,
     FileFieldData,
     FileFieldExtractedData,
     GenericFieldData,
+    LargeComputedMetadata,
+    LinkExtractedData,
     LinkFieldData,
     LinkFieldExtractedData,
     QueueType,
@@ -48,6 +54,7 @@ from nucliadb_models.resource import (
     ResourceData,
     TextFieldData,
     TextFieldExtractedData,
+    VectorObject,
 )
 from nucliadb_models.search import ResourceProperties
 from nucliadb_models.security import ResourceSecurity
@@ -56,84 +63,15 @@ from nucliadb_protos.writer_pb2 import FieldStatus
 from nucliadb_utils.utilities import get_storage
 
 
-async def set_resource_field_extracted_data(
-    field: Field,
-    field_data: ExtractedDataType,
-    field_type_name: FieldTypeName,
-    wanted_extracted_data: list[ExtractedDataTypeName],
-) -> None:
-    if field_data is None:
-        return
-
-    if ExtractedDataTypeName.TEXT in wanted_extracted_data:
-        data_et = await field.get_extracted_text()
-        if data_et is not None:
-            field_data.text = from_proto.extracted_text(data_et)
-
-    metadata_wanted = ExtractedDataTypeName.METADATA in wanted_extracted_data
-    shortened_metadata_wanted = ExtractedDataTypeName.SHORTENED_METADATA in wanted_extracted_data
-    if metadata_wanted or shortened_metadata_wanted:
-        data_fcm = await field.get_field_metadata()
-
-        if data_fcm is not None:
-            field_data.metadata = from_proto.field_computed_metadata(
-                data_fcm, shortened=shortened_metadata_wanted and not metadata_wanted
-            )
-
-    if ExtractedDataTypeName.LARGE_METADATA in wanted_extracted_data:
-        data_lcm = await field.get_large_field_metadata()
-        if data_lcm is not None:
-            field_data.large_metadata = from_proto.large_computed_metadata(data_lcm)
-
-    if ExtractedDataTypeName.VECTOR in wanted_extracted_data:
-        # XXX: our extracted API is not vectorset-compatible, so we'll get the
-        # first vectorset and return the values. Ideally, we should provide a
-        # way to select a vectorset
-        vectorset_id = None
-        async with datamanagers.with_ro_transaction() as txn:
-            async for vectorset_id, vs in datamanagers.vectorsets.iter(
-                txn=txn,
-                kbid=field.resource.kb.kbid,
-            ):
-                break
-        assert vectorset_id is not None, "All KBs must have at least a vectorset"
-        data_vec = await field.get_vectors(vectorset_id, vs.storage_key_kind)
-        if data_vec is not None:
-            field_data.vectors = from_proto.vector_object(data_vec)
-
-    if ExtractedDataTypeName.QA in wanted_extracted_data:
-        qa = await field.get_question_answers()
-        if qa is not None:
-            field_data.question_answers = from_proto.field_question_answers(qa)
-
-    if (
-        isinstance(field, File)
-        and isinstance(field_data, FileFieldExtractedData)
-        and ExtractedDataTypeName.FILE in wanted_extracted_data
-    ):
-        data_fed = await field.get_file_extracted_data()
-        if data_fed is not None:
-            field_data.file = from_proto.file_extracted_data(data_fed)
-
-    if (
-        isinstance(field, Link)
-        and isinstance(field_data, LinkFieldExtractedData)
-        and ExtractedDataTypeName.LINK in wanted_extracted_data
-    ):
-        data_led = await field.get_link_extracted_data()
-        if data_led is not None:
-            field_data.link = from_proto.link_extracted_data(data_led)
-
-
 async def serialize(
     kbid: str,
-    rid: Optional[str],
+    rid: str | None,
     show: list[ResourceProperties],
     field_type_filter: list[FieldTypeName],
     extracted: list[ExtractedDataTypeName],
-    service_name: Optional[str] = None,
-    slug: Optional[str] = None,
-) -> Optional[Resource]:
+    service_name: str | None = None,
+    slug: str | None = None,
+) -> Resource | None:
     driver = get_driver()
     async with driver.ro_transaction() as txn:
         return await managed_serialize(
@@ -148,45 +86,56 @@ async def serialize(
         )
 
 
-async def serialize_field_errors(
-    field: Field,
-    serialized: Union[
-        TextFieldData, FileFieldData, LinkFieldData, ConversationFieldData, GenericFieldData
-    ],
-):
-    status = await field.get_status()
-    if status is None:
-        status = FieldStatus()
-    serialized.status = status.Status.Name(status.status)
-    if status.errors:
-        serialized.errors = []
-        for error in status.errors:
-            serialized.errors.append(
-                Error(
-                    body=error.source_error.error,
-                    code=error.source_error.code,
-                    code_str=writer_pb2.Error.ErrorCode.Name(error.source_error.code),
-                    created=error.created.ToDatetime(),
-                    severity=writer_pb2.Error.Severity.Name(error.source_error.severity),
-                )
-            )
-        serialized.error = serialized.errors[-1]
-
-
 async def managed_serialize(
     txn: Transaction,
     kbid: str,
-    rid: Optional[str],
+    rid: str | None,
     show: list[ResourceProperties],
     field_type_filter: list[FieldTypeName],
     extracted: list[ExtractedDataTypeName],
-    service_name: Optional[str] = None,
-    slug: Optional[str] = None,
-) -> Optional[Resource]:
+    service_name: str | None = None,
+    slug: str | None = None,
+) -> Resource | None:
     orm_resource = await get_orm_resource(txn, kbid, rid=rid, slug=slug, service_name=service_name)
     if orm_resource is None:
         return None
 
+    return await serialize_resource(orm_resource, show, field_type_filter, extracted)
+
+
+async def get_orm_resource(
+    txn: Transaction,
+    kbid: str,
+    rid: str | None,
+    slug: str | None = None,
+    service_name: str | None = None,
+) -> ORMResource | None:
+    storage = await get_storage(service_name=service_name)
+
+    kb = KnowledgeBox(txn, storage, kbid)
+
+    if rid is None:
+        if slug is None:
+            raise ValueError("Either rid or slug parameters should be used")
+
+        rid = await kb.get_resource_uuid_by_slug(slug)
+        if rid is None:
+            # Could not find resource uuid from slug
+            return None
+
+    orm_resource = await kb.get(rid)
+    if orm_resource is None:
+        return None
+
+    return orm_resource
+
+
+async def serialize_resource(
+    orm_resource: ORMResource,
+    show: list[ResourceProperties],
+    field_type_filter: list[FieldTypeName],
+    extracted: list[ExtractedDataTypeName],
+) -> Resource:
     resource = Resource(id=orm_resource.uuid)
 
     include_values = ResourceProperties.VALUES in show
@@ -230,29 +179,18 @@ async def managed_serialize(
             resource.queue = QueueType[orm_resource.basic.QueueType.Name(orm_resource.basic.queue)]
 
             if ResourceProperties.RELATIONS in show:
-                relations = await orm_resource.get_user_relations()
-                resource.usermetadata.relations = [
-                    from_proto.relation(rel) for rel in relations.relations
-                ]
+                resource.usermetadata.relations = await serialize_user_relations(orm_resource)
 
     if ResourceProperties.ORIGIN in show:
-        await orm_resource.get_origin()
-        if orm_resource.origin is not None:
-            resource.origin = from_proto.origin(orm_resource.origin)
+        resource.origin = await serialize_origin(orm_resource)
 
     if ResourceProperties.EXTRA in show:
-        await orm_resource.get_extra()
-        if orm_resource.extra is not None:
-            resource.extra = from_proto.extra(orm_resource.extra)
+        resource.extra = await serialize_extra(orm_resource)
 
     include_errors = ResourceProperties.ERRORS in show
 
     if ResourceProperties.SECURITY in show:
-        await orm_resource.get_security()
-        resource.security = ResourceSecurity(access_groups=[])
-        if orm_resource.security is not None:
-            for gid in orm_resource.security.access_groups:
-                resource.security.access_groups.append(gid)
+        resource.security = await serialize_security(orm_resource)
 
     if (field_type_filter and (include_values or include_extracted_data)) or include_errors:
         await orm_resource.get_fields()
@@ -360,38 +298,157 @@ async def managed_serialize(
     return resource
 
 
-async def get_orm_resource(
-    txn: Transaction,
-    kbid: str,
-    rid: Optional[str],
-    slug: Optional[str] = None,
-    service_name: Optional[str] = None,
-) -> Optional[ORMResource]:
-    storage = await get_storage(service_name=service_name)
+async def serialize_origin(resource: ORMResource) -> Origin | None:
+    origin = await resource.get_origin()
+    if origin is None:
+        return None
 
-    kb = KnowledgeBox(txn, storage, kbid)
+    return from_proto.origin(origin)
 
-    if rid is None:
-        if slug is None:
-            raise ValueError("Either rid or slug parameters should be used")
 
-        rid = await kb.get_resource_uuid_by_slug(slug)
-        if rid is None:
-            # Could not find resource uuid from slug
-            return None
+async def serialize_extra(resource: ORMResource) -> Extra | None:
+    extra = await resource.get_extra()
+    if extra is None:
+        return None
+    return from_proto.extra(extra)
 
-    orm_resource = await kb.get(rid)
-    if orm_resource is None:
+
+async def serialize_user_relations(resource: ORMResource) -> list[Relation]:
+    relations = await resource.get_user_relations()
+    return [from_proto.relation(rel) for rel in relations.relations]
+
+
+async def serialize_security(resource: ORMResource) -> ResourceSecurity:
+    security = ResourceSecurity(access_groups=[])
+
+    security_pb = await resource.get_security()
+    if security_pb is not None:
+        for gid in security_pb.access_groups:
+            security.access_groups.append(gid)
+
+    return security
+
+
+async def serialize_field_errors(
+    field: Field,
+    serialized: (
+        TextFieldData | FileFieldData | LinkFieldData | ConversationFieldData | GenericFieldData
+    ),
+):
+    status = await field.get_status()
+    if status is None:
+        status = FieldStatus()
+    serialized.status = status.Status.Name(status.status)
+    if status.errors:
+        serialized.errors = []
+        for error in status.errors:
+            serialized.errors.append(
+                Error(
+                    body=error.source_error.error,
+                    code=error.source_error.code,
+                    code_str=writer_pb2.Error.ErrorCode.Name(error.source_error.code),
+                    created=error.created.ToDatetime(),
+                    severity=writer_pb2.Error.Severity.Name(error.source_error.severity),
+                )
+            )
+        serialized.error = serialized.errors[-1]
+
+
+async def set_resource_field_extracted_data(
+    field: Field,
+    field_data: ExtractedDataType,
+    field_type_name: FieldTypeName,
+    wanted_extracted_data: list[ExtractedDataTypeName],
+) -> None:
+    if field_data is None:
+        return
+
+    if ExtractedDataTypeName.TEXT in wanted_extracted_data:
+        field_data.text = await serialize_extracted_text(field)
+
+    metadata_wanted = ExtractedDataTypeName.METADATA in wanted_extracted_data
+    shortened_metadata_wanted = ExtractedDataTypeName.SHORTENED_METADATA in wanted_extracted_data
+    if metadata_wanted or shortened_metadata_wanted:
+        field_data.metadata = await serialize_extracted_metadata(
+            field, shortened=shortened_metadata_wanted and not metadata_wanted
+        )
+
+    if ExtractedDataTypeName.LARGE_METADATA in wanted_extracted_data:
+        field_data.large_metadata = await serialize_extracted_large_metadata(field)
+
+    if ExtractedDataTypeName.VECTOR in wanted_extracted_data:
+        field_data.vectors = await serialize_extracted_vectors(field)
+
+    if ExtractedDataTypeName.QA in wanted_extracted_data:
+        field_data.question_answers = await serialize_extracted_question_answers(field)
+
+    if (
+        isinstance(field, File)
+        and isinstance(field_data, FileFieldExtractedData)
+        and ExtractedDataTypeName.FILE in wanted_extracted_data
+    ):
+        field_data.file = await serialize_file_extracted_data(field)
+
+    if (
+        isinstance(field, Link)
+        and isinstance(field_data, LinkFieldExtractedData)
+        and ExtractedDataTypeName.LINK in wanted_extracted_data
+    ):
+        field_data.link = await serialize_link_extracted_data(field)
+
+
+async def serialize_extracted_text(field: Field) -> ExtractedText | None:
+    data_et = await field.get_extracted_text()
+    if data_et is None:
         return None
+    return from_proto.extracted_text(data_et)
 
-    return orm_resource
 
+async def serialize_extracted_metadata(field: Field, *, shortened: bool) -> FieldComputedMetadata | None:
+    data_fcm = await field.get_field_metadata()
+    if data_fcm is None:
+        return None
+    return from_proto.field_computed_metadata(data_fcm, shortened)
 
-async def get_resource_uuid_by_slug(
-    kbid: str, slug: str, service_name: Optional[str] = None
-) -> Optional[str]:
-    storage = await get_storage(service_name=service_name)
-    driver = get_driver()
-    async with driver.ro_transaction() as txn:
-        kb = KnowledgeBox(txn, storage, kbid)
-        return await kb.get_resource_uuid_by_slug(slug)
+
+async def serialize_extracted_large_metadata(field: Field) -> LargeComputedMetadata | None:
+    data_lcm = await field.get_large_field_metadata()
+    if data_lcm is None:
+        return None
+    return from_proto.large_computed_metadata(data_lcm)
+
+
+async def serialize_extracted_vectors(field: Field) -> VectorObject | None:
+    # XXX: our extracted API is not vectorset-compatible, so we'll get the
+    # first vectorset and return the values. Ideally, we should provide a
+    # way to select a vectorset
+    vectorset_id = None
+    async with datamanagers.with_ro_transaction() as txn:
+        async for vectorset_id, vs in datamanagers.vectorsets.iter(txn=txn, kbid=field.kbid):
+            break
+    assert vectorset_id is not None, "All KBs must have at least a vectorset"
+    data_vec = await field.get_vectors(vectorset_id, vs.storage_key_kind)
+    if data_vec is None:
+        return None
+    return from_proto.vector_object(data_vec)
+
+
+async def serialize_extracted_question_answers(field: Field) -> FieldQuestionAnswers | None:
+    qa = await field.get_question_answers()
+    if qa is None:
+        return None
+    return from_proto.field_question_answers(qa)
+
+
+async def serialize_file_extracted_data(field: File) -> FileExtractedData | None:
+    data_fed = await field.get_file_extracted_data()
+    if data_fed is None:
+        return None
+    return from_proto.file_extracted_data(data_fed)
+
+
+async def serialize_link_extracted_data(field: Link) -> LinkExtractedData | None:
+    data_led = await field.get_link_extracted_data()
+    if data_led is None:
+        return None
+    return from_proto.link_extracted_data(data_led)

nucliadb/ingest/service/writer.py

@@ -18,12 +18,11 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import uuid
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
 
 from nucliadb.backups import tasks as backup_tasks
 from nucliadb.backups import utils as backup_utils
 from nucliadb.common import datamanagers
-from nucliadb.common.cluster.exceptions import AlreadyExists, EntitiesGroupNotFound
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
 from nucliadb.common.external_index_providers.exceptions import ExternalIndexCreationError
@@ -49,7 +48,6 @@ from nucliadb_protos.knowledgebox_pb2 import (
 )
 from nucliadb_protos.writer_pb2 import (
     BrokerMessage,
-    DelEntitiesRequest,
     GetEntitiesGroupRequest,
     GetEntitiesGroupResponse,
     GetEntitiesRequest,
@@ -58,12 +56,7 @@ from nucliadb_protos.writer_pb2 import (
     IndexStatus,
     ListEntitiesGroupsRequest,
     ListEntitiesGroupsResponse,
-    NewEntitiesGroupRequest,
-    NewEntitiesGroupResponse,
     OpStatusWriter,
-    SetEntitiesRequest,
-    UpdateEntitiesGroupRequest,
-    UpdateEntitiesGroupResponse,
     WriterStatusRequest,
     WriterStatusResponse,
 )
@@ -126,6 +119,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
                 external_index_provider=request.external_index_provider,
                 hidden_resources_enabled=request.hidden_resources_enabled,
                 hidden_resources_hide_on_creation=request.hidden_resources_hide_on_creation,
+                prewarm_enabled=request.prewarm_enabled,
             )
 
         except KnowledgeBoxConflict:
@@ -167,11 +161,17 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
             )
 
         try:
-            async with self.driver.rw_transaction() as txn:
-                kbid = await KnowledgeBoxORM.update(
-                    txn, uuid=request.uuid, slug=request.slug, config=request.config
-                )
-                await txn.commit()
+            kbid = await KnowledgeBoxORM.update(
+                self.driver,
+                kbid=request.uuid,
+                slug=request.slug,
+                title=request.config.title or None,
+                description=request.config.description or None,
+                external_index_provider=request.config.external_index_provider or None,
+                hidden_resources_enabled=request.config.hidden_resources_enabled,
+                hidden_resources_hide_on_creation=request.config.hidden_resources_hide_on_creation,
+                prewarm_enabled=request.config.prewarm_enabled,
+            )
         except KnowledgeBoxNotFound:
             return UpdateKnowledgeBoxResponse(status=KnowledgeBoxResponseStatus.NOTFOUND)
         except Exception:
@@ -217,29 +217,6 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
             logger.info(f"Processed {message.uuid}")
         return response
 
-    async def NewEntitiesGroup(  # type: ignore
-        self, request: NewEntitiesGroupRequest, context=None
-    ) -> NewEntitiesGroupResponse:
-        response = NewEntitiesGroupResponse()
-        async with self.driver.ro_transaction() as ro_txn:
-            kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
-            if kbobj is None:
-                response.status = NewEntitiesGroupResponse.Status.KB_NOT_FOUND
-                return response
-
-        async with self.driver.rw_transaction() as txn:
-            kbobj.txn = txn
-            entities_manager = EntitiesManager(kbobj, txn)
-            try:
-                await entities_manager.create_entities_group(request.group, request.entities)
-            except AlreadyExists:
-                response.status = NewEntitiesGroupResponse.Status.ALREADY_EXISTS
-                return response
-
-            await txn.commit()
-            response.status = NewEntitiesGroupResponse.Status.OK
-            return response
-
     async def GetEntities(  # type: ignore
         self, request: GetEntitiesRequest, context=None
     ) -> GetEntitiesResponse:
@@ -313,81 +290,6 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
 
             return response
 
-    async def SetEntities(self, request: SetEntitiesRequest, context=None) -> OpStatusWriter:  # type: ignore
-        response = OpStatusWriter()
-        async with self.driver.ro_transaction() as ro_txn:
-            kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
-            if kbobj is None:
-                response.status = OpStatusWriter.Status.NOTFOUND
-                return response
-
-        async with self.driver.rw_transaction() as txn:
-            kbobj.txn = txn
-            entities_manager = EntitiesManager(kbobj, txn)
-            try:
-                await entities_manager.set_entities_group(request.group, request.entities)
-            except Exception as e:
-                errors.capture_exception(e)
-                logger.error("Error in ingest gRPC servicer", exc_info=True)
-                response.status = OpStatusWriter.Status.ERROR
-            else:
-                response.status = OpStatusWriter.Status.OK
-                await txn.commit()
-            return response
-
-    async def UpdateEntitiesGroup(  # type: ignore
-        self, request: UpdateEntitiesGroupRequest, context=None
-    ) -> UpdateEntitiesGroupResponse:
-        response = UpdateEntitiesGroupResponse()
-        async with self.driver.ro_transaction() as ro_txn:
-            kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
-            if kbobj is None:
-                response.status = UpdateEntitiesGroupResponse.Status.KB_NOT_FOUND
-                return response
-
-        async with self.driver.rw_transaction() as txn:
-            kbobj.txn = txn
-            entities_manager = EntitiesManager(kbobj, txn)
-            try:
-                await entities_manager.set_entities_group_metadata(
-                    request.group,
-                    title=request.title,
-                    color=request.color,
-                )
-                updates = {**request.add, **request.update}
-                await entities_manager.update_entities(request.group, updates)
-                await entities_manager.delete_entities(request.group, request.delete)  # type: ignore
-            except EntitiesGroupNotFound:
-                response.status = UpdateEntitiesGroupResponse.Status.ENTITIES_GROUP_NOT_FOUND
-                return response
-
-            await txn.commit()
-            response.status = UpdateEntitiesGroupResponse.Status.OK
-            return response
-
-    async def DelEntities(self, request: DelEntitiesRequest, context=None) -> OpStatusWriter:  # type: ignore
-        response = OpStatusWriter()
-
-        async with self.driver.ro_transaction() as ro_txn:
-            kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
-            if kbobj is None:
-                response.status = OpStatusWriter.Status.NOTFOUND
-                return response
-
-        async with self.driver.rw_transaction() as txn:
-            kbobj.txn = txn
-            entities_manager = EntitiesManager(kbobj, txn)
-            try:
-                await entities_manager.delete_entities_group(request.group)
-            except Exception as e:
-                errors.capture_exception(e)
-                logger.error("Error in ingest gRPC servicer", exc_info=True)
-                response.status = OpStatusWriter.Status.ERROR
-            else:
-                await txn.commit()
-                response.status = OpStatusWriter.Status.OK
-            return response
-
     async def Status(  # type: ignore
         self, request: WriterStatusRequest, context=None
     ) -> WriterStatusResponse:
@@ -406,8 +308,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
 
     async def Index(self, request: IndexResource, context=None) -> IndexStatus:  # type: ignore
         async with self.driver.ro_transaction() as txn:
-            kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
-            resobj = ResourceORM(txn, self.storage, kbobj, request.rid)
+            resobj = ResourceORM(txn, self.storage, request.kbid, request.rid)
             bm = await generate_broker_message(resobj)
             transaction = get_transaction_utility()
             partitioning = get_partitioning()
@@ -421,7 +322,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
         try:
             async with self.driver.rw_transaction() as txn:
                 kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
-                resobj = ResourceORM(txn, self.storage, kbobj, request.rid)
+                resobj = ResourceORM(txn, self.storage, request.kbid, request.rid)
                 resobj.disable_vectors = not request.reindex_vectors
                 index_message = await get_resource_index_message(resobj, reindex=True)
                 shard = await self.proc.get_or_assign_resource_shard(txn, kbobj, request.rid)