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

nucliadb/ingest/orm/resource.py

@@ -22,8 +22,9 @@ from __future__ import annotations
 import asyncio
 import logging
 from collections import defaultdict
+from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor
-from typing import TYPE_CHECKING, Any, Optional, Sequence, Type
+from typing import Any
 
 from nucliadb.common import datamanagers
 from nucliadb.common.datamanagers.resources import KB_RESOURCE_SLUG
@@ -68,13 +69,11 @@ from nucliadb_protos.resources_pb2 import Origin as PBOrigin
 from nucliadb_protos.resources_pb2 import Relations as PBRelations
 from nucliadb_protos.writer_pb2 import BrokerMessage
 from nucliadb_utils.storages.storage import Storage
-
-if TYPE_CHECKING:  # pragma: no cover
-    from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
+from nucliadb_utils.utilities import get_storage
 
 logger = logging.getLogger(__name__)
 
-KB_FIELDS: dict[int, Type] = {
+KB_FIELDS: dict[int, type] = {
     FieldType.TEXT: Text,
     FieldType.FILE: File,
     FieldType.LINK: Link,
@@ -104,40 +103,55 @@ class Resource:
         self,
         txn: Transaction,
         storage: Storage,
-        kb: KnowledgeBox,
+        kbid: str,
         uuid: str,
-        basic: Optional[PBBasic] = None,
+        basic: PBBasic | None = None,
         disable_vectors: bool = True,
     ):
         self.fields: dict[tuple[FieldType.ValueType, str], Field] = {}
         self.conversations: dict[int, PBConversation] = {}
-        self.relations: Optional[PBRelations] = None
-        self.all_fields_keys: Optional[list[tuple[FieldType.ValueType, str]]] = None
-        self.origin: Optional[PBOrigin] = None
-        self.extra: Optional[PBExtra] = None
-        self.security: Optional[utils_pb2.Security] = None
+        self.relations: PBRelations | None = None
+        self.all_fields_keys: list[tuple[FieldType.ValueType, str]] | None = None
+        self.origin: PBOrigin | None = None
+        self.extra: PBExtra | None = None
+        self.security: utils_pb2.Security | None = None
         self.modified: bool = False
         self._modified_extracted_text: list[FieldID] = []
 
         self.txn = txn
         self.storage = storage
-        self.kb = kb
+        self.kbid = kbid
         self.uuid = uuid
         self.basic = basic
         self.disable_vectors = disable_vectors
-        self._previous_status: Optional[Metadata.Status.ValueType] = None
-        self.user_relations: Optional[PBRelations] = None
+        self._previous_status: Metadata.Status.ValueType | None = None
+        self.user_relations: PBRelations | None = None
         self.locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
 
+    @classmethod
+    async def get(cls, txn: Transaction, kbid: str, rid: str) -> Resource | None:
+        basic = await datamanagers.resources.get_basic(txn, kbid=kbid, rid=rid)
+        if basic is None:
+            return None
+        storage = await get_storage()
+        return cls(
+            txn=txn,
+            storage=storage,
+            kbid=kbid,
+            uuid=rid,
+            basic=basic,
+            disable_vectors=False,
+        )
+
     async def set_slug(self):
         basic = await self.get_basic()
-        new_key = KB_RESOURCE_SLUG.format(kbid=self.kb.kbid, slug=basic.slug)
+        new_key = KB_RESOURCE_SLUG.format(kbid=self.kbid, slug=basic.slug)
         await self.txn.set(new_key, self.uuid.encode())
 
     # Basic
-    async def get_basic(self) -> Optional[PBBasic]:
+    async def get_basic(self) -> PBBasic:
         if self.basic is None:
-            basic = await datamanagers.resources.get_basic(self.txn, kbid=self.kb.kbid, rid=self.uuid)
+            basic = await datamanagers.resources.get_basic(self.txn, kbid=self.kbid, rid=self.uuid)
             self.basic = basic if basic is not None else PBBasic()
         return self.basic
 
@@ -159,7 +173,7 @@ class Resource:
     async def set_basic(
         self,
         payload: PBBasic,
-        deleted_fields: Optional[list[FieldID]] = None,
+        deleted_fields: list[FieldID] | None = None,
     ):
         await self.get_basic()
 
@@ -212,49 +226,43 @@ class Resource:
         if deleted_fields is not None and len(deleted_fields) > 0:
             delete_basic_computedmetadata_classifications(self.basic, deleted_fields=deleted_fields)
 
-        await datamanagers.resources.set_basic(
-            self.txn, kbid=self.kb.kbid, rid=self.uuid, basic=self.basic
-        )
+        await datamanagers.resources.set_basic(self.txn, kbid=self.kbid, rid=self.uuid, basic=self.basic)
         self.modified = True
 
     # Origin
-    async def get_origin(self) -> Optional[PBOrigin]:
+    async def get_origin(self) -> PBOrigin | None:
         if self.origin is None:
-            origin = await datamanagers.resources.get_origin(self.txn, kbid=self.kb.kbid, rid=self.uuid)
+            origin = await datamanagers.resources.get_origin(self.txn, kbid=self.kbid, rid=self.uuid)
             self.origin = origin
         return self.origin
 
     async def set_origin(self, payload: PBOrigin):
-        await datamanagers.resources.set_origin(
-            self.txn, kbid=self.kb.kbid, rid=self.uuid, origin=payload
-        )
+        await datamanagers.resources.set_origin(self.txn, kbid=self.kbid, rid=self.uuid, origin=payload)
         self.modified = True
         self.origin = payload
 
     # Extra
-    async def get_extra(self) -> Optional[PBExtra]:
+    async def get_extra(self) -> PBExtra | None:
         if self.extra is None:
-            extra = await datamanagers.resources.get_extra(self.txn, kbid=self.kb.kbid, rid=self.uuid)
+            extra = await datamanagers.resources.get_extra(self.txn, kbid=self.kbid, rid=self.uuid)
             self.extra = extra
         return self.extra
 
     async def set_extra(self, payload: PBExtra):
-        await datamanagers.resources.set_extra(self.txn, kbid=self.kb.kbid, rid=self.uuid, extra=payload)
+        await datamanagers.resources.set_extra(self.txn, kbid=self.kbid, rid=self.uuid, extra=payload)
         self.modified = True
         self.extra = payload
 
     # Security
-    async def get_security(self) -> Optional[utils_pb2.Security]:
+    async def get_security(self) -> utils_pb2.Security | None:
         if self.security is None:
-            security = await datamanagers.resources.get_security(
-                self.txn, kbid=self.kb.kbid, rid=self.uuid
-            )
+            security = await datamanagers.resources.get_security(self.txn, kbid=self.kbid, rid=self.uuid)
             self.security = security
         return self.security
 
     async def set_security(self, payload: utils_pb2.Security) -> None:
         await datamanagers.resources.set_security(
-            self.txn, kbid=self.kb.kbid, rid=self.uuid, security=payload
+            self.txn, kbid=self.kbid, rid=self.uuid, security=payload
         )
         self.modified = True
         self.security = payload
@@ -262,7 +270,7 @@ class Resource:
     # Relations
     async def get_user_relations(self) -> PBRelations:
         if self.user_relations is None:
-            sf = self.storage.user_relations(self.kb.kbid, self.uuid)
+            sf = self.storage.user_relations(self.kbid, self.uuid)
             relations = await self.storage.download_pb(sf, PBRelations)
             if relations is None:
                 # Key not found = no relations
@@ -272,7 +280,7 @@ class Resource:
         return self.user_relations
 
     async def set_user_relations(self, payload: PBRelations):
-        sf = self.storage.user_relations(self.kb.kbid, self.uuid)
+        sf = self.storage.user_relations(self.kbid, self.uuid)
         await self.storage.upload_pb(sf, payload)
         self.modified = True
         self.user_relations = payload
@@ -354,25 +362,34 @@ class Resource:
 
         await field_obj.delete()
 
+    async def field_exists(self, type: FieldType.ValueType, field: str) -> bool:
+        """Return whether this resource has this field or not."""
+        all_fields_ids = await self.get_fields_ids()
+        for field_type, field_id in all_fields_ids:
+            if field_type == type and field_id == field:
+                return True
+        return False
+
     def has_field(self, type: FieldType.ValueType, field: str) -> bool:
+        # REVIEW: are we sure we don't want to actually check this?
         return (type, field) in self.fields
 
-    async def get_all_field_ids(self, *, for_update: bool) -> Optional[PBAllFieldIDs]:
+    async def get_all_field_ids(self, *, for_update: bool) -> PBAllFieldIDs | None:
         return await datamanagers.resources.get_all_field_ids(
-            self.txn, kbid=self.kb.kbid, rid=self.uuid, for_update=for_update
+            self.txn, kbid=self.kbid, rid=self.uuid, for_update=for_update
         )
 
     async def set_all_field_ids(self, all_fields: PBAllFieldIDs):
         return await datamanagers.resources.set_all_field_ids(
-            self.txn, kbid=self.kb.kbid, rid=self.uuid, allfields=all_fields
+            self.txn, kbid=self.kbid, rid=self.uuid, allfields=all_fields
         )
 
     async def update_all_field_ids(
         self,
         *,
-        updated: Optional[list[FieldID]] = None,
-        deleted: Optional[list[FieldID]] = None,
-        errors: Optional[list[writer_pb2.Error]] = None,
+        updated: list[FieldID] | None = None,
+        deleted: list[FieldID] | None = None,
+        errors: list[writer_pb2.Error] | None = None,
     ):
         needs_update = False
         all_fields = await self.get_all_field_ids(for_update=True)
@@ -451,7 +468,7 @@ class Resource:
 
         # If this message comes from the processor (not a DA worker), we clear all previous errors
         # TODO: When generated_by is populated with DA tasks by processor, remove only related errors
-        from_processor = any((x.WhichOneof("generator") == "processor" for x in message.generated_by))
+        from_processor = any(x.WhichOneof("generator") == "processor" for x in message.generated_by)
 
         for (field_type, field), errors in errors_by_field.items():
             field_obj = await self.get_field(field, field_type, load=False)
@@ -471,7 +488,7 @@ class Resource:
             # We infer the status for processor messages
             if message.source == BrokerMessage.MessageSource.PROCESSOR:
                 if any(
-                    (e.source_error.severity == writer_pb2.Error.Severity.ERROR for e in status.errors)
+                    e.source_error.severity == writer_pb2.Error.Severity.ERROR for e in status.errors
                 ):
                     status.status = writer_pb2.FieldStatus.Status.ERROR
                 else:
@@ -501,25 +518,21 @@ class Resource:
             return
 
         field_statuses = await datamanagers.fields.get_statuses(
-            self.txn, kbid=self.kb.kbid, rid=self.uuid, fields=field_ids.fields
+            self.txn, kbid=self.kbid, rid=self.uuid, fields=field_ids.fields
         )
 
         # If any field is processing -> PENDING
-        if any((f.status == writer_pb2.FieldStatus.Status.PENDING for f in field_statuses)):
+        if any(f.status == writer_pb2.FieldStatus.Status.PENDING for f in field_statuses):
             self.basic.metadata.status = PBMetadata.Status.PENDING
         # If we have any non-DA error -> ERROR
         elif any(
-            (
-                f.status == writer_pb2.FieldStatus.Status.ERROR
-                and any(
-                    (
-                        e.source_error.severity == writer_pb2.Error.Severity.ERROR
-                        and e.source_error.code != writer_pb2.Error.ErrorCode.DATAAUGMENTATION
-                        for e in f.errors
-                    )
-                )
-                for f in field_statuses
+            f.status == writer_pb2.FieldStatus.Status.ERROR
+            and any(
+                e.source_error.severity == writer_pb2.Error.Severity.ERROR
+                and e.source_error.code != writer_pb2.Error.ErrorCode.DATAAUGMENTATION
+                for e in f.errors
             )
+            for f in field_statuses
         ):
             self.basic.metadata.status = PBMetadata.Status.ERROR
         # Otherwise (everything processed or we only have DA errors) -> PROCESSED
@@ -642,7 +655,7 @@ class Resource:
             FieldType.LINK,
             load=False,
         )
-        maybe_update_basic_thumbnail(self.basic, link_extracted_data.link_thumbnail, self.kb.kbid)
+        maybe_update_basic_thumbnail(self.basic, link_extracted_data.link_thumbnail, self.kbid)
 
         await field_link.set_link_extracted_data(link_extracted_data)
 
@@ -669,7 +682,7 @@ class Resource:
             return
         logger.info(
             "Updating resource title from link extracted data",
-            extra={"kbid": self.kb.kbid, "field": link_extracted_data.field, "rid": self.uuid},
+            extra={"kbid": self.kbid, "field": link_extracted_data.field, "rid": self.uuid},
         )
         title = link_extracted_data.title
         await self.update_resource_title(title)
@@ -711,7 +724,7 @@ class Resource:
         # uri can change after extraction
         await field_file.set_file_extracted_data(file_extracted_data)
         maybe_update_basic_icon(self.basic, file_extracted_data.icon)
-        maybe_update_basic_thumbnail(self.basic, file_extracted_data.file_thumbnail, self.kb.kbid)
+        maybe_update_basic_thumbnail(self.basic, file_extracted_data.file_thumbnail, self.kbid)
         self.modified = True
 
     async def _should_update_resource_title_from_file_metadata(self) -> bool:
@@ -733,7 +746,7 @@ class Resource:
         filenames = set()
         for (field_type, _), field_obj in fields.items():
             if field_type == FieldType.FILE:
-                field_value: Optional[FieldFile] = await field_obj.get_value()
+                field_value: FieldFile | None = await field_obj.get_value()
                 if field_value is not None:
                     if field_value.file.filename not in ("", None):
                         filenames.add(field_value.file.filename)
@@ -758,7 +771,7 @@ class Resource:
             fid = FieldId.from_pb(rid=self.uuid, field_type=FieldType.FILE, key=fed.field)
             logger.info(
                 "Updating resource title from file extracted data",
-                extra={"kbid": self.kb.kbid, "field": fid.full(), "new_title": fed.title},
+                extra={"kbid": self.kbid, "field": fid.full(), "new_title": fed.title},
             )
             await self.update_resource_title(fed.title)
             await self.unmark_title_for_reset()
@@ -776,9 +789,7 @@ class Resource:
         )
         await field_obj.set_field_metadata(field_metadata)
 
