nucliadb 6.3.5.post3985__py3-none-any.whl → 6.3.5.post3990__py3-none-any.whl

migrations/0023_backfill_pg_catalog.py

@@ -29,6 +29,7 @@ from typing import cast
 
 from nucliadb.common import datamanagers
 from nucliadb.common.maindb.pg import PGDriver, PGTransaction
+from nucliadb.ingest.orm.index_message import get_resource_index_message
 from nucliadb.ingest.orm.processor.pgcatalog import pgcatalog_update
 from nucliadb.migrator.context import ExecutionContext
 
@@ -73,8 +74,8 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
                         logger.warning(f"Could not load resource {rid} for kbid {kbid}")
                         continue
 
-                    await resource.compute_global_tags(resource.indexer)
-                    await pgcatalog_update(txn, kbid, resource)
+                    index_message = await get_resource_index_message(resource, reindex=False)
+                    await pgcatalog_update(txn, kbid, resource, index_message)
 
                 await txn.commit()
                 continue_sql = f"AND key > '/kbs/{kbid}/r/{rid}'"

nucliadb/common/cluster/rollover.py

@@ -38,7 +38,7 @@ from nucliadb_telemetry import errors
 from .utils import (
     delete_resource_from_shard,
     get_resource,
-    get_resource_index_message,
+    get_rollover_resource_index_message,
     index_resource_to_shard,
 )
 
@@ -288,7 +288,7 @@ async def _index_resource_to_rollover_index(
                 f"Shard {shard_id} not found. Was a new one created during migration?"
             )
         resource = await get_resource(kbid, resource_id)
-        index_message = await get_resource_index_message(kbid, resource_id)
+        index_message = await get_rollover_resource_index_message(kbid, resource_id)
         if resource is None or index_message is None:
             # resource no longer existing, remove indexing and carry on
             async with datamanagers.with_transaction() as txn:
@@ -503,7 +503,7 @@ async def validate_indexed_data(
                 await txn.commit()
             continue
 
-        index_message = await get_resource_index_message(kbid, resource_id)
+        index_message = await get_rollover_resource_index_message(kbid, resource_id)
         if index_message is None:
             logger.error(
                 "Resource index message not found while validating, skipping",

nucliadb/common/cluster/utils.py

@@ -28,6 +28,7 @@ from nucliadb.common.cluster.manager import (
     StandaloneKBShardManager,
 )
 from nucliadb.common.cluster.settings import settings
+from nucliadb.ingest.orm import index_message
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb_protos import nodereader_pb2, writer_pb2
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
@@ -80,7 +81,9 @@ async def get_resource(kbid: str, resource_id: str) -> Optional[Resource]:
 
 
 @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=8)
-async def get_resource_index_message(kbid: str, resource_id: str) -> Optional[nodereader_pb2.Resource]:
+async def get_rollover_resource_index_message(
+    kbid: str, resource_id: str
+) -> Optional[nodereader_pb2.Resource]:
     async with datamanagers.with_ro_transaction() as txn:
         resource = await datamanagers.resources.get_resource(txn, kbid=kbid, rid=resource_id)
         if resource is None:
@@ -89,8 +92,9 @@ async def get_resource_index_message(kbid: str, resource_id: str) -> Optional[no
                 extra={"kbid": kbid, "resource_id": resource_id},
             )
             return None
-        resource_index_message = (await resource.generate_index_message(reindex=False)).brain
-        return resource_index_message
+        # We set the reindex=False because we are indexing the resource for the first time in the
+        # newly created shards.
+        return await index_message.get_resource_index_message(resource, reindex=False)
 
 
 @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=8)
