hyperboria/nexus/actions/update_document.py

from aiosumma import SummaHttpClient
from nexus.cognitron.schema import coders
from nexus.models.proto.operation_pb2 import \
    DocumentOperation as DocumentOperationPb

from .base import BaseAction


class SendDocumentOperationUpdateDocumentPbToSummaAction(BaseAction):
    def __init__(self, summa):
        super().__init__()
        self.summa_client = SummaHttpClient(**summa)
        self.waits.append(self.summa_client)

    async def do(self, document_operation_pb: DocumentOperationPb) -> DocumentOperationPb:
        update_document_pb = document_operation_pb.update_document
        schema = update_document_pb.typed_document.WhichOneof('document')
        document = getattr(update_document_pb.typed_document, schema)
        original_id = getattr(document, 'original_id', None)
        if not update_document_pb.reindex or original_id:
            return document_operation_pb
        document_tantivy = coders[schema].encode_document(document)
        await self.summa_client.put_document(schema, document_tantivy)
        if update_document_pb.commit:
            await self.summa_client.commit(schema)
        return document_operation_pb