-        maybe_update_basic_thumbnail(
-            self.basic, field_metadata.metadata.metadata.thumbnail, self.kb.kbid
-        )
+        maybe_update_basic_thumbnail(self.basic, field_metadata.metadata.metadata.thumbnail, self.kbid)
 
         update_basic_computedmetadata_classifications(self.basic, field_metadata)
         self.modified = True
@@ -790,7 +801,7 @@ class Resource:
         await self.get_fields(force=True)
         vectorsets = {
             vectorset_id: vs
-            async for vectorset_id, vs in datamanagers.vectorsets.iter(self.txn, kbid=self.kb.kbid)
+            async for vectorset_id, vs in datamanagers.vectorsets.iter(self.txn, kbid=self.kbid)
         }
 
         for field_vectors in fields_vectors:
@@ -799,13 +810,13 @@ class Resource:
                 assert len(vectorsets) == 1, (
                     "Invalid broker message, can't ingest vectors from unknown vectorset to KB with multiple vectorsets"
                 )
-                vectorset = list(vectorsets.values())[0]
+                vectorset = next(iter(vectorsets.values()))
 
             else:
                 if field_vectors.vectorset_id not in vectorsets:
                     logger.warning(
                         "Dropping extracted vectors for unknown vectorset",
-                        extra={"kbid": self.kb.kbid, "vectorset": field_vectors.vectorset_id},
+                        extra={"kbid": self.kbid, "vectorset": field_vectors.vectorset_id},
                     )
                     continue
 
