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

nucliadb/writer/resource/field.py

@@ -21,36 +21,48 @@ 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
 
 
+async def extract_file_field_from_pb(field_pb: resources_pb2.FieldFile) -> str:
+    processing = get_processing()
+
+    if field_pb.file.source == resources_pb2.CloudFile.Source.EXTERNAL:
+        file_field = models.FileField(
+            language=field_pb.language,
+            password=field_pb.password,
+            file=models.File(payload=None, uri=field_pb.file.uri),
+        )
+        return processing.convert_external_filefield_to_str(file_field)
+    else:
+        storage = await get_storage(service_name=SERVICE_NAME)
+        return await processing.convert_internal_filefield_to_str(field_pb, storage)
+
+
 async def extract_file_field(
     field_id: str,
     resource: ORMResource,
     toprocess: PushPayload,
     password: Optional[str] = None,
 ):
-    processing = get_processing()
-    storage = await get_storage(service_name=SERVICE_NAME)
-
     field_type = resources_pb2.FieldType.FILE
     field = await resource.get_field(field_id, field_type)
     field_pb = await field.get_value()
@@ -60,9 +72,7 @@ async def extract_file_field(
     if password is not None:
         field_pb.password = password
 
-    toprocess.filefield[field_id] = await processing.convert_internal_filefield_to_str(
-        field_pb, storage
-    )
+    toprocess.filefield[field_id] = await extract_file_field_from_pb(field_pb)
 
 
 async def extract_fields(resource: ORMResource, toprocess: PushPayload):
@@ -70,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,
         }:
@@ -84,9 +93,7 @@ async def extract_fields(resource: ORMResource, toprocess: PushPayload):
         field_pb = await field.get_value()
 
         if field_type_name is FieldTypeName.FILE:
-            toprocess.filefield[
-                field_id
-            ] = await processing.convert_internal_filefield_to_str(field_pb, storage)
+            toprocess.filefield[field_id] = await extract_file_field_from_pb(field_pb)
 
         if field_type_name is FieldTypeName.LINK:
             parsed_link = MessageToDict(
@@ -106,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
@@ -148,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(
+                    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)
-                    )
+                    full_conversation.messages.append(models.PushMessage(**parsed_message))
             toprocess.conversationfield[field_id] = full_conversation
 
 
