nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/ingest/orm/processor/auditing.py

@@ -0,0 +1,117 @@
+# 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 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
+
+
+async def collect_audit_fields(
+    driver: Driver, storage: Storage, message: writer_pb2.BrokerMessage
+) -> list[audit_pb2.AuditField]:
+    if message.type == writer_pb2.BrokerMessage.MessageType.DELETE:
+        # If we are fully deleting a resource we won't iterate the delete_fields (if any).
+        # Make no sense as we already collected all resource fields as deleted
+        return []
+
+    audit_storage_fields: list[audit_pb2.AuditField] = []
+    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()
+
+        for field_id, field_type in iterate_auditable_fields(field_keys, message):
+            auditfield = audit_pb2.AuditField()
+            auditfield.field_type = field_type
+            auditfield.field_id = field_id
+            if field_type is writer_pb2.FieldType.FILE:
+                auditfield.filename = message.files[field_id].file.filename
+            # The field did exist, so we are overwriting it, with a modified file
+            # in case of a file
+            auditfield.action = audit_pb2.AuditField.FieldAction.MODIFIED
+            if field_type is writer_pb2.FieldType.FILE:
+                auditfield.size = message.files[field_id].file.size
+
+            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)
+            audit_field = audit_pb2.AuditField()
+            audit_field.action = audit_pb2.AuditField.FieldAction.DELETED
+            audit_field.field_id = fieldid.field
+            audit_field.field_type = fieldid.field_type
+            if fieldid.field_type is writer_pb2.FieldType.FILE:
+                val = await field.get_value()
+                audit_field.size = 0
+                if val is not None:
+                    audit_field.filename = val.file.filename
+            audit_storage_fields.append(audit_field)
+
+    return audit_storage_fields
+
+
+def iterate_auditable_fields(
+    resource_keys: list[tuple[FieldType.ValueType, str]],
+    message: writer_pb2.BrokerMessage,
+):
+    """
+    Generator that emits the combined list of field ids from both
+    the existing resource and message that needs to be considered
+    in the audit of fields.
+    """
+    yielded = set()
+
+    # Include all fields present in the message we are processing
+    for field_id in message.files.keys():
+        key = (field_id, writer_pb2.FieldType.FILE)
+        yield key
+        yielded.add(key)
+
+    for field_id in message.conversations.keys():
+        key = (field_id, writer_pb2.FieldType.CONVERSATION)
+        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.links.keys():
+        key = (field_id, writer_pb2.FieldType.LINK)
+        yield key
+        yielded.add(key)
+
+    for field_type, field_id in resource_keys:
+        if field_type is writer_pb2.FieldType.GENERIC:
+            continue
+
+        if not (
+            field_id in message.files or message.type is writer_pb2.BrokerMessage.MessageType.DELETE
+        ):
+            continue
+
+        # Avoid duplicates
+        if (field_type, field_id) in yielded:
+            continue
+
+        yield (field_id, field_type)

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}
+        )