@@ -916,7 +927,7 @@ def maybe_update_basic_summary(basic: PBBasic, summary_text: str) -> bool:
     return True
 
 
-def maybe_update_basic_icon(basic: PBBasic, mimetype: Optional[str]) -> bool:
+def maybe_update_basic_icon(basic: PBBasic, mimetype: str | None) -> bool:
     if basic.icon not in (None, "", "application/octet-stream", GENERIC_MIME_TYPE):
         # Icon already set or detected
         return False
@@ -935,7 +946,7 @@ def maybe_update_basic_icon(basic: PBBasic, mimetype: Optional[str]) -> bool:
     return True
 
 
-def maybe_update_basic_thumbnail(basic: PBBasic, thumbnail: Optional[CloudFile], kbid: str) -> bool:
+def maybe_update_basic_thumbnail(basic: PBBasic, thumbnail: CloudFile | None, kbid: str) -> bool:
     if basic.thumbnail or thumbnail is None:
         return False
     basic.thumbnail = CloudLink.format_reader_download_uri(thumbnail.uri)
@@ -972,7 +983,7 @@ def update_basic_languages(basic: Basic, languages: list[str]) -> bool:
     return updated
 
 
-def get_text_field_mimetype(bm: BrokerMessage) -> Optional[str]:
+def get_text_field_mimetype(bm: BrokerMessage) -> str | None:
     if len(bm.texts) == 0:
         return None
     text_format = next(iter(bm.texts.values())).format