@@ -168,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)
@@ -178,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(
@@ -200,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
@@ -317,84 +287,16 @@ def parse_link_field(
     if link_field.css_selector is not None:
         writer.links[key].css_selector = link_field.css_selector
 
+    if link_field.xpath is not None:
+        writer.links[key].xpath = link_field.xpath
+
     toprocess.linkfield[key] = models.LinkUpload(
         link=link_field.uri,
         headers=link_field.headers or {},
         cookies=link_field.cookies or {},
         localstorage=link_field.localstorage or {},
         css_selector=link_field.css_selector,
-    )
-
-
-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
+        xpath=link_field.xpath,
     )
 
 
@@ -429,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

@@ -19,7 +19,8 @@
 #
 from typing import Optional
 
-from pydantic import BaseSettings, Field
+from pydantic import Field
+from pydantic_settings import BaseSettings
 
 
 class Settings(BaseSettings):
@@ -32,10 +33,10 @@ class BackPressureSettings(BaseSettings):
     enabled: bool = Field(
         default=False,
         description="Enable or disable back pressure.",
-        env=["back_pressure_enabled"],
+        alias="back_pressure_enabled",
     )
     indexing_rate: float = Field(
-        default=2,
+        default=4,
         description="Estimation of the indexing rate in messages per second. This is used to calculate the try again in time",  # noqa
     )
     ingest_rate: float = Field(
@@ -47,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=1_000,
+        # 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,
@@ -66,6 +71,10 @@ class BackPressureSettings(BaseSettings):
         default=30,
         description="Interval in seconds to check the ingest pending messages",
     )
+    max_wait_time: int = Field(
+        default=60,
+        description="Max time in seconds to wait before trying again after back pressure",
+    )
 
 
 settings = Settings()

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

nucliadb/writer/tus/azure.py

@@ -0,0 +1,111 @@
+# 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 __future__ import annotations
+
+from typing import Optional
+
+from nucliadb.writer import logger
+from nucliadb.writer.tus.dm import FileDataManager
+from nucliadb.writer.tus.storage import BlobStore, FileStorageManager
+from nucliadb_protos.resources_pb2 import CloudFile
+from nucliadb_utils.storages import CHUNK_SIZE
+from nucliadb_utils.storages.azure import AzureObjectStore
+from nucliadb_utils.storages.exceptions import ObjectNotFoundError
+from nucliadb_utils.storages.utils import ObjectMetadata
+
+
+class AzureBlobStore(BlobStore):
+    async def finalize(self):
+        if self._object_store is None:
+            return
+        try:
+            await self._object_store.finalize()
+        except Exception:
+            logger.exception("Error closing AzureBlobStore")
+        self._object_store = None
+
+    async def initialize(self, account_url: str, connection_string: Optional[str] = None):
+        self.bucket = "nucliadb-{kbid}"
+        self.source = CloudFile.Source.AZURE
+        self._object_store = AzureObjectStore(account_url, connection_string=connection_string)
+        await self._object_store.initialize()
+
+    @property
+    def object_store(self) -> AzureObjectStore:
+        assert self._object_store is not None
+        return self._object_store
+
+    async def check_exists(self, bucket_name: str) -> bool:
+        return await self.object_store.bucket_exists(bucket_name)
+
+    async def create_bucket(self, bucket_name: str) -> bool:
+        created = await self.object_store.bucket_create(bucket_name)
+        return not created
+
+
+class AzureFileStorageManager(FileStorageManager):
+    storage: AzureBlobStore
+    chunk_size = CHUNK_SIZE
+    min_upload_size = None
+
+    @property
+    def object_store(self) -> AzureObjectStore:
+        return self.storage.object_store
+
+    async def start(self, dm: FileDataManager, path: str, kbid: str):
+        bucket = self.storage.get_bucket_name(kbid)
+        if dm.filename == 0:
+            filename = "file"
+        else:
+            filename = dm.filename
+        metadata = ObjectMetadata(
+            filename=filename,
+            content_type=dm.content_type,
+            size=dm.size,
+        )
+        await self.object_store.upload_multipart_start(bucket, path, metadata)
+        await dm.update(path=path, bucket=bucket)
+
+    async def delete_upload(self, uri: str, kbid: str) -> None:
+        bucket = self.storage.get_bucket_name(kbid)
+        try:
+            await self.object_store.delete(bucket, uri)
+        except ObjectNotFoundError:
+            logger.warning(
+                "Attempt to delete an upload but not found",
+                extra={"uri": uri, "kbid": kbid, "bucket": bucket},
+            )
+
+    async def append(self, dm: FileDataManager, iterable, offset: int) -> int:
+        bucket = dm.get("bucket")
+        assert bucket is not None
+        path = dm.get("path")
+        assert path is not None
+        uploaded_bytes = await self.object_store.upload_multipart_append(bucket, path, iterable)
+        await dm.update(offset=offset)
+        return uploaded_bytes
+
+    async def finish(self, dm: FileDataManager):
+        path = dm.get("path")
+        await dm.finish()
+        return path
+
+    def validate_intermediate_chunk(self, uploaded_bytes: int):
+        pass

nucliadb/writer/tus/dm.py

@@ -20,6 +20,7 @@
 import time
 from typing import Any, Optional
 
+import backoff
 import orjson
 from redis import asyncio as aioredis
 from starlette.requests import Request
@@ -33,6 +34,11 @@ class NoRedisConfigured(Exception):
     pass
 
 
+RETRIABLE_REDIS_ERRORS = (
+    aioredis.ConnectionError,
+    aioredis.TimeoutError,
+)
+
 DATA: dict[str, Any] = {}
 
 
@@ -59,10 +65,7 @@ class FileDataManager:
             # someone else
             last_activity: Optional[int] = self._data.get("last_activity")
             if last_activity and (time.time() - last_activity) < self._ttl:
-                if (
-                    request.headers
-                    and request.headers.get("tus-override-upload", "0") != "1"
-                ):
+                if request.headers and request.headers.get("tus-override-upload", "0") != "1":
                     raise HTTPPreconditionFailed(
                         detail="There is already an active tusupload that conflicts with this one."
                     )
@@ -136,7 +139,7 @@ class RedisFileDataManagerFactory:
 
     async def finalize(self):
         try:
-            await self.redis.close(close_connection_pool=True)
+            await self.redis.aclose(close_connection_pool=True)
         except Exception:
             logger.warning("Error closing redis connection", exc_info=True)
             pass
@@ -146,6 +149,9 @@ class RedisFileDataManager(FileDataManager):
     def __init__(self, redis: aioredis.Redis):
         self.redis = redis
 
+    @backoff.on_exception(
+        backoff.expo, RETRIABLE_REDIS_ERRORS, jitter=backoff.random_jitter, max_tries=3
+    )
     async def load(self, key):
         # preload data
         self.key = key
@@ -157,6 +163,9 @@ class RedisFileDataManager(FileDataManager):
                 self._data = orjson.loads(data)
                 self._loaded = True
 
+    @backoff.on_exception(
+        backoff.expo, RETRIABLE_REDIS_ERRORS, jitter=backoff.random_jitter, max_tries=3
+    )
     async def save(self):
         if self.key is None:
             raise Exception("Not initialized")
@@ -164,6 +173,9 @@ class RedisFileDataManager(FileDataManager):
         value = orjson.dumps(self._data)
         await self.redis.set(self.key, value, ex=self._ttl)
 
+    @backoff.on_exception(
+        backoff.expo, RETRIABLE_REDIS_ERRORS, jitter=backoff.random_jitter, max_tries=3
+    )
     async def _delete_key(self):
         if self.key is None:
             raise Exception("Not initialized")

nucliadb/writer/tus/exceptions.py

@@ -31,9 +31,7 @@ class HTTPException(StarletteHTTPException):
 
     def __init__(self, detail: Optional[str] = None):
         if self._status_code:
-            super(HTTPException, self).__init__(
-                status_code=self._status_code, detail=detail
-            )
+            super(HTTPException, self).__init__(status_code=self._status_code, detail=detail)
         else:
             raise AttributeError("Status code not defined")