nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.whl

nucliadb/ingest/orm/processor/auditing.py

@@ -17,12 +17,11 @@
 # 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 nucliadb_protos.resources_pb2 import FieldType
-
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb_protos import audit_pb2, writer_pb2
+from nucliadb_protos.resources_pb2 import FieldType
 from nucliadb_utils.storages.storage import Storage
 
 
@@ -35,7 +34,7 @@ async def collect_audit_fields(
         return []
 
     audit_storage_fields: list[audit_pb2.AuditField] = []
-    async with driver.transaction() as txn:
+    async with driver.transaction(read_only=True) as txn:
         kb = KnowledgeBox(txn, storage, message.kbid)
         resource = Resource(txn, storage, kb, message.uuid)
         field_keys = await resource.get_fields_ids()
@@ -55,9 +54,7 @@ async def collect_audit_fields(
             audit_storage_fields.append(auditfield)
 
         for fieldid in message.delete_fields or []:
-            field = await resource.get_field(
-                fieldid.field, writer_pb2.FieldType.FILE, load=True
-            )
+            field = await resource.get_field(fieldid.field, writer_pb2.FieldType.FILE, load=True)
             audit_field = audit_pb2.AuditField()
             audit_field.action = audit_pb2.AuditField.FieldAction.DELETED
             audit_field.field_id = fieldid.field
@@ -94,26 +91,11 @@ def iterate_auditable_fields(
         yield key
         yielded.add(key)
 
-    for field_id in message.layouts.keys():
-        key = (field_id, writer_pb2.FieldType.LAYOUT)
-        yield key
-        yielded.add(key)
-
     for field_id in message.texts.keys():
         key = (field_id, writer_pb2.FieldType.TEXT)
         yield key
         yielded.add(key)
 
-    for field_id in message.keywordsets.keys():
-        key = (field_id, writer_pb2.FieldType.KEYWORDSET)
-        yield key
-        yielded.add(key)
-
-    for field_id in message.datetimes.keys():
-        key = (field_id, writer_pb2.FieldType.DATETIME)
-        yield key
-        yielded.add(key)
-
     for field_id in message.links.keys():
         key = (field_id, writer_pb2.FieldType.LINK)
         yield key
@@ -124,8 +106,7 @@ def iterate_auditable_fields(
             continue
 
         if not (
-            field_id in message.files
-            or message.type is writer_pb2.BrokerMessage.MessageType.DELETE
+            field_id in message.files or message.type is writer_pb2.BrokerMessage.MessageType.DELETE
         ):
             continue
 

nucliadb/ingest/orm/processor/data_augmentation.py

@@ -0,0 +1,164 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+
+import logging
+from dataclasses import dataclass, field
+from typing import Optional, Sequence
+
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.ingest.processing import ProcessingEngine, PushPayload, Source
+from nucliadb_models.text import PushTextFormat, Text
+from nucliadb_protos import resources_pb2, writer_pb2
+from nucliadb_protos.resources_pb2 import FieldID, FieldType
+from nucliadb_utils.utilities import Utility, get_partitioning, get_utility
+
+logger = logging.getLogger("ingest-processor")
+
+
+@dataclass
+class GeneratedFields:
+    texts: list[str] = field(default_factory=list)
+    links: list[str] = field(default_factory=list)
+    files: list[str] = field(default_factory=list)
+    conversations: list[str] = field(default_factory=list)
+
+    def is_not_empty(self) -> bool:
+        return (len(self.texts) + len(self.links) + len(self.files) + len(self.conversations)) > 0
+
+
+async def get_generated_fields(bm: writer_pb2.BrokerMessage, resource: Resource) -> GeneratedFields:
+    """Processing can send messages with generated fields. Those can be
+    generated with a data augmentation task and, as learning can't queue it to
+    process, nucliadb is responsible to send the generated field to process (and
+    ingest the processed thing later).
+
+    Given a broker message and a resource, this function returns the list of
+    generated fields, that can be empty.
+
+    """
+    generated_fields = GeneratedFields()
+
+    # only messages from processor can add generated fields
+    if bm.source != writer_pb2.BrokerMessage.MessageSource.PROCESSOR:
+        return generated_fields
+
+    # search all fields
+
+    all_fields = await resource.get_all_field_ids(for_update=False)
+    fields: Sequence[FieldID]
+    if all_fields is None:
+        fields = []
+    else:
+        fields = all_fields.fields
+
+    for field_id in bm.texts:
+        field = FieldID(field_type=FieldType.TEXT, field=field_id)
+        if field not in fields:
+            generated_fields.texts.append(field_id)
+
+    for field_id in bm.links:
+        field = FieldID(field_type=FieldType.LINK, field=field_id)
+        if field not in fields:
+            generated_fields.links.append(field_id)
+
+    for field_id in bm.files:
+        field = FieldID(field_type=FieldType.FILE, field=field_id)
+        if field not in fields:
+            generated_fields.files.append(field_id)
+
+    for field_id in bm.conversations:
+        field = FieldID(field_type=FieldType.CONVERSATION, field=field_id)
+        if field not in fields:
+            generated_fields.conversations.append(field_id)
+
+    return generated_fields
+
+
+async def send_generated_fields_to_process(
+    kbid: str,
+    resource: Resource,
+    generated_fields: GeneratedFields,
+    bm: writer_pb2.BrokerMessage,
+):
+    partitioning = get_partitioning()
+    partition = partitioning.generate_partition(kbid, resource.uuid)
+
+    processing: ProcessingEngine = get_utility(Utility.PROCESSING)
+
+    payload = _generate_processing_payload_for_fields(kbid, resource.uuid, generated_fields, bm)
+    if payload is not None:
+        processing_info = await processing.send_to_process(payload, partition)
+        logger.info(
+            "Sent generated fields to process",
+            extra={"processing_info": processing_info},
+        )
+
+
+def _generate_processing_payload_for_fields(
+    kbid: str,
+    rid: str,
+    fields: GeneratedFields,
+    bm: writer_pb2.BrokerMessage,
+) -> Optional[PushPayload]:
+    partitioning = get_partitioning()
+    partition = partitioning.generate_partition(kbid, rid)
+
+    payload = PushPayload(kbid=kbid, uuid=rid, userid="nucliadb-ingest", partition=partition)
+
+    payload.kbid = bm.kbid
+    payload.uuid = rid
+    payload.source = Source.INGEST
+    payload.slug = bm.slug
+
+    # populate generated fields
+
+    for text in fields.texts:
+        payload.textfield[text] = _bm_text_field_to_processing(bm.texts[text])
+
+    for file in fields.files:
+        logger.warning(
+            "Ingest received a broker message from processor with a new file field! Skipping",
+            extra={"kbid": kbid, "rid": rid, "field_id": file},
+        )
+        pass
+
+    for link in fields.links:
+        logger.warning(
+            "Ingest received a broker message from processor with a new link field!",
+            extra={"kbid": kbid, "rid": rid, "field_id": link},
+        )
+        pass
+
+    for conversation in fields.conversations:
+        logger.warning(
+            "Ingest received a broker message from processor with a new conversation field! Skipping",
+            extra={"kbid": kbid, "rid": rid, "field_id": conversation},
+        )
+        pass
+
+    if len(fields.texts) > 0:
+        return payload
+    else:
+        # we don't want to send weird empty messages to processing
+        return None
+
+
+def _bm_text_field_to_processing(text_field: resources_pb2.FieldText) -> Text:
+    return Text(body=text_field.body, format=PushTextFormat(text_field.format))

nucliadb/ingest/orm/processor/pgcatalog.py

@@ -0,0 +1,84 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 cast
+
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb.common.maindb.pg import PGDriver, PGTransaction
+from nucliadb.common.maindb.utils import get_driver
+from nucliadb_telemetry import metrics
+
+from ..resource import Resource
+
+observer = metrics.Observer("pg_catalog_write", labels={"type": ""})
+
+
+def _pg_transaction(txn: Transaction) -> PGTransaction:
+    return cast(PGTransaction, txn)
+
+
+def pgcatalog_enabled(kbid):
+    return isinstance(get_driver(), PGDriver)
+
+
+@observer.wrap({"type": "update"})
+async def pgcatalog_update(txn: Transaction, kbid: str, resource: Resource):
+    if not pgcatalog_enabled(kbid):
+        return
+
+    if resource.basic is None:
+        raise ValueError("Cannot index into the catalog a resource without basic metadata ")
+
+    created_at = resource.basic.created.ToDatetime()
+    modified_at = resource.basic.modified.ToDatetime()
+    if modified_at < created_at:
+        modified_at = created_at
+
+    async with _pg_transaction(txn).connection.cursor() as cur:
+        await cur.execute(
+            """
+            INSERT INTO catalog
+            (kbid, rid, title, created_at, modified_at, labels)
+            VALUES
+            (%(kbid)s, %(rid)s, %(title)s, %(created_at)s, %(modified_at)s, %(labels)s)
+            ON CONFLICT (kbid, rid) DO UPDATE SET
+            title = excluded.title,
+            created_at = excluded.created_at,
+            modified_at = excluded.modified_at,
+            labels = excluded.labels""",
+            {
+                "kbid": resource.kb.kbid,
+                "rid": resource.uuid,
+                "title": resource.basic.title,
+                "created_at": created_at,
+                "modified_at": modified_at,
+                "labels": list(resource.indexer.brain.labels),
+            },
+        )
+
+
+@observer.wrap({"type": "delete"})
+async def pgcatalog_delete(txn: Transaction, kbid: str, rid: str):
+    if not pgcatalog_enabled(kbid):
+        return
+    async with _pg_transaction(txn).connection.cursor() as cur:
+        await cur.execute(
+            "DELETE FROM catalog where kbid = %(kbid)s AND rid = %(rid)s", {"kbid": kbid, "rid": rid}
+        )