nucliadb/ingest/orm/utils.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import urllib.parse
-from typing import Sequence
+from collections.abc import Sequence
 
 from nucliadb.models.internal.processing import PushPayload, PushTextFormat, Text
 from nucliadb_protos.resources_pb2 import (

nucliadb/ingest/partitions.py

@@ -25,12 +25,17 @@ from nucliadb.ingest.settings import Settings
 
 
 def assign_partitions(settings: Settings):
+    """
+    This function dynamically assigns the partitions to the current ingest sts
+    replica based on its hostname, typically (ingest-0, ingest-1, etc).
+    """
     # partitions start from 1, instead of 0
     all_partitions = [str(part + 1) for part in range(settings.nuclia_partitions)]
 
     # get replica number and total replicas from environment
     logger.info(f"PARTITIONS: Total Replicas = {settings.total_replicas}")
     if settings.replica_number == -1:
+        # Get replica number from hostname
         hostname = os.environ.get("HOSTNAME")
         if hostname is not None:
             sts_values = hostname.split("-")
@@ -39,10 +44,16 @@ def assign_partitions(settings: Settings):
                     settings.replica_number = int(sts_values[-1])
                 except Exception:
                     logger.error(f"Could not extract replica number from hostname: {hostname}")
-                    pass
+            else:
+                logger.warning(f"Could not determine replica number from hostname: {hostname}")
+        else:
+            logger.warning(f"Could not determine replica number from hostname.")
 
         if settings.replica_number == -1:
             settings.replica_number = 0
+    else:
+        # We assume that replica numbers are set manually via env variables
+        pass
     logger.info(f"PARTITIONS: Replica Number = {settings.replica_number}")
 
     # calculate assigned partitions based on total replicas and own replica number

nucliadb/ingest/processing.py

@@ -25,7 +25,7 @@ import uuid
 from collections import defaultdict
 from contextlib import AsyncExitStack
 from enum import Enum
-from typing import Any, Optional
+from typing import Any
 
 import aiohttp
 import backoff
@@ -132,19 +132,19 @@ def to_processing_driver_type(file_backend_driver: FileBackendConfig) -> Process
 class ProcessingEngine:
     def __init__(
         self,
-        nuclia_service_account: Optional[str] = None,
-        nuclia_zone: Optional[str] = None,
-        nuclia_public_url: Optional[str] = None,
-        nuclia_processing_cluster_url: Optional[str] = None,
-        onprem: Optional[bool] = False,
-        nuclia_jwt_key: Optional[str] = None,
+        nuclia_service_account: str | None = None,
+        nuclia_zone: str | None = None,
+        nuclia_public_url: str | None = None,
+        nuclia_processing_cluster_url: str | None = None,
+        onprem: bool | None = False,
+        nuclia_jwt_key: str | None = None,
         days_to_keep: int = 3,
         driver: FileBackendConfig = FileBackendConfig.GCS,
     ):
         self.nuclia_service_account = nuclia_service_account
         self.nuclia_zone = nuclia_zone
         if nuclia_public_url is not None:
-            self.nuclia_public_url: Optional[str] = nuclia_public_url.format(zone=nuclia_zone)
+            self.nuclia_public_url: str | None = nuclia_public_url.format(zone=nuclia_zone)
         else:
             self.nuclia_public_url = None
 
@@ -196,7 +196,7 @@ class ProcessingEngine:
         return jwt.encode(payload, self.nuclia_jwt_key, algorithm="HS256")
 
     def generate_file_token_from_fieldfile(
-        self, file: FieldFilePB, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file: FieldFilePB, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         if self.nuclia_jwt_key is None:
             raise AttributeError("Nuclia JWT key not set")
@@ -235,7 +235,7 @@ class ProcessingEngine:
     )
     @processing_observer.wrap({"type": "file_field_upload"})
     async def convert_filefield_to_str(
-        self, file: models.FileField, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file: models.FileField, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         # Upload file without storing on Nuclia DB
         headers = {}
@@ -273,7 +273,7 @@ class ProcessingEngine:
         ).decode()
 
     def convert_external_filefield_to_str(
-        self, file_field: models.FileField, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file_field: models.FileField, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         if self.nuclia_jwt_key is None:
             raise AttributeError("Nuclia JWT key not set")
@@ -313,7 +313,7 @@ class ProcessingEngine:
         self,
         file: FieldFilePB,
         storage: Storage,
-        classif_labels: Optional[list[ClassificationLabel]] = None,
+        classif_labels: list[ClassificationLabel] | None = None,
     ) -> str:
         """It's already an internal file that needs to be uploaded"""
         if self.onprem is False:
@@ -438,7 +438,7 @@ class ProcessingEngine:
             queue=QueueType(queue_type) if queue_type is not None else None,
         )
 
