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

nucliadb/export_import/utils.py

@@ -18,25 +18,27 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import functools
-from io import BytesIO
 from typing import AsyncGenerator, AsyncIterator, Callable, Optional
 
-import nats.errors
 from google.protobuf.message import DecodeError as ProtobufDecodeError
 
+from nucliadb import learning_proxy
 from nucliadb.common import datamanagers
 from nucliadb.common.context import ApplicationContext
 from nucliadb.export_import import logger
 from nucliadb.export_import.datamanager import ExportImportDataManager
 from nucliadb.export_import.exceptions import (
     ExportStreamExhausted,
+    IncompatibleExport,
     WrongExportStreamFormat,
 )
 from nucliadb.export_import.models import ExportedItemType, ExportItem, Metadata
+from nucliadb.ingest.orm.broker_message import generate_broker_message
 from nucliadb_models.export_import import Status
 from nucliadb_protos import knowledgebox_pb2 as kb_pb2
 from nucliadb_protos import resources_pb2, writer_pb2
 from nucliadb_utils.const import Streams
+from nucliadb_utils.transaction import MaxTransactionSizeExceededError
 
 BinaryStream = AsyncGenerator[bytes, None]
 BinaryStreamGenerator = Callable[[int], BinaryStream]
@@ -50,7 +52,6 @@ PROCESSING_BM_FIELDS = [
     "field_metadata",
     "field_vectors",
     "field_large_metadata",
-    "user_vectors",
 ]
 
 # Broker message fields that are populated by the nucliadb writer component
@@ -59,9 +60,6 @@ WRITER_BM_FIELDS = [
     "files",
     "texts",
     "conversations",
-    "layouts",
-    "keywordsets",
-    "datetimes",
 ]
 
 
@@ -88,7 +86,7 @@ async def transaction_commit(
             wait=False,
             target_subject=Streams.INGEST_PROCESSED.subject,
         )
-    except nats.errors.MaxPayloadError:
+    except MaxTransactionSizeExceededError:
         stored_key = await context.blob_storage.set_stream_message(
             kbid=bm.kbid, rid=bm.uuid, data=bm.SerializeToString()
         )
@@ -151,23 +149,17 @@ async def set_entities_groups(
     context: ApplicationContext, kbid: str, entities_groups: kb_pb2.EntitiesGroups
 ) -> None:
     async with datamanagers.with_transaction() as txn:
-        await datamanagers.entities.set_entities_groups(
-            txn, kbid=kbid, entities_groups=entities_groups
-        )
+        await datamanagers.entities.set_entities_groups(txn, kbid=kbid, entities_groups=entities_groups)
         await txn.commit()
 
 
-async def set_labels(
-    context: ApplicationContext, kbid: str, labels: kb_pb2.Labels
-) -> None:
+async def set_labels(context: ApplicationContext, kbid: str, labels: kb_pb2.Labels) -> None:
     async with datamanagers.with_transaction() as txn:
         await datamanagers.labels.set_labels(txn, kbid=kbid, labels=labels)
         await txn.commit()
 
 
-async def iter_kb_resource_uuids(
-    context: ApplicationContext, kbid: str
-) -> AsyncGenerator[str, None]:
+async def iter_kb_resource_uuids(context: ApplicationContext, kbid: str) -> AsyncGenerator[str, None]:
     async for rid in datamanagers.resources.iterate_resource_ids(kbid=kbid):
         yield rid
 
@@ -175,8 +167,13 @@ async def iter_kb_resource_uuids(
 async def get_broker_message(
     context: ApplicationContext, kbid: str, rid: str
 ) -> Optional[writer_pb2.BrokerMessage]:
-    async with datamanagers.with_transaction() as txn:
-        return await datamanagers.resources.get_broker_message(txn, kbid=kbid, rid=rid)
+    async with datamanagers.with_ro_transaction() as txn:
+        resource = await datamanagers.resources.get_resource(txn, kbid=kbid, rid=rid)
+        if resource is None:
+            return None
+        resource.disable_vectors = False
+        resource.txn = txn
+        return await generate_broker_message(resource)
 
 
 def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFile]:
@@ -184,6 +181,10 @@ def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFil
     binaries: list[resources_pb2.CloudFile] = []
     for file_field in bm.files.values():
         if file_field.HasField("file"):
+            if file_field.file.source == resources_pb2.CloudFile.Source.EXTERNAL:
+                # Binaries of externally hosted files are not
+                # to be downloaded and stored in the export file
+                continue
             _clone_collect_cf(binaries, file_field.file)
 
     for conversation in bm.conversations.values():
