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

nucliadb/ingest/orm/processor/processor.py

@@ -19,7 +19,6 @@
 #
 import asyncio
 import logging
-from typing import Optional
 
 import aiohttp.client_exceptions
 import nats.errors
@@ -28,6 +27,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 +61,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,16 +140,13 @@ class Processor:
     and can not use the txn id
     """
 
-    messages: dict[str, list[writer_pb2.BrokerMessage]]
-
     def __init__(
         self,
         driver: Driver,
         storage: Storage,
-        pubsub: Optional[PubSubDriver] = None,
-        partition: Optional[str] = None,
+        pubsub: PubSubDriver | None = None,
+        partition: str | None = None,
     ):
-        self.messages = {}
         self.driver = driver
         self.storage = storage
         self.partition = partition
@@ -162,7 +157,7 @@ class Processor:
         self,
         message: writer_pb2.BrokerMessage,
         seqid: int,
-        partition: Optional[str] = None,
+        partition: str | None = None,
         transaction_check: bool = True,
     ) -> None:
         partition = partition if self.partition is None else self.partition
@@ -180,18 +175,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:
@@ -227,7 +216,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 +232,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 +245,6 @@ class Processor:
         await self.notify_commit(
             partition=partition,
             seqid=seqid,
-            multi=message.multiid,
             message=message,
             write_type=writer_pb2.Notification.WriteType.DELETED,
         )
@@ -277,15 +265,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:
@@ -296,58 +281,55 @@ class Processor:
 
         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])
-                resource: Optional[Resource] = None
+                uuid = await self.get_resource_uuid(kb, message)
+
+                resource: Resource | None = 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 +339,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 +356,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 +367,6 @@ class Processor:
                     await self.notify_commit(
                         partition=partition,
                         seqid=seqid,
-                        multi=multi,
                         message=message,
                         write_type=(
                             writer_pb2.Notification.WriteType.CREATED
@@ -399,7 +379,6 @@ class Processor:
                     await self.notify_abort(
                         partition=partition,
                         seqid=seqid,
-                        multi=multi,
                         kbid=kbid,
                         rid=uuid,
                         source=message.source,
@@ -419,7 +398,6 @@ class Processor:
                 await self.notify_abort(
                     partition=partition,
                     seqid=seqid,
-                    multi=multi,
                     kbid=kbid,
                     rid=uuid,
                     source=message.source,
@@ -429,11 +407,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 +445,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 +501,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 +563,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 +624,6 @@ class Processor:
         *,
         partition: str,
         seqid: int,
-        multi: str,
         message: writer_pb2.BrokerMessage,
         write_type: writer_pb2.Notification.WriteType.ValueType,
     ):
@@ -678,7 +631,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 +651,6 @@ class Processor:
         *,
         partition: str,
         seqid: int,
-        multi: str,
         kbid: str,
         rid: str,
         source: writer_pb2.BrokerMessage.MessageSource.ValueType,
@@ -706,7 +658,7 @@ class Processor:
         message = writer_pb2.Notification(
             partition=int(partition),
             seqid=seqid,
-            multi=multi,
+            multi="",
             uuid=rid,
             kbid=kbid,
             action=writer_pb2.Notification.ABORT,
@@ -722,7 +674,7 @@ class Processor:
             await self.pubsub.publish(channel, payload)
 
     async def _mark_resource_error(
-        self, kb: KnowledgeBox, resource: Optional[Resource], partition: str, seqid: int
+        self, kb: KnowledgeBox, resource: Resource | None, partition: str, seqid: int
     ) -> None:
         """
         Unhandled error processing, try to mark resource as error
@@ -743,8 +695,8 @@ class Processor:
     # XXX: Why are these utility functions here?
     async def get_kb_obj(
         self, txn: Transaction, kbid: knowledgebox_pb2.KnowledgeBoxID
-    ) -> Optional[KnowledgeBox]:
-        uuid: Optional[str] = kbid.uuid
+    ) -> KnowledgeBox | None:
+        uuid: str | None = kbid.uuid
         if uuid == "":
             uuid = await datamanagers.kb.get_kb_uuid(txn, slug=kbid.slug)
 
@@ -759,23 +711,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

@@ -17,14 +17,13 @@
 # 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/>.
 #
-from typing import Optional
 
 from nucliadb.common.maindb.driver import Driver, Transaction
 
 TXNID = "/internal/worker/{worker}"
 
 
-async def get_last_seqid(driver: Driver, worker: str) -> Optional[int]:
+async def get_last_seqid(driver: Driver, worker: str) -> int | None:
     """
     Get last stored sequence id for a worker.