-    async def delete_from_processing(self, *, kbid: str, resource_id: Optional[str] = None) -> None:
+    async def delete_from_processing(self, *, kbid: str, resource_id: str | None = None) -> None:
         """
         Delete a resource from processing. This prevents inflight resources from being processed
         and wasting resources.
@@ -479,7 +479,7 @@ class DummyProcessingEngine(ProcessingEngine):
         pass
 
     async def convert_filefield_to_str(
-        self, file: models.FileField, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file: models.FileField, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         self.calls.append([file])
         index = len(self.values["convert_filefield_to_str"])
@@ -487,7 +487,7 @@ class DummyProcessingEngine(ProcessingEngine):
         return f"convert_filefield_to_str,{index}"
 
     def convert_external_filefield_to_str(
-        self, file_field: models.FileField, classif_labels: Optional[list[ClassificationLabel]] = None
+        self, file_field: models.FileField, classif_labels: list[ClassificationLabel] | None = None
     ) -> str:
         self.calls.append([file_field])
         index = len(self.values["convert_external_filefield_to_str"])
@@ -498,7 +498,7 @@ class DummyProcessingEngine(ProcessingEngine):
         self,
         file: FieldFilePB,
         storage: Storage,
-        classif_labels: Optional[list[ClassificationLabel]] = None,
+        classif_labels: list[ClassificationLabel] | None = None,
     ) -> str:
         self.calls.append([file, storage])
         index = len(self.values["convert_internal_filefield_to_str"])
@@ -516,5 +516,5 @@ class DummyProcessingEngine(ProcessingEngine):
         self.values["send_to_process"].append([item, partition])
         return ProcessingInfo(seqid=len(self.calls), account_seq=0, queue=QueueType.SHARED)
 
-    async def delete_from_processing(self, *, kbid: str, resource_id: Optional[str] = None) -> None:
+    async def delete_from_processing(self, *, kbid: str, resource_id: str | None = None) -> None:
         self.calls.append([kbid, resource_id])