@@ -106,7 +110,7 @@ async def index_resource_to_shard(
     partitioning = app_context.partitioning
 
     if resource_index_message is None:
-        resource_index_message = await get_resource_index_message(kbid, resource_id)
+        resource_index_message = await get_rollover_resource_index_message(kbid, resource_id)
         if resource_index_message is None:
             return
 

nucliadb/common/external_index_providers/pinecone.py

@@ -36,7 +36,7 @@ from nucliadb.common.external_index_providers.base import (
     VectorsetExternalIndex,
 )
 from nucliadb.common.external_index_providers.exceptions import ExternalIndexCreationError
-from nucliadb.common.ids import FieldId, ParagraphId, VectorId
+from nucliadb.common.ids import ParagraphId, VectorId
 from nucliadb_models.search import SCORE_TYPE, TextPosition
 from nucliadb_protos import knowledgebox_pb2 as kb_pb2
 from nucliadb_protos import utils_pb2
@@ -418,52 +418,17 @@ class PineconeIndexManager(ExternalIndexManager):
         if len(delete_tasks) > 0:
             await asyncio.gather(*delete_tasks)
 
-    def get_vectorsets_in_resource(self, index_data: Resource) -> set[str]:
-        vectorsets: set[str] = set()
-        for _, paragraph in iter_paragraphs(index_data):
-            if not paragraph.sentences and not paragraph.vectorsets_sentences:
-                continue
-            if paragraph.sentences and self.default_vectorset:
-                vectorsets.add(self.default_vectorset)
-            for vectorset_id, vectorsets_sentences in paragraph.vectorsets_sentences.items():
-                if vectorsets_sentences.sentences:
-                    vectorsets.add(vectorset_id)
-            # Once we have found at least one paragraph with vectors, we can stop iterating
-            return vectorsets
-        return vectorsets
-
     def get_index_host(self, vectorset_id: str, rollover: bool = False) -> str:
         if rollover:
             return self.rollover_indexes[vectorset_id].index_host
         else:
             return self.indexes[vectorset_id].index_host
 
-    def get_prefixes_to_delete(self, index_data: Resource) -> set[str]:
-        prefixes_to_delete = set()
-        # TODO: migrate to vector_prefixes_to_delete
-        for field_id in index_data.sentences_to_delete:
-            try:
-                delete_vid = VectorId.from_string(field_id)
-                prefixes_to_delete.add(delete_vid.field_id.full())
-            except ValueError:  # pragma: no cover
-                try:
-                    delete_field = FieldId.from_string(field_id)
-                    prefixes_to_delete.add(delete_field.full())
-                except ValueError:
-                    logger.warning(f"Invalid id to delete sentences from: {field_id}.")
-                    continue
-        for paragraph_id in index_data.paragraphs_to_delete:
-            try:
-                delete_pid = ParagraphId.from_string(paragraph_id)
-                prefixes_to_delete.add(delete_pid.field_id.full())
-            except ValueError:  # pragma: no cover
-                try:
-                    delete_field = FieldId.from_string(paragraph_id)
-                    prefixes_to_delete.add(delete_field.full())
-                except ValueError:
-                    logger.warning(f"Invalid id to delete: {paragraph_id}. ParagraphId expected.")
-                    continue
-        return prefixes_to_delete
+    def get_prefixes_to_delete(self, index_data: Resource) -> dict[str, set[str]]:
+        return {
+            vectorset_id: set(prefixes_list.items)
+            for vectorset_id, prefixes_list in index_data.vector_prefixes_to_delete.items()
+        }
 
     async def _index_resource(
         self, resource_uuid: str, index_data: Resource, to_rollover_indexes: bool = False
@@ -480,10 +445,8 @@ class PineconeIndexManager(ExternalIndexManager):
           metadata with any specific sentence metadata. This is done for each vectorset.
         - Finally, upsert the vectors to each vectorset index in parallel.
         """
-        vectorsets = self.get_vectorsets_in_resource(index_data)
-        prefixes_to_delete = self.get_prefixes_to_delete(index_data)
         delete_tasks = []
-        for vectorset in vectorsets:
+        for vectorset, prefixes_to_delete in self.get_prefixes_to_delete(index_data).items():
             index_host = self.get_index_host(vectorset_id=vectorset, rollover=to_rollover_indexes)
             delete_tasks.append(
                 asyncio.create_task(

nucliadb/ingest/fields/exceptions.py

@@ -29,3 +29,7 @@ class InvalidPBClass(Exception):
         self.source = source
         self.destination = destination
         super().__init__(f"Source and destination does not match {self.source} - {self.destination}")
+
+
+class FieldAuthorNotFound(Exception):
+    pass