@@ -191,11 +192,6 @@ def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFil
             for attachment in message.content.attachments:
                 _clone_collect_cf(binaries, attachment)
 
-    for layout in bm.layouts.values():
-        for block in layout.body.blocks.values():
-            if block.HasField("file"):
-                _clone_collect_cf(binaries, block.file)
-
     for field_extracted_data in bm.file_extracted_data:
         if field_extracted_data.HasField("file_thumbnail"):
             _clone_collect_cf(binaries, field_extracted_data.file_thumbnail)
@@ -213,6 +209,8 @@ def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFil
             _clone_collect_cf(binaries, link_extracted_data.link_preview)
         if link_extracted_data.HasField("link_image"):
             _clone_collect_cf(binaries, link_extracted_data.link_image)
+        for file_generated in link_extracted_data.file_generated.values():
+            _clone_collect_cf(binaries, file_generated)
 
     for field_metadata in bm.field_metadata:
         if field_metadata.metadata.metadata.HasField("thumbnail"):
@@ -225,9 +223,7 @@ def get_cloud_files(bm: writer_pb2.BrokerMessage) -> list[resources_pb2.CloudFil
     return binaries
 
 
-def _clone_collect_cf(
-    binaries: list[resources_pb2.CloudFile], origin: resources_pb2.CloudFile
-):
+def _clone_collect_cf(binaries: list[resources_pb2.CloudFile], origin: resources_pb2.CloudFile):
     cf = resources_pb2.CloudFile()
     cf.CopyFrom(origin)
     # Mark the cloud file of the broker message being exported as export source
@@ -244,43 +240,21 @@ async def download_binary(
 
 
 async def get_entities(context: ApplicationContext, kbid: str) -> kb_pb2.EntitiesGroups:
-    async with datamanagers.with_transaction() as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         return await datamanagers.entities.get_entities_groups(txn, kbid=kbid)
 
 
 async def get_labels(context: ApplicationContext, kbid: str) -> kb_pb2.Labels:
-    async with datamanagers.with_transaction() as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         return await datamanagers.labels.get_labels(txn, kbid=kbid)
 
 
-class EndOfStream(Exception):
-    ...
+class EndOfStream(Exception): ...
 
 
 class ExportStream:
     """
     Models a stream of export bytes that can be read from asynchronously.
-    """
-
-    def __init__(self, export: BytesIO):
-        self.export = export
-        self.read_bytes = 0
-        self._length = len(export.getvalue())
-
-    async def read(self, n_bytes):
-        """
-        Reads n_bytes from the export stream.
-        Raises ExportStreamExhausted if there are no more bytes to read.
-        """
-        if self.read_bytes == self._length:
-            raise ExportStreamExhausted()
-        chunk = self.export.read(n_bytes)
-        self.read_bytes += len(chunk)
-        return chunk
-
-
-class IteratorExportStream(ExportStream):
-    """
     Adapts the parent class to be able to read bytes yielded from an async iterator.
     """
 
@@ -325,8 +299,8 @@ class ExportStreamReader:
     yields the deserialized export items ready to be imported.
     """
 
-    def __init__(self, export_stream: ExportStream):
-        self.stream = export_stream
+    def __init__(self, stream: AsyncGenerator[bytes, None]):
+        self.stream = ExportStream(stream)
 
     @property
     def read_bytes(self) -> int:
@@ -401,6 +375,26 @@ class ExportStreamReader:
             raise WrongExportStreamFormat() from ex
         return labels
 
+    async def maybe_read_learning_config(
+        self,
+    ) -> tuple[Optional[learning_proxy.LearningConfiguration], bytes]:
+        """
+        Tries to read a learning config from the beginning of the stream.
+        Returs the learning config if found. It also returns any leftover bytes that
+        may have been read from the network into memory that need to be yielded and imported.
+        """
+        # We assume that the learning config is the first item in the export stream.
+        try:
+            type_bytes = await self.stream.read(3)
+        except ExportStreamExhausted:
+            return None, self.stream.buffer
+        if type_bytes != ExportedItemType.LEARNING_CONFIG.value.encode():
+            # Backward compatible code for old exports that don't have a learning config.
+            return None, type_bytes + self.stream.buffer
+        data = await self.read_item()
+        lconfig = learning_proxy.LearningConfiguration.model_validate_json(data)
+        return lconfig, self.stream.buffer
+
     async def iter_items(self) -> AsyncGenerator[ExportItem, None]:
         while True:
             try:
@@ -411,7 +405,7 @@ class ExportStreamReader:
                     ExportedItemType.ENTITIES: self.read_entities,
                     ExportedItemType.LABELS: self.read_labels,
                 }[item_type]
-                data = await read_data_func()  # type: ignore
+                data = await read_data_func()
                 yield item_type, data
             except ExportStreamExhausted:
                 break
@@ -476,3 +470,49 @@ class TaskRetryHandler:
                 await self.dm.set_metadata(self.type, metadata)
 
         return wrapper
+
+
+async def get_learning_config(
+    kbid: str,
+) -> Optional[learning_proxy.LearningConfiguration]:
+    return await learning_proxy.get_configuration(kbid)
+
+
+def stream_compatible_with_kb(
+    kbid: str, stream: AsyncGenerator[bytes, None]
+) -> AsyncGenerator[bytes, None]:
+    """
+    Wrapper around an export stream that checks if the export is compatible with the destination knowledge box.
+    """
+
+    async def wrapped() -> AsyncGenerator[bytes, None]:
+        # Read the a few bytes from the beginning of the stream to check the semantic model.
+        # If the semantic model is not compatible, raise an exception.
+        # If there are leftover bytes, yield them.
+        leftover_bytes = await _check_semantic_model_compatibility(kbid, stream)
+        if len(leftover_bytes) > 0:
+            yield leftover_bytes
+
+        # Now yield the rest of the stream
+        async for chunk in stream:
+            yield chunk
+
+    return wrapped()
+
+
+async def _check_semantic_model_compatibility(kbid: str, stream: AsyncGenerator[bytes, None]) -> bytes:
+    stream_reader = ExportStreamReader(stream)
+    lconfig, leftover_bytes = await stream_reader.maybe_read_learning_config()
+    if lconfig is None:
+        logger.warning("Learning config not found on the export stream. Export may be incompatible.")
+        return leftover_bytes
+    kb_lconfig = await get_learning_config(kbid)
+    if kb_lconfig is None:
+        logger.warning("No learning config found on the knowledge box. Export may be incompatible.")
+        return leftover_bytes
+    if kb_lconfig.semantic_model == lconfig.semantic_model:
+        logger.info(f"Semantic model match: {kb_lconfig.semantic_model}")
+        return leftover_bytes
+    raise IncompatibleExport(
+        f"Cannot import. Semantic model mismatch: {kb_lconfig.semantic_model} != {lconfig.semantic_model}"
+    )

nucliadb/health.py

@@ -78,9 +78,7 @@ async def grpc_health_check(health_servicer) -> None:
         for check in _health_checks:
             if not check():
                 logger.info(f"Health check failed on {check.__name__}")
-                await health_servicer.set(
-                    "", health_pb2.HealthCheckResponse.NOT_SERVING
-                )
+                await health_servicer.set("", health_pb2.HealthCheckResponse.NOT_SERVING)
                 break
         else:
             await health_servicer.set("", health_pb2.HealthCheckResponse.SERVING)

nucliadb/ingest/app.py

@@ -18,10 +18,9 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
+import importlib.metadata
 from typing import Awaitable, Callable
 
-import pkg_resources
-
 from nucliadb import health
 from nucliadb.common.cluster.discovery.utils import (
     setup_cluster_discovery,
@@ -30,10 +29,12 @@ from nucliadb.common.cluster.discovery.utils import (
 from nucliadb.common.cluster.settings import settings as cluster_settings
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.context import ApplicationContext
+from nucliadb.common.nidx import start_nidx_utility
 from nucliadb.export_import.tasks import get_exports_consumer, get_imports_consumer
 from nucliadb.ingest import SERVICE_NAME
 from nucliadb.ingest.consumer import service as consumer_service
 from nucliadb.ingest.partitions import assign_partitions
+from nucliadb.ingest.processing import start_processing_engine, stop_processing_engine
 from nucliadb.ingest.service import start_grpc
 from nucliadb.ingest.settings import settings
 from nucliadb_telemetry import errors
@@ -46,10 +47,12 @@ from nucliadb_utils.utilities import (
     start_audit_utility,
     start_indexing_utility,
     start_nats_manager,
+    start_partitioning_utility,
     start_transaction_utility,
     stop_audit_utility,
     stop_indexing_utility,
     stop_nats_manager,
+    stop_partitioning_utility,
     stop_transaction_utility,
 )
 
@@ -59,15 +62,17 @@ async def initialize() -> list[Callable[[], Awaitable[None]]]:
 
     await setup_cluster()
     await start_transaction_utility(SERVICE_NAME)
-    if (
-        not cluster_settings.standalone_mode
-        and indexing_settings.index_jetstream_servers is not None
-    ):
+    if not cluster_settings.standalone_mode and indexing_settings.index_jetstream_servers is not None:
         await start_indexing_utility(SERVICE_NAME)
 
+    start_partitioning_utility()
+
+    await start_nidx_utility()
+
     await start_audit_utility(SERVICE_NAME)
 
     finalizers = [
+        stop_partitioning_utility,
         stop_transaction_utility,
         stop_indexing_utility,
         stop_audit_utility,
@@ -123,8 +128,7 @@ async def main_consumer():  # pragma: no cover
     ingest_consumers = await consumer_service.start_ingest_consumers(SERVICE_NAME)
 
     await run_until_exit(
-        [grpc_health_finalizer, pull_workers, ingest_consumers, metrics_server.shutdown]
-        + finalizers
+        [grpc_health_finalizer, pull_workers, ingest_consumers, metrics_server.shutdown] + finalizers
     )
 
 
@@ -138,12 +142,13 @@ async def main_orm_grpc():  # pragma: no cover
 async def main_ingest_processed_consumer():  # pragma: no cover
     finalizers = await initialize()
 
+    await start_processing_engine()
     metrics_server = await serve_metrics()
     grpc_health_finalizer = await health.start_grpc_health_service(settings.grpc_port)
     consumer = await consumer_service.start_ingest_processed_consumer(SERVICE_NAME)
 
     await run_until_exit(
-        [grpc_health_finalizer, consumer, metrics_server.shutdown] + finalizers
+        [grpc_health_finalizer, consumer, metrics_server.shutdown, stop_processing_engine] + finalizers
     )
 
 
@@ -181,10 +186,9 @@ async def main_subscriber_workers():  # pragma: no cover
 
 def setup_configuration():  # pragma: no cover
     setup_logging()
-
     assign_partitions(settings)
 
-    errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
+    errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
     if asyncio._get_running_loop() is not None:
         raise RuntimeError("cannot be called from a running event loop")

nucliadb/ingest/consumer/auditing.py

@@ -23,14 +23,11 @@ import logging
 import uuid
 from functools import partial
 
-from nucliadb_protos.resources_pb2 import FieldType
-
+from nucliadb.common import datamanagers
 from nucliadb.common.cluster.exceptions import ShardsNotFound
 from nucliadb.common.cluster.manager import choose_node
 from nucliadb.common.cluster.utils import get_shard_manager
-from nucliadb.common.maindb.driver import Driver
-from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
-from nucliadb.ingest.orm.resource import Resource
+from nucliadb.common.constants import AVG_PARAGRAPH_SIZE_BYTES
 from nucliadb_protos import audit_pb2, nodereader_pb2, noderesources_pb2, writer_pb2
 from nucliadb_utils import const
 from nucliadb_utils.audit.audit import AuditStorage
@@ -62,12 +59,10 @@ class IndexAuditHandler:
     def __init__(
         self,
         *,
-        driver: Driver,
         audit: AuditStorage,
         pubsub: PubSubDriver,
         check_delay: float = 5.0,
     ):
-        self.driver = driver
         self.audit = audit
         self.pubsub = pubsub
         self.shard_manager = get_shard_manager()
@@ -98,17 +93,15 @@ class IndexAuditHandler:
             metrics.total_messages.inc({"action": "ignored", "type": "audit_counter"})
             return
 
-        self.task_handler.schedule(
-            notification.kbid, partial(self.process_kb, notification.kbid)
-        )
+        self.task_handler.schedule(notification.kbid, partial(self.process_kb, notification.kbid))
         metrics.total_messages.inc({"action": "scheduled", "type": "audit_counter"})
 
     @metrics.handler_histo.wrap({"type": "audit_counter"})
     async def process_kb(self, kbid: str) -> None:
         try:
-            shard_groups: list[
-                writer_pb2.ShardObject
-            ] = await self.shard_manager.get_shards_by_kbid(kbid)
+            shard_groups: list[writer_pb2.ShardObject] = await self.shard_manager.get_shards_by_kbid(
+                kbid
+            )
         except ShardsNotFound:
             logger.warning(f"No shards found for kbid {kbid}, skipping")
             return
@@ -119,7 +112,8 @@ class IndexAuditHandler:
         total_paragraphs = 0
 
         for shard_obj in shard_groups:
-            node, shard_id = choose_node(shard_obj)
+            # TODO: Uses node for auditing, don't want to suddenly change metrics
+            node, shard_id = choose_node(shard_obj, use_nidx=False)
             shard: nodereader_pb2.Shard = await node.reader.GetShard(
                 nodereader_pb2.GetShardRequest(shard_id=noderesources_pb2.ShardId(id=shard_id))  # type: ignore
             )
@@ -127,12 +121,18 @@ class IndexAuditHandler:
             total_fields += shard.fields
             total_paragraphs += shard.paragraphs
 
-        await self.audit.report(
+        async with datamanagers.with_ro_transaction() as txn:
+            num_vectorsets = (
+                len([vs async for vs in datamanagers.vectorsets.iter(txn=txn, kbid=kbid)]) or 1
+            )
+
+        self.audit.report_storage(
             kbid=kbid,
-            audit_type=audit_pb2.AuditRequest.AuditType.INDEXED,
-            kb_counter=audit_pb2.AuditKBCounter(
-                fields=total_fields, paragraphs=total_paragraphs
-            ),
+            paragraphs=total_paragraphs,
+            fields=total_fields,
+            bytes=total_paragraphs  # This is an estimation of bytes stored in a KB
+            * AVG_PARAGRAPH_SIZE_BYTES
+            * num_vectorsets,
         )
 
 
@@ -147,12 +147,10 @@ class ResourceWritesAuditHandler:
     def __init__(
         self,
         *,
-        driver: Driver,
         storage: Storage,
         audit: AuditStorage,
         pubsub: PubSubDriver,
     ):
-        self.driver = driver
         self.storage = storage
         self.audit = audit
         self.pubsub = pubsub
@@ -169,117 +167,6 @@ class ResourceWritesAuditHandler:
     async def finalize(self) -> None:
         await self.pubsub.unsubscribe(self.subscription_id)
 
-    def iterate_auditable_fields(
-        self,
-        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.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
-            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)
-
-    async def collect_audit_fields(
-        self, 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 self.driver.transaction() as txn:
-            kb = KnowledgeBox(txn, self.storage, message.kbid)
-            resource = Resource(txn, self.storage, kb, message.uuid)
-            field_keys = await resource.get_fields_ids()
-
-            for field_id, field_type in self.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
-
     async def handle_message(self, raw_data) -> None:
         data = self.pubsub.parse(raw_data)
         notification = writer_pb2.Notification()
@@ -289,27 +176,23 @@ class ResourceWritesAuditHandler:
             metrics.total_messages.inc({"action": "ignored", "type": "audit_fields"})
             return
 
-        message = notification.message
-        if message.source == message.MessageSource.PROCESSOR:
+        message_audit: writer_pb2.Audit = notification.message_audit
+        if message_audit.message_source == writer_pb2.BrokerMessage.MessageSource.PROCESSOR:
             metrics.total_messages.inc({"action": "ignored", "type": "audit_fields"})
             return
 
-        logger.info(
-            {"message": "Processing field audit for kbid", "kbid": notification.kbid}
-        )
+        logger.info({"message": "Processing field audit for kbid", "kbid": notification.kbid})
 
         metrics.total_messages.inc({"action": "scheduled", "type": "audit_fields"})
         with metrics.handler_histo({"type": "audit_fields"}):
-            audit_fields = await self.collect_audit_fields(message)
-            field_metadata = [fi.field for fi in message.field_metadata]
-            when = message.audit.when if message.audit.HasField("when") else None
-            await self.audit.report(
-                kbid=message.kbid,
+            when = message_audit.when if message_audit.HasField("when") else None
+            self.audit.report_and_send(
+                kbid=message_audit.kbid,
                 when=when,
-                user=message.audit.user,
-                rid=message.uuid,
-                origin=message.audit.origin,
-                field_metadata=field_metadata,
+                user=message_audit.user,
+                rid=message_audit.uuid,
+                origin=message_audit.origin,
+                field_metadata=list(message_audit.field_metadata),
                 audit_type=AUDIT_TYPES.get(notification.write_type),
-                audit_fields=audit_fields,
+                audit_fields=list(message_audit.audit_fields),
             )