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

nucliadb/writer/lifecycle.py

@@ -17,13 +17,20 @@
 # 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.ingest.processing import start_processing_engine
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
+from nucliadb.common.context.fastapi import inject_app_context
+from nucliadb.ingest.processing import start_processing_engine, stop_processing_engine
 from nucliadb.ingest.utils import start_ingest, stop_ingest
 from nucliadb.writer import SERVICE_NAME
+from nucliadb.writer.back_pressure import start_materializer, stop_materializer
+from nucliadb.writer.settings import back_pressure_settings
 from nucliadb.writer.tus import finalize as storage_finalize
 from nucliadb.writer.tus import initialize as storage_initialize
-from nucliadb.writer.utilities import get_processing
 from nucliadb_telemetry.utils import clean_telemetry, setup_telemetry
+from nucliadb_utils.settings import is_onprem_nucliadb
 from nucliadb_utils.utilities import (
     finalize_utilities,
     start_partitioning_utility,
@@ -32,29 +39,28 @@ from nucliadb_utils.utilities import (
 )
 
 
-async def initialize():
-    await setup_telemetry(SERVICE_NAME)
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    back_pressure_enabled = back_pressure_settings.enabled and not is_onprem_nucliadb()
 
+    await setup_telemetry(SERVICE_NAME)
     await start_ingest(SERVICE_NAME)
-
     await start_processing_engine()
-
     start_partitioning_utility()
-
     await start_transaction_utility(SERVICE_NAME)
     await storage_initialize()
 
+    # Inject application context into the fastapi app's state
+    async with inject_app_context(app) as context:
+        if back_pressure_enabled:
+            await start_materializer(context)
+        yield
 
-async def finalize():
+    if back_pressure_enabled:
+        await stop_materializer()
     await stop_transaction_utility()
-
     await stop_ingest()
-    processing = get_processing()
-    if processing is not None:
-        await processing.finalize()
-
+    await stop_processing_engine()
     await storage_finalize()
-
     await clean_telemetry(SERVICE_NAME)
-
     await finalize_utilities()

nucliadb/writer/resource/audit.py

@@ -19,9 +19,10 @@
 #
 from datetime import datetime
 
-from nucliadb_protos.writer_pb2 import Audit
 from starlette.requests import Request
 
+from nucliadb_protos.writer_pb2 import Audit
+
 
 def parse_audit(audit: Audit, request: Request):
     audit.user = request.headers.get("X-NUCLIADB-USER", "")

nucliadb/writer/resource/basic.py

@@ -18,8 +18,31 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from datetime import datetime
+from typing import Optional
 
 from fastapi import HTTPException
+
+from nucliadb.common.models_utils import to_proto
+from nucliadb.common.models_utils.from_proto import (
+    RelationNodeTypeMap,
+    RelationTypeMap,
+)
+from nucliadb.ingest.orm.utils import set_title
+from nucliadb.ingest.processing import PushPayload
+from nucliadb_models.content_types import GENERIC_MIME_TYPE
+from nucliadb_models.file import FileField
+from nucliadb_models.link import LinkField
+from nucliadb_models.metadata import (
+    ParagraphAnnotation,
+    QuestionAnswerAnnotation,
+)
+from nucliadb_models.text import TEXT_FORMAT_TO_MIMETYPE, PushTextFormat, Text
+from nucliadb_models.writer import (
+    ComingResourcePayload,
+    CreateResourcePayload,
+    UpdateResourcePayload,
+)
+from nucliadb_protos.knowledgebox_pb2 import KnowledgeBoxConfig
 from nucliadb_protos.resources_pb2 import (
     Answers,
     Basic,
@@ -30,38 +53,19 @@ from nucliadb_protos.resources_pb2 import (
     Metadata,
     PageSelections,
     Paragraph,
+    TokenSplit,
+    UserFieldMetadata,
+    VisualSelection,
 )
 from nucliadb_protos.resources_pb2 import ParagraphAnnotation as PBParagraphAnnotation
 from nucliadb_protos.resources_pb2 import (
     QuestionAnswerAnnotation as PBQuestionAnswerAnnotation,
 )
-from nucliadb_protos.resources_pb2 import TokenSplit, UserFieldMetadata, VisualSelection
 from nucliadb_protos.utils_pb2 import Relation, RelationNode
 from nucliadb_protos.writer_pb2 import BrokerMessage
 
-from nucliadb.ingest.orm.utils import set_title
-from nucliadb.ingest.processing import PushPayload
-from nucliadb_models.common import FIELD_TYPES_MAP_REVERSE
-from nucliadb_models.file import FileField
-from nucliadb_models.link import LinkField
-from nucliadb_models.metadata import (
-    ParagraphAnnotation,
-    QuestionAnswerAnnotation,
-    RelationNodeTypeMap,
-    RelationTypeMap,
-)
-from nucliadb_models.text import TEXT_FORMAT_TO_MIMETYPE, PushTextFormat, Text
-from nucliadb_models.writer import (
-    GENERIC_MIME_TYPE,
-    ComingResourcePayload,
-    CreateResourcePayload,
-    UpdateResourcePayload,
-)
-
 
-def parse_basic_modify(
-    bm: BrokerMessage, item: ComingResourcePayload, toprocess: PushPayload
-):
+def parse_basic_modify(bm: BrokerMessage, item: ComingResourcePayload, toprocess: PushPayload):
     bm.basic.modified.FromDatetime(datetime.now())
     if item.title:
         set_title(bm, toprocess, item.title)
@@ -83,14 +87,9 @@ def parse_basic_modify(
         bm.basic.metadata.useful = True
         bm.basic.metadata.status = Metadata.Status.PENDING
 
-        toprocess.genericfield["summary"] = Text(
-            body=item.summary, format=PushTextFormat.PLAIN
-        )
+        toprocess.genericfield["summary"] = Text(body=item.summary, format=PushTextFormat.PLAIN)
     if item.thumbnail:
         bm.basic.thumbnail = item.thumbnail
-    if item.layout:
-        bm.basic.layout = item.layout
-
     if item.metadata is not None:
         bm.basic.metadata.metadata.update(item.metadata.metadata)
         if item.metadata.language:
@@ -147,9 +146,8 @@ def parse_basic_modify(
                 userfieldmetadata.question_answers.append(qa_annotation_pb)
 
             userfieldmetadata.field.field = fieldmetadata.field.field
-            userfieldmetadata.field.field_type = FIELD_TYPES_MAP_REVERSE[  # type: ignore
-                fieldmetadata.field.field_type.value
-            ]
+
+            userfieldmetadata.field.field_type = to_proto.field_type(fieldmetadata.field.field_type)
 
             bm.basic.fieldmetadata.append(userfieldmetadata)
 
@@ -167,9 +165,7 @@ def parse_basic_modify(
             ]
         )
 
-        relation_node_resource = RelationNode(
-            value=bm.uuid, ntype=RelationNode.NodeType.RESOURCE
-        )
+        relation_node_resource = RelationNode(value=bm.uuid, ntype=RelationNode.NodeType.RESOURCE)
         relations = []
         for relation in item.usermetadata.relations:
             if relation.from_ is None:
@@ -205,8 +201,16 @@ def parse_basic_modify(
         unique_groups = list(set(item.security.access_groups))
         bm.security.access_groups.extend(unique_groups)
 
+    if item.hidden is not None:
+        bm.basic.hidden = item.hidden
+
 
-def parse_basic(bm: BrokerMessage, item: CreateResourcePayload, toprocess: PushPayload):
+def parse_basic_creation(
+    bm: BrokerMessage,
+    item: CreateResourcePayload,
+    toprocess: PushPayload,
+    kb_config: Optional[KnowledgeBoxConfig],
+):
     bm.basic.created.FromDatetime(datetime.now())
 
     if item.title is None:
@@ -215,6 +219,10 @@ def parse_basic(bm: BrokerMessage, item: CreateResourcePayload, toprocess: PushP
 
     parse_basic_modify(bm, item, toprocess)
 
+    if item.hidden is None:
+        if kb_config and kb_config.hidden_resources_hide_on_creation:
+            bm.basic.hidden = True
+
 
 def set_status(basic: Basic, item: CreateResourcePayload):
     basic.metadata.status = Metadata.Status.PENDING
@@ -227,15 +235,11 @@ def set_status_modify(basic: Basic, item: UpdateResourcePayload):
 def validate_classifications(paragraph: ParagraphAnnotation):
     classifications = paragraph.classifications
     if len(classifications) == 0:
-        raise HTTPException(
-            status_code=422, detail="ensure classifications has at least 1 items"
-        )
+        raise HTTPException(status_code=422, detail="ensure classifications has at least 1 items")
 
-    unique_classifications = {tuple(cf.dict().values()) for cf in classifications}
+    unique_classifications = {tuple(cf.model_dump().values()) for cf in classifications}
     if len(unique_classifications) != len(classifications):
-        raise HTTPException(
-            status_code=422, detail="Paragraph classifications need to be unique"
-        )
+        raise HTTPException(status_code=422, detail="Paragraph classifications need to be unique")
 
 
 def compute_title(item: CreateResourcePayload, rid: str) -> str:
@@ -273,9 +277,7 @@ def build_question_answer_annotation_pb(
     pb.cancelled_by_user = qa_annotation.cancelled_by_user
     pb.question_answer.question.text = qa_annotation.question_answer.question.text
     if qa_annotation.question_answer.question.language is not None:
-        pb.question_answer.question.language = (
-            qa_annotation.question_answer.question.language
-        )
+        pb.question_answer.question.language = qa_annotation.question_answer.question.language
     pb.question_answer.question.ids_paragraphs.extend(
         qa_annotation.question_answer.question.ids_paragraphs
     )

nucliadb/writer/resource/field.py

@@ -21,23 +21,23 @@ from datetime import datetime
 from typing import Optional, Union
 
 from google.protobuf.json_format import MessageToDict
-from nucliadb_protos.writer_pb2 import BrokerMessage
 
 import nucliadb_models as models
+from nucliadb.common.models_utils import from_proto, to_proto
 from nucliadb.ingest.fields.conversation import Conversation
 from nucliadb.ingest.orm.resource import Resource as ORMResource
 from nucliadb.ingest.processing import PushPayload
 from nucliadb.writer import SERVICE_NAME
-from nucliadb.writer.layouts import serialize_blocks
 from nucliadb.writer.utilities import get_processing
-from nucliadb_models.common import FIELD_TYPES_MAP, FieldTypeName
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.content_types import GENERIC_MIME_TYPE
 from nucliadb_models.conversation import PushConversation
 from nucliadb_models.writer import (
-    GENERIC_MIME_TYPE,
     CreateResourcePayload,
     UpdateResourcePayload,
 )
 from nucliadb_protos import resources_pb2
+from nucliadb_protos.writer_pb2 import BrokerMessage
 from nucliadb_utils.storages.storage import StorageField
 from nucliadb_utils.utilities import get_storage
 
@@ -80,12 +80,11 @@ async def extract_fields(resource: ORMResource, toprocess: PushPayload):
     storage = await get_storage(service_name=SERVICE_NAME)
     await resource.get_fields()
     for (field_type, field_id), field in resource.fields.items():
-        field_type_name = FIELD_TYPES_MAP[field_type]
+        field_type_name = from_proto.field_type_name(field_type)
 
         if field_type_name not in {
             FieldTypeName.TEXT,
             FieldTypeName.FILE,
-            FieldTypeName.LAYOUT,
             FieldTypeName.CONVERSATION,
             FieldTypeName.LINK,
         }:
@@ -114,28 +113,7 @@ async def extract_fields(resource: ORMResource, toprocess: PushPayload):
             parsed_text["format"] = models.PushTextFormat[parsed_text["format"]]
             toprocess.textfield[field_id] = models.Text(**parsed_text)
 
-        if field_type_name is FieldTypeName.LAYOUT:
-            parsed_layout = MessageToDict(
-                field_pb,
-                preserving_proto_field_name=True,
-                including_default_value_fields=True,
-            )
-            parsed_layout["format"] = resources_pb2.FieldLayout.Format.Value(
-                parsed_layout["format"]
-            )
-
-            for blockid, block in parsed_layout["body"]["blocks"].items():
-                cf = field_pb.body.blocks[blockid].file
-                block["file"] = await processing.convert_internal_cf_to_str(cf, storage)
-
-            parsed_layout["blocks"] = parsed_layout.get("body", {}).get("blocks", {})
-            del parsed_layout["body"]
-
-            toprocess.layoutfield[field_id] = models.LayoutDiff(**parsed_layout)
-
-        if field_type_name is FieldTypeName.CONVERSATION and isinstance(
-            field, Conversation
-        ):
+        if field_type_name is FieldTypeName.CONVERSATION and isinstance(field, Conversation):
             metadata = await field.get_metadata()
             if metadata.pages == 0:
                 continue
@@ -156,14 +134,13 @@ async def extract_fields(resource: ORMResource, toprocess: PushPayload):
                         await processing.convert_internal_cf_to_str(cf, storage)
                         for cf in message.content.attachments
                     ]
-                    parsed_message["content"]["format"] = (
-                        resources_pb2.MessageContent.Format.Value(
-                            parsed_message["content"]["format"]
-                        )
-                    )
-                    full_conversation.messages.append(
-                        models.PushMessage(**parsed_message)
+                    if "attachments_fields" in parsed_message["content"]:
+                        # Not defined on the push payload
+                        del parsed_message["content"]["attachments_fields"]
+                    parsed_message["content"]["format"] = resources_pb2.MessageContent.Format.Value(
+                        parsed_message["content"]["format"]
                     )
+                    full_conversation.messages.append(models.PushMessage(**parsed_message))
             toprocess.conversationfield[field_id] = full_conversation
 
 
@@ -176,9 +153,7 @@ async def parse_fields(
     x_skip_store: bool,
 ):
     for key, file_field in item.files.items():
-        await parse_file_field(
-            key, file_field, writer, toprocess, kbid, uuid, skip_store=x_skip_store
-        )
+        await parse_file_field(key, file_field, writer, toprocess, kbid, uuid, skip_store=x_skip_store)
 
     for key, link_field in item.links.items():
         parse_link_field(key, link_field, writer, toprocess)
@@ -186,19 +161,8 @@ async def parse_fields(
     for key, text_field in item.texts.items():
         parse_text_field(key, text_field, writer, toprocess)
 
-    for key, layout_field in item.layouts.items():
-        await parse_layout_field(key, layout_field, writer, toprocess, kbid, uuid)
-
     for key, conversation_field in item.conversations.items():
-        await parse_conversation_field(
-            key, conversation_field, writer, toprocess, kbid, uuid
-        )
-
-    for key, datetime_field in item.datetimes.items():
-        parse_datetime_field(key, datetime_field, writer, toprocess)
-
-    for key, keywordset_field in item.keywordsets.items():
-        parse_keywordset_field(key, keywordset_field, writer, toprocess)
+        await parse_conversation_field(key, conversation_field, writer, toprocess, kbid, uuid)
 
 
 def parse_text_field(
@@ -208,9 +172,7 @@ def parse_text_field(
     toprocess: PushPayload,
 ) -> None:
     writer.texts[key].body = text_field.body
-    writer.texts[key].format = resources_pb2.FieldText.Format.Value(
-        text_field.format.value
-    )
+    writer.texts[key].format = resources_pb2.FieldText.Format.Value(text_field.format.value)
     etw = resources_pb2.ExtractedTextWrapper()
     etw.field.field = key
     etw.field.field_type = resources_pb2.FieldType.TEXT
@@ -338,78 +300,6 @@ def parse_link_field(
     )
 
 
-def parse_keywordset_field(
-    key: str,
-    keywordset_field: models.FieldKeywordset,
-    writer: BrokerMessage,
-    toprocess: PushPayload,
-) -> None:
-    if keywordset_field.keywords is None:
-        return
-
-    for keyword in keywordset_field.keywords:
-        fieldpb = resources_pb2.Keyword()
-        fieldpb.value = keyword.value
-        writer.keywordsets[key].keywords.append(fieldpb)
-
-
-def parse_datetime_field(
-    key: str,
-    datetime_field: models.FieldDatetime,
-    writer: BrokerMessage,
-    toprocess: PushPayload,
-) -> None:
-    if datetime_field.value is None:
-        return
-
-    writer.datetimes[key].value.FromDatetime(datetime_field.value)
-
-
-async def parse_layout_field(
-    key: str,
-    layout_field: models.InputLayoutField,
-    writer: BrokerMessage,
-    toprocess: PushPayload,
-    kbid: str,
-    uuid: str,
-) -> None:
-    storage = await get_storage(service_name=SERVICE_NAME)
-    processing = get_processing()
-
-    lc: resources_pb2.FieldLayout = await serialize_blocks(
-        layout_field, kbid, uuid, key, storage
-    )
-    writer.layouts[key].CopyFrom(lc)
-
-    toprocess_blocks = {}
-    for blockid, block in layout_field.body.blocks.items():
-        sf_conv_field: StorageField = storage.layout_field(
-            kbid, uuid, field=key, ident=block.ident
-        )
-        cf_conv_field = await storage.upload_b64file_to_cloudfile(
-            sf_conv_field,
-            block.file.payload.encode(),
-            block.file.filename,
-            block.file.content_type,
-            block.file.md5,
-        )
-
-        toprocess_blocks[blockid] = models.PushLayoutBlock(
-            x=block.x,
-            y=block.y,
-            cols=block.cols,
-            rows=block.rows,
-            type=block.type,
-            ident=block.ident,
-            payload=block.payload,
-            file=await processing.convert_internal_cf_to_str(cf_conv_field, storage),
-        )
-
-    toprocess.layoutfield[key] = models.LayoutDiff(
-        format=lc.format, blocks=toprocess_blocks  # type: ignore
-    )
-
-
 async def parse_conversation_field(
     key: str,
     conversation_field: models.InputConversationField,
@@ -441,8 +331,16 @@ async def parse_conversation_field(
         )
 
         cm.content.text = message.content.text
-        cm.content.format = resources_pb2.MessageContent.Format.Value(
-            message.content.format.value
+        cm.content.format = resources_pb2.MessageContent.Format.Value(message.content.format.value)
+        cm.content.attachments_fields.extend(
+            [
+                resources_pb2.FieldRef(
+                    field_type=to_proto.field_type_name(attachment.field_type),
+                    field_id=attachment.field_id,
+                    split=attachment.split if attachment.split is not None else "",
+                )
+                for attachment in message.content.attachments_fields
+            ]
         )
 
         for count, file in enumerate(message.content.attachments):

nucliadb/writer/resource/origin.py

@@ -16,10 +16,9 @@
 #
 # 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 Origin
-
 from nucliadb_models import Extra, InputOrigin
 from nucliadb_protos import resources_pb2
+from nucliadb_protos.resources_pb2 import Origin
 
 
 def parse_origin(origin: Origin, origin_payload: InputOrigin):

nucliadb/writer/settings.py

@@ -48,16 +48,20 @@ class BackPressureSettings(BaseSettings):
         description="Estimation of the processing rate in messages per second. This is used to calculate the try again in time",  # noqa
     )
     max_indexing_pending: int = Field(
-        default=100,
+        default=200,
         description="Max number of messages pending to index in a node queue before rate limiting writes. Set to 0 to disable indexing back pressure checks",  # noqa
+        alias="back_pressure_max_indexing_pending",
     )
     max_ingest_pending: int = Field(
-        default=1000,
+        # Disabled by default
+        default=0,
         description="Max number of messages pending to be ingested by processed consumers before rate limiting writes. Set to 0 to disable ingest back pressure checks",  # noqa
+        alias="back_pressure_max_ingest_pending",
     )
     max_processing_pending: int = Field(
         default=1000,
         description="Max number of messages pending to process per Knowledge Box before rate limiting writes. Set to 0 to disable processing back pressure checks",  # noqa
+        alias="back_pressure_max_processing_pending",
     )
     indexing_check_interval: int = Field(
         default=30,

nucliadb/writer/tus/__init__.py

@@ -23,10 +23,6 @@ from typing import Optional
 from nucliadb.writer.settings import settings as writer_settings
 from nucliadb.writer.tus.dm import FileDataManager, RedisFileDataManagerFactory
 from nucliadb.writer.tus.exceptions import ManagerNotAvailable
-from nucliadb.writer.tus.gcs import GCloudBlobStore, GCloudFileStorageManager
-from nucliadb.writer.tus.local import LocalBlobStore, LocalFileStorageManager
-from nucliadb.writer.tus.pg import PGBlobStore, PGFileStorageManager
-from nucliadb.writer.tus.s3 import S3BlobStore, S3FileStorageManager
 from nucliadb.writer.tus.storage import BlobStore, FileStorageManager
 from nucliadb_utils.exceptions import ConfigurationError
 from nucliadb_utils.settings import FileBackendConfig, storage_settings
@@ -48,6 +44,8 @@ REDIS_FILE_DATA_MANAGER_FACTORY: Optional[RedisFileDataManagerFactory] = None
 async def initialize():
     global DRIVER
     if storage_settings.file_backend == FileBackendConfig.GCS:
+        from nucliadb.writer.tus.gcs import GCloudBlobStore, GCloudFileStorageManager
+
         storage_backend = GCloudBlobStore()
 
         await storage_backend.initialize(
@@ -64,6 +62,8 @@ async def initialize():
         DRIVER = TusStorageDriver(backend=storage_backend, manager=storage_manager)
 
     elif storage_settings.file_backend == FileBackendConfig.S3:
+        from nucliadb.writer.tus.s3 import S3BlobStore, S3FileStorageManager
+
         storage_backend = S3BlobStore()
 
         await storage_backend.initialize(
@@ -83,6 +83,8 @@ async def initialize():
         DRIVER = TusStorageDriver(backend=storage_backend, manager=storage_manager)
 
     elif storage_settings.file_backend == FileBackendConfig.LOCAL:
+        from nucliadb.writer.tus.local import LocalBlobStore, LocalFileStorageManager
+
         storage_backend = LocalBlobStore(storage_settings.local_files)
 
         await storage_backend.initialize()
@@ -91,12 +93,18 @@ async def initialize():
 
         DRIVER = TusStorageDriver(backend=storage_backend, manager=storage_manager)
 
-    elif storage_settings.file_backend == FileBackendConfig.PG:
-        storage_backend = PGBlobStore(storage_settings.driver_pg_url)
+    elif storage_settings.file_backend == FileBackendConfig.AZURE:
+        from nucliadb.writer.tus.azure import AzureBlobStore, AzureFileStorageManager
 
-        await storage_backend.initialize()
+        if storage_settings.azure_account_url is None:
+            raise ConfigurationError("AZURE_ACCOUNT_URL env variable not configured")
 
-        storage_manager = PGFileStorageManager(storage_backend)
+        storage_backend = AzureBlobStore()
+        await storage_backend.initialize(
+            storage_settings.azure_account_url,
+            connection_string=storage_settings.azure_connection_string,
+        )
+        storage_manager = AzureFileStorageManager(storage_backend)
 
         DRIVER = TusStorageDriver(backend=storage_backend, manager=storage_manager)
 
@@ -117,7 +125,7 @@ async def finalize():
         REDIS_FILE_DATA_MANAGER_FACTORY = None
 
 
-def get_dm() -> FileDataManager:  # type: ignore
+def get_dm() -> FileDataManager:
     if writer_settings.dm_enabled:
         global REDIS_FILE_DATA_MANAGER_FACTORY
         if REDIS_FILE_DATA_MANAGER_FACTORY is None:
@@ -136,9 +144,3 @@ def get_storage_manager() -> FileStorageManager:
     if DRIVER is None:
         raise ManagerNotAvailable()
     return DRIVER.manager
-
-
-def clear_storage():
-    global DRIVER
-
-    DRIVER = None