nucliadb 6.7.2.post4862__py3-none-any.whl → 6.9.2.post5282__py3-none-any.whl

nucliadb/ingest/orm/processor/processor.py

@@ -28,6 +28,7 @@ from nidx_protos import noderesources_pb2, nodewriter_pb2
 from nidx_protos.noderesources_pb2 import Resource as PBBrainResource
 
 from nucliadb.common import datamanagers, locking
+from nucliadb.common.catalog import catalog_delete, catalog_update
 from nucliadb.common.cluster.settings import settings as cluster_settings
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.external_index_providers.base import ExternalIndexManager
@@ -61,8 +62,6 @@ from nucliadb_utils.cache.pubsub import PubSubDriver
 from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.utilities import get_storage, has_feature
 
-from .pgcatalog import pgcatalog_delete, pgcatalog_update
-
 logger = logging.getLogger("ingest-processor")
 
 MESSAGE_TO_NOTIFICATION_SOURCE = {
@@ -142,8 +141,6 @@ class Processor:
     and can not use the txn id
     """
 
-    messages: dict[str, list[writer_pb2.BrokerMessage]]
-
     def __init__(
         self,
         driver: Driver,
@@ -151,7 +148,6 @@ class Processor:
         pubsub: Optional[PubSubDriver] = None,
         partition: Optional[str] = None,
     ):
-        self.messages = {}
         self.driver = driver
         self.storage = storage
         self.partition = partition
@@ -180,18 +176,12 @@ class Processor:
         if message.type == writer_pb2.BrokerMessage.MessageType.DELETE:
             await self.delete_resource(message, seqid, partition, transaction_check)
         elif message.type == writer_pb2.BrokerMessage.MessageType.AUTOCOMMIT:
-            await self.txn([message], seqid, partition, transaction_check)
-        elif message.type == writer_pb2.BrokerMessage.MessageType.MULTI:
-            # XXX Not supported right now
-            # MULTI, COMMIT and ROLLBACK are all not supported in transactional mode right now
-            # This concept is probably not tenable with current architecture because
-            # of how nats works and how we would need to manage rollbacks.
-            # XXX Should this be removed?
-            await self.multi(message, seqid)
-        elif message.type == writer_pb2.BrokerMessage.MessageType.COMMIT:
-            await self.commit(message, seqid, partition)
-        elif message.type == writer_pb2.BrokerMessage.MessageType.ROLLBACK:
-            await self.rollback(message, seqid, partition)
+            await self.txn(message, seqid, partition, transaction_check)
+        else:  # pragma: no cover
+            logger.error(
+                f"Unsupported message type: {message.type}",
+                extra={"seqid": seqid, "partition": partition},
+            )
 
     async def get_resource_uuid(self, kb: KnowledgeBox, message: writer_pb2.BrokerMessage) -> str:
         if message.uuid is None:
@@ -208,7 +198,7 @@ class Processor:
         partition: str,
         transaction_check: bool = True,
     ) -> None:
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             try:
                 kb = KnowledgeBox(txn, self.storage, message.kbid)
 
@@ -227,7 +217,8 @@ class Processor:
                     shard = await kb.get_resource_shard(shard_id)
                     if shard is None:
                         raise AttributeError("Shard not available")
-                    await pgcatalog_delete(txn, message.kbid, uuid)
+
+                    await catalog_delete(txn, message.kbid, uuid)
                     external_index_manager = await get_external_index_manager(kbid=message.kbid)
                     if external_index_manager is not None:
                         await self.external_index_delete_resource(external_index_manager, uuid)
@@ -242,7 +233,6 @@ class Processor:
                         await self.notify_abort(
                             partition=partition,
                             seqid=seqid,
-                            multi=message.multiid,
                             kbid=message.kbid,
                             rid=message.uuid,
                             source=message.source,
@@ -256,7 +246,6 @@ class Processor:
         await self.notify_commit(
             partition=partition,
             seqid=seqid,
-            multi=message.multiid,
             message=message,
             write_type=writer_pb2.Notification.WriteType.DELETED,
         )
@@ -267,7 +256,7 @@ class Processor:
         # so we commit it in a different transaction to make it as short as possible
         prev_txn = resource.txn
         try:
-            async with self.driver.transaction() as txn:
+            async with self.driver.rw_transaction() as txn:
                 resource.txn = txn
                 await resource.set_slug()
                 await txn.commit()
@@ -277,15 +266,12 @@ class Processor:
     @processor_observer.wrap({"type": "txn"})
     async def txn(
         self,
-        messages: list[writer_pb2.BrokerMessage],
+        message: writer_pb2.BrokerMessage,
         seqid: int,
         partition: str,
         transaction_check: bool = True,
     ) -> None:
-        if len(messages) == 0:
-            return None
-
-        kbid = messages[0].kbid
+        kbid = message.kbid
         if not await datamanagers.atomic.kb.exists_kb(kbid=kbid):
             logger.info(f"KB {kbid} is deleted: skiping txn")
             if transaction_check:
@@ -294,60 +280,57 @@ class Processor:
                     await txn.commit()
             return None
 
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             try:
-                multi = messages[0].multiid
                 kb = KnowledgeBox(txn, self.storage, kbid)
-                uuid = await self.get_resource_uuid(kb, messages[0])
+                uuid = await self.get_resource_uuid(kb, message)
+
                 resource: Optional[Resource] = None
                 handled_exception = None
                 created = False
 
-                for message in messages:
-                    if resource is not None:
-                        assert resource.uuid == message.uuid
-
-                    if message.source == writer_pb2.BrokerMessage.MessageSource.WRITER:
-                        resource = await kb.get(uuid)
-                        if resource is None:
-                            # It's a new resource
-                            resource = await kb.add_resource(uuid, message.slug, message.basic)
-                            created = True
-                        else:
-                            # It's an update from writer for an existing resource
-                            ...
-
-                    elif message.source == writer_pb2.BrokerMessage.MessageSource.PROCESSOR:
-                        resource = await kb.get(uuid)
-                        if resource is None:
-                            logger.info(
-                                f"Secondary message for resource {message.uuid} and resource does not exist, ignoring"
-                            )
-                            continue
-                        else:
-                            # It's an update from processor for an existing resource
-                            ...
-
-                        generated_fields = await get_generated_fields(message, resource)
-                        if generated_fields.is_not_empty():
-                            await send_generated_fields_to_process(
-                                kbid, resource, generated_fields, message
-                            )
-                            # TODO: remove this when processor sends the field set
-                            for generated_text in generated_fields.texts:
-                                message.texts[
-                                    generated_text
-                                ].generated_by.data_augmentation.SetInParent()
-
+                if message.source == writer_pb2.BrokerMessage.MessageSource.WRITER:
+                    resource = await kb.get(uuid)
+                    if resource is None:
+                        # It's a new resource
+                        resource = await kb.add_resource(uuid, message.slug, message.basic)
+                        created = True
                     else:
-                        raise InvalidBrokerMessage(f"Unknown broker message source: {message.source}")
+                        # It's an update from writer for an existing resource
+                        ...
+
+                elif message.source == writer_pb2.BrokerMessage.MessageSource.PROCESSOR:
+                    resource = await kb.get(uuid)
+                    if resource is None:
+                        logger.info(
+                            f"Processor message for resource received but the resource does not exist, ignoring.",
+                            extra={
+                                "kbid": kbid,
+                                "rid": uuid,
+                                "seqid": seqid,
+                            },
+                        )
+                        return None
+                    else:
+                        # It's an update from processor for an existing resource
+                        ...
+
+                    generated_fields = await get_generated_fields(message, resource)
+                    if generated_fields.is_not_empty():
+                        await send_generated_fields_to_process(kbid, resource, generated_fields, message)
+                        # TODO: remove this when processor sends the field set
+                        for generated_text in generated_fields.texts:
+                            message.texts[generated_text].generated_by.data_augmentation.SetInParent()
+
+                else:  # pragma: no cover
+                    raise InvalidBrokerMessage(f"Unknown broker message source: {message.source}")
 
-                    # apply changes from the broker message to the resource
-                    await self.apply_resource(message, resource, update=(not created))
+                # apply changes from the broker message to the resource
+                await self.apply_resource(message, resource, update=(not created))
 
                 # index message
                 if resource and resource.modified:
-                    index_message = await self.generate_index_message(resource, messages, created)
+                    index_message = await self.generate_index_message(resource, message, created)
                     try:
                         warnings = await self.index_resource(
                             index_message=index_message,
@@ -357,7 +340,7 @@ class Processor:
                             seqid=seqid,
                             partition=partition,
                             kb=kb,
-                            source=messages_source(messages),
+                            source=to_index_message_source(message),
                         )
                         # Save indexing warnings
                         for field_id, warning in warnings:
@@ -374,8 +357,7 @@ class Processor:
                             index_message.labels.remove(current_status[0])
                             index_message.labels.append("/n/s/ERROR")
 
-                    await pgcatalog_update(txn, kbid, resource, index_message)
-
+                    await catalog_update(txn, kbid, resource, index_message)
                     if transaction_check:
                         await sequence_manager.set_last_seqid(txn, partition, seqid)
                     await txn.commit()
@@ -386,7 +368,6 @@ class Processor:
                     await self.notify_commit(
                         partition=partition,
                         seqid=seqid,
-                        multi=multi,
                         message=message,
                         write_type=(
                             writer_pb2.Notification.WriteType.CREATED
@@ -399,7 +380,6 @@ class Processor:
                     await self.notify_abort(
                         partition=partition,
                         seqid=seqid,
-                        multi=multi,
                         kbid=kbid,
                         rid=uuid,
                         source=message.source,
@@ -419,7 +399,6 @@ class Processor:
                 await self.notify_abort(
                     partition=partition,
                     seqid=seqid,
-                    multi=multi,
                     kbid=kbid,
                     rid=uuid,
                     source=message.source,
@@ -429,11 +408,10 @@ class Processor:
                 # As we are in the middle of a transaction, we cannot let the exception raise directly
                 # as we need to do some cleanup. The exception will be reraised at the end of the function
                 # and then handled by the top caller, so errors can be handled in the same place.
-                await self.deadletter(messages, partition, seqid)
+                await self.deadletter(message, partition, seqid)
                 await self.notify_abort(
                     partition=partition,
                     seqid=seqid,
-                    multi=multi,
                     kbid=kbid,
                     rid=uuid,
                     source=message.source,
@@ -468,22 +446,27 @@ class Processor:
             # a resource was move to another shard while it was being indexed
             shard_id = await datamanagers.resources.get_resource_shard_id(txn, kbid=kbid, rid=uuid)
 
-        shard = None
-        if shard_id is not None:
-            # Resource already has a shard assigned
-            shard = await kb.get_resource_shard(shard_id)
-            if shard is None:
-                raise AttributeError("Shard not available")
-        else:
-            # It's a new resource, get KB's current active shard to place new resource on
-            shard = await self.index_node_shard_manager.get_current_active_shard(txn, kbid)
-            if shard is None:
-                # No current shard available, create a new one
-                shard = await self.index_node_shard_manager.create_shard_by_kbid(txn, kbid)
-            await datamanagers.resources.set_resource_shard_id(
-                txn, kbid=kbid, rid=uuid, shard=shard.shard
-            )
-        return shard
+            shard = None
+            if shard_id is not None:
+                # Resource already has a shard assigned
+                shard = await kb.get_resource_shard(shard_id)
+                if shard is None:
+                    raise AttributeError("Shard not available")
+            else:
+                # It's a new resource, get KB's current active shard to place new resource on
+                shard = await self.index_node_shard_manager.get_current_active_shard(txn, kbid)
+                if shard is None:
+                    # No current shard available, create a new one
+                    async with locking.distributed_lock(locking.NEW_SHARD_LOCK.format(kbid=kbid)):
+                        kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
+                        prewarm = kb_config is not None and kb_config.prewarm_enabled
+                        shard = await self.index_node_shard_manager.create_shard_by_kbid(
+                            txn, kbid, prewarm_enabled=prewarm
+                        )
+                await datamanagers.resources.set_resource_shard_id(
+                    txn, kbid=kbid, rid=uuid, shard=shard.shard
+                )
+            return shard
 
     @processor_observer.wrap({"type": "index_resource"})
     async def index_resource(
@@ -519,17 +502,16 @@ class Processor:
     async def generate_index_message(
         self,
         resource: Resource,
-        messages: list[writer_pb2.BrokerMessage],
+        message: writer_pb2.BrokerMessage,
         resource_created: bool,
     ) -> PBBrainResource:
         builder = IndexMessageBuilder(resource)
-        message_source = messages_source(messages)
-        if message_source == nodewriter_pb2.IndexMessageSource.WRITER:
-            return await builder.for_writer_bm(messages, resource_created)
-        elif message_source == nodewriter_pb2.IndexMessageSource.PROCESSOR:
-            return await builder.for_processor_bm(messages)
+        if message.source == writer_pb2.BrokerMessage.MessageSource.WRITER:
+            return await builder.for_writer_bm(message, resource_created)
+        elif message.source == writer_pb2.BrokerMessage.MessageSource.PROCESSOR:
+            return await builder.for_processor_bm(message)
         else:  # pragma: no cover
-            raise InvalidBrokerMessage(f"Unknown broker message source: {message_source}")
+            raise InvalidBrokerMessage(f"Unknown broker message source: {message.source}")
 
     async def external_index_delete_resource(
         self, external_index_manager: ExternalIndexManager, resource_uuid: str
@@ -582,35 +564,8 @@ class Processor:
             resource_uuid=resource_uuid, resource_data=index_message
         )
 
-    async def multi(self, message: writer_pb2.BrokerMessage, seqid: int) -> None:
-        self.messages.setdefault(message.multiid, []).append(message)
-
-    async def commit(self, message: writer_pb2.BrokerMessage, seqid: int, partition: str) -> None:
-        if message.multiid not in self.messages:
-            # Error
-            logger.error(f"Closed multi {message.multiid}")
-            await self.deadletter([message], partition, seqid)
-        else:
-            await self.txn(self.messages[message.multiid], seqid, partition)
-
-    async def rollback(self, message: writer_pb2.BrokerMessage, seqid: int, partition: str) -> None:
-        # Error
-        logger.error(f"Closed multi {message.multiid}")
-        del self.messages[message.multiid]
-        await self.notify_abort(
-            partition=partition,
-            seqid=seqid,
-            multi=message.multiid,
-            kbid=message.kbid,
-            rid=message.uuid,
-            source=message.source,
-        )
-
-    async def deadletter(
-        self, messages: list[writer_pb2.BrokerMessage], partition: str, seqid: int
-    ) -> None:
-        for seq, message in enumerate(messages):
-            await self.storage.deadletter(message, seq, seqid, partition)
+    async def deadletter(self, message: writer_pb2.BrokerMessage, partition: str, seqid: int) -> None:
+        await self.storage.deadletter(message, 0, seqid, partition)
 
     @processor_observer.wrap({"type": "apply_resource"})
     async def apply_resource(
@@ -670,7 +625,6 @@ class Processor:
         *,
         partition: str,
         seqid: int,
-        multi: str,
         message: writer_pb2.BrokerMessage,
         write_type: writer_pb2.Notification.WriteType.ValueType,
     ):
@@ -678,7 +632,7 @@ class Processor:
         notification = writer_pb2.Notification(
             partition=int(partition),
             seqid=seqid,
-            multi=multi,
+            multi="",
             uuid=message.uuid,
             kbid=message.kbid,
             action=writer_pb2.Notification.Action.COMMIT,
@@ -698,7 +652,6 @@ class Processor:
         *,
         partition: str,
         seqid: int,
-        multi: str,
         kbid: str,
         rid: str,
         source: writer_pb2.BrokerMessage.MessageSource.ValueType,
@@ -706,7 +659,7 @@ class Processor:
         message = writer_pb2.Notification(
             partition=int(partition),
             seqid=seqid,
-            multi=multi,
+            multi="",
             uuid=rid,
             kbid=kbid,
             action=writer_pb2.Notification.ABORT,
@@ -731,7 +684,7 @@ class Processor:
             logger.info(f"Skip when resource does not even have basic metadata: {resource}")
             return
         try:
-            async with self.driver.transaction() as txn:
+            async with self.driver.rw_transaction() as txn:
                 kb.txn = resource.txn = txn
                 resource.basic.metadata.status = resources_pb2.Metadata.Status.ERROR
                 await resource.set_basic(resource.basic)
@@ -759,23 +712,16 @@ class Processor:
         return kbobj
 
 
-def messages_source(messages: list[writer_pb2.BrokerMessage]):
-    from_writer = all(
-        (message.source == writer_pb2.BrokerMessage.MessageSource.WRITER for message in messages)
-    )
-    from_processor = all(
-        (message.source == writer_pb2.BrokerMessage.MessageSource.PROCESSOR for message in messages)
-    )
-    if from_writer:
-        source = nodewriter_pb2.IndexMessageSource.WRITER
-    elif from_processor:
-        source = nodewriter_pb2.IndexMessageSource.PROCESSOR
+def to_index_message_source(message: writer_pb2.BrokerMessage):
+    if message.source == writer_pb2.BrokerMessage.MessageSource.WRITER:
+        return nodewriter_pb2.IndexMessageSource.WRITER
+    elif message.source == writer_pb2.BrokerMessage.MessageSource.PROCESSOR:
+        return nodewriter_pb2.IndexMessageSource.PROCESSOR
     else:  # pragma: no cover
-        msg = "Processor received multiple broker messages with different sources in the same txn!"
+        msg = f"Processor received a broker message with unexpected source! {message.source}"
         logger.error(msg)
         errors.capture_exception(Exception(msg))
-        source = nodewriter_pb2.IndexMessageSource.PROCESSOR
-    return source
+        return nodewriter_pb2.IndexMessageSource.PROCESSOR
 
 
 def has_vectors_operation(index_message: PBBrainResource) -> bool:

nucliadb/ingest/orm/processor/sequence_manager.py

@@ -31,7 +31,7 @@ async def get_last_seqid(driver: Driver, worker: str) -> Optional[int]:
     This is oriented towards the ingest consumer and processor,
     which is the only one that should be writing to this key.
     """
-    async with driver.transaction(read_only=True) as txn:
+    async with driver.ro_transaction() as txn:
         key = TXNID.format(worker=worker)
         last_seq = await txn.get(key)
         if not last_seq:

nucliadb/ingest/orm/resource.py

@@ -135,7 +135,7 @@ class Resource:
         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)
             self.basic = basic if basic is not None else PBBasic()
@@ -354,7 +354,16 @@ 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]:

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/serialize.py

@@ -135,7 +135,7 @@ async def serialize(
     slug: Optional[str] = None,
 ) -> Optional[Resource]:
     driver = get_driver()
-    async with driver.transaction(read_only=True) as txn:
+    async with driver.ro_transaction() as txn:
         return await managed_serialize(
             txn,
             kbid,
@@ -392,6 +392,6 @@ async def get_resource_uuid_by_slug(
 ) -> Optional[str]:
     storage = await get_storage(service_name=service_name)
     driver = get_driver()
-    async with driver.transaction(read_only=True) as txn:
+    async with driver.ro_transaction() as txn:
         kb = KnowledgeBox(txn, storage, kbid)
         return await kb.get_resource_uuid_by_slug(slug)