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

nucliadb/writer/api/v1/export_import.py

@@ -17,6 +17,7 @@
 # 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 datetime import datetime
 from typing import AsyncGenerator
 from uuid import uuid4
 
@@ -29,24 +30,32 @@ from nucliadb.common.context import ApplicationContext
 from nucliadb.common.context.fastapi import get_app_context
 from nucliadb.export_import import importer
 from nucliadb.export_import.datamanager import ExportImportDataManager
-from nucliadb.export_import.exceptions import IncompatibleExport
+from nucliadb.export_import.exceptions import (
+    IncompatibleExport,
+)
 from nucliadb.export_import.models import (
     ExportMetadata,
     ImportMetadata,
     NatsTaskMessage,
 )
 from nucliadb.export_import.tasks import get_exports_producer, get_imports_producer
-from nucliadb.export_import.utils import stream_compatible_with_kb
+from nucliadb.export_import.utils import ExportStreamReader, stream_compatible_with_kb
 from nucliadb.models.responses import HTTPClientError
 from nucliadb.writer import logger
-from nucliadb.writer.api.v1.router import KB_PREFIX, api
+from nucliadb.writer.api.utils import only_for_onprem
+from nucliadb.writer.api.v1.knowledgebox import create_kb
+from nucliadb.writer.api.v1.router import KB_PREFIX, KBS_PREFIX, api
 from nucliadb.writer.back_pressure import maybe_back_pressure
 from nucliadb_models.export_import import (
     CreateExportResponse,
     CreateImportResponse,
+    NewImportedKbResponse,
     Status,
 )
-from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_models.resource import (
+    KnowledgeBoxConfig,
+    NucliaDBRoles,
+)
 from nucliadb_telemetry import errors
 from nucliadb_utils.authentication import requires_one
 
@@ -75,6 +84,56 @@ async def start_kb_export_endpoint(request: Request, kbid: str):
         return CreateExportResponse(export_id=export_id)
 
 
+@only_for_onprem
+@api.post(
+    f"/{KBS_PREFIX}/import",
+    summary="Create a KB from an export and import its content",
+    tags=["Knowledge Boxes"],
+    response_model=NewImportedKbResponse,
+    openapi_extra={"x-hidden-operation": True},
+)
+@requires_one([NucliaDBRoles.MANAGER, NucliaDBRoles.WRITER])
+@version(1)
+async def kb_create_and_import_endpoint(request: Request):
+    context = get_app_context(request.app)
+
+    # Read stream and parse learning configuration
+    stream = request.stream()
+    stream_reader = ExportStreamReader(stream)
+    learning_config, leftover_bytes = await stream_reader.maybe_read_learning_config()
+    if learning_config is None:
+        return HTTPClientError(
+            status_code=400,
+            detail="Trying to import an export missing learning config. Try using import on an existing KB or use a newer export",
+        )
+
+    # Create a KB with the import learning config
+
+    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    import_kb_config = KnowledgeBoxConfig(
+        title=f"Imported KB - {now}",
+        learning_configuration=learning_config.dict(),
+    )
+    kbid, slug = await create_kb(import_kb_config)
+
+    # Import contents to the new KB
+
+    async def stream_with_leftovers(leftovers: bytes, stream: AsyncGenerator[bytes, None]):
+        if len(leftovers) > 0:
+            yield leftovers
+        async for chunk in stream:
+            yield chunk
+
+    await importer.import_kb(
+        context=context, kbid=kbid, stream=stream_with_leftovers(leftover_bytes, stream)
+    )
+
+    return NewImportedKbResponse(
+        kbid=kbid,
+        slug=slug,
+    )
+
+
 @api.post(
     f"/{KB_PREFIX}/{{kbid}}/import",
     status_code=200,
@@ -139,30 +198,24 @@ async def start_export_task(context: ApplicationContext, kbid: str, export_id: s
         producer = await get_exports_producer(context)
         msg = NatsTaskMessage(kbid=kbid, id=export_id)
         seqid = await producer(msg)  # type: ignore
-        logger.info(
-            f"Export task produced. seqid={seqid} kbid={kbid} export_id={export_id}"
-        )
+        logger.info(f"Export task produced. seqid={seqid} kbid={kbid} export_id={export_id}")
     except Exception as e:
         errors.capture_exception(e)
         await dm.delete_metadata("export", metadata)
         raise
 
 
-async def start_import_task(
-    context: ApplicationContext, kbid: str, import_id: str, import_size: int
-):
+async def start_import_task(context: ApplicationContext, kbid: str, import_id: str, import_size: int):
     dm = ExportImportDataManager(context.kv_driver, context.blob_storage)
     metadata = ImportMetadata(kbid=kbid, id=import_id)
     metadata.task.status = Status.SCHEDULED
-    metadata.total = import_size
+    metadata.total = import_size or 0
     await dm.set_metadata("import", metadata)
     try:
         producer = await get_imports_producer(context)
         msg = NatsTaskMessage(kbid=kbid, id=import_id)
         seqid = await producer(msg)  # type: ignore
-        logger.info(
-            f"Import task produced. seqid={seqid} kbid={kbid} import_id={import_id}"
-        )
+        logger.info(f"Import task produced. seqid={seqid} kbid={kbid} import_id={import_id}")
     except Exception as e:
         errors.capture_exception(e)
         await dm.delete_metadata("import", metadata)

nucliadb/writer/api/v1/field.py

@@ -21,9 +21,7 @@ from inspect import iscoroutinefunction
 from typing import TYPE_CHECKING, Callable, Optional, Type, Union
 
 from fastapi import HTTPException, Response
-from fastapi_versioning import version  # type: ignore
-from nucliadb_protos.resources_pb2 import FieldID, Metadata
-from nucliadb_protos.writer_pb2 import BrokerMessage
+from fastapi_versioning import version
 from starlette.requests import Request
 
 import nucliadb_models as models
@@ -36,6 +34,7 @@ from nucliadb.writer.api.constants import (
     X_FILE_PASSWORD,
     X_NUCLIADB_USER,
 )
+from nucliadb.writer.api.v1 import transaction
 from nucliadb.writer.api.v1.resource import (
     get_rid_from_slug_or_raise_error,
     validate_rid_exists_or_raise_error,
@@ -46,10 +45,7 @@ from nucliadb.writer.resource.audit import parse_audit
 from nucliadb.writer.resource.field import (
     extract_file_field,
     parse_conversation_field,
-    parse_datetime_field,
     parse_file_field,
-    parse_keywordset_field,
-    parse_layout_field,
     parse_link_field,
     parse_text_field,
 )
@@ -58,28 +54,23 @@ from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_models.utils import FieldIdString
 from nucliadb_models.writer import ResourceFieldAdded, ResourceUpdated
 from nucliadb_protos import resources_pb2
+from nucliadb_protos.resources_pb2 import FieldID, Metadata
+from nucliadb_protos.writer_pb2 import BrokerMessage
 from nucliadb_utils.authentication import requires
 from nucliadb_utils.exceptions import LimitsExceededError, SendToProcessError
-from nucliadb_utils.transaction import TransactionCommitTimeoutError
 from nucliadb_utils.utilities import (
     get_partitioning,
     get_storage,
-    get_transaction_utility,
 )
 
 if TYPE_CHECKING:  # pragma: no cover
-    FIELD_TYPE_NAME_TO_FIELD_TYPE_MAP: dict[
-        models.FieldTypeName, resources_pb2.FieldType.V
-    ]
+    FIELD_TYPE_NAME_TO_FIELD_TYPE_MAP: dict[models.FieldTypeName, resources_pb2.FieldType.ValueType]
 else:
     FIELD_TYPE_NAME_TO_FIELD_TYPE_MAP: dict[models.FieldTypeName, int]
 
 FieldModelType = Union[
     models.TextField,
     models.LinkField,
-    models.FieldKeywordset,
-    models.FieldDatetime,
-    models.InputLayoutField,
     models.InputConversationField,
     models.FileField,
 ]
@@ -87,10 +78,7 @@ FieldModelType = Union[
 FIELD_TYPE_NAME_TO_FIELD_TYPE_MAP = {
     models.FieldTypeName.FILE: resources_pb2.FieldType.FILE,
     models.FieldTypeName.LINK: resources_pb2.FieldType.LINK,
-    models.FieldTypeName.DATETIME: resources_pb2.FieldType.DATETIME,
-    models.FieldTypeName.KEYWORDSET: resources_pb2.FieldType.KEYWORDSET,
     models.FieldTypeName.TEXT: resources_pb2.FieldType.TEXT,
-    models.FieldTypeName.LAYOUT: resources_pb2.FieldType.LAYOUT,
     # models.FieldTypeName.GENERIC: resources_pb2.FieldType.GENERIC,
     models.FieldTypeName.CONVERSATION: resources_pb2.FieldType.CONVERSATION,
 }
@@ -128,24 +116,12 @@ async def add_field_to_resource(
 
     parse_field = FIELD_PARSERS_MAP[type(field_payload)]
     if iscoroutinefunction(parse_field):
-        await parse_field(
-            kbid, rid, field_id, field_payload, writer, toprocess, **parser_kwargs
-        )
+        await parse_field(kbid, rid, field_id, field_payload, writer, toprocess, **parser_kwargs)
     else:
-        parse_field(
-            kbid, rid, field_id, field_payload, writer, toprocess, **parser_kwargs
-        )
+        parse_field(kbid, rid, field_id, field_payload, writer, toprocess, **parser_kwargs)
 
-    transaction = get_transaction_utility()
     processing = get_processing()
-    try:
-        writer.source = BrokerMessage.MessageSource.WRITER
-        await transaction.commit(writer, partition, wait=True)
-    except TransactionCommitTimeoutError:
-        raise HTTPException(
-            status_code=501,
-            detail="Inconsistent write. This resource will not be processed and may not be stored.",
-        )
+    await transaction.commit(writer, partition)
     try:
         processing_info = await processing.send_to_process(toprocess, partition)
     except LimitsExceededError as exc:
@@ -168,9 +144,7 @@ async def add_field_to_resource_by_slug(
     **parser_kwargs,
 ):
     rid = await get_rid_from_slug_or_raise_error(kbid, slug)
-    return await add_field_to_resource(
-        request, kbid, rid, field_id, field_payload, **parser_kwargs
-    )
+    return await add_field_to_resource(request, kbid, rid, field_id, field_payload, **parser_kwargs)
 
 
 async def delete_resource_field(
@@ -183,8 +157,6 @@ async def delete_resource_field(
     await validate_rid_exists_or_raise_error(kbid, rid)
 
     partitioning = get_partitioning()
-    transaction = get_transaction_utility()
-
     partition = partitioning.generate_partition(kbid, rid)
     writer = BrokerMessage()
 
@@ -197,15 +169,7 @@ async def delete_resource_field(
 
     writer.delete_fields.append(pb_field_id)
     parse_audit(writer.audit, request)
-
-    try:
-        await transaction.commit(writer, partition, wait=True)
-    except TransactionCommitTimeoutError:
-        raise HTTPException(
-            status_code=501,
-            detail="Inconsistent write. This resource will not be processed and may not be stored.",
-        )
-
+    await transaction.commit(writer, partition)
     return Response(status_code=204)
 
 
@@ -251,41 +215,6 @@ def parse_link_field_adapter(
     return parse_link_field(field_id, field_payload, writer, toprocess)
 
 
-def parse_keywordset_field_adapter(
-    _kbid: str,
-    _rid: str,
-    field_id: FieldIdString,
-    field_payload: models.FieldKeywordset,
-    writer: BrokerMessage,
-    toprocess: PushPayload,
-):
-    return parse_keywordset_field(field_id, field_payload, writer, toprocess)
-
-
-def parse_datetime_field_adapter(
-    _kbid: str,
-    _rid: str,
-    field_id: FieldIdString,
-    field_payload: models.FieldDatetime,
-    writer: BrokerMessage,
-    toprocess: PushPayload,
-):
-    return parse_datetime_field(field_id, field_payload, writer, toprocess)
-
-
-async def parse_layout_field_adapter(
-    kbid: str,
-    rid: str,
-    field_id: FieldIdString,
-    field_payload: models.InputLayoutField,
-    writer: BrokerMessage,
-    toprocess: PushPayload,
-):
-    return await parse_layout_field(
-        field_id, field_payload, writer, toprocess, kbid, rid
-    )
-
-
 async def parse_conversation_field_adapter(
     kbid: str,
     rid: str,
@@ -294,9 +223,7 @@ async def parse_conversation_field_adapter(
     writer: BrokerMessage,
     toprocess: PushPayload,
 ):
-    return await parse_conversation_field(
-        field_id, field_payload, writer, toprocess, kbid, rid
-    )
+    return await parse_conversation_field(field_id, field_payload, writer, toprocess, kbid, rid)
 
 
 async def parse_file_field_adapter(
@@ -316,9 +243,6 @@ async def parse_file_field_adapter(
 FIELD_PARSERS_MAP: dict[Type, Union[Callable]] = {
     models.TextField: parse_text_field_adapter,
     models.LinkField: parse_link_field_adapter,
-    models.FieldKeywordset: parse_keywordset_field_adapter,
-    models.FieldDatetime: parse_datetime_field_adapter,
-    models.InputLayoutField: parse_layout_field_adapter,
     models.InputConversationField: parse_conversation_field_adapter,
     models.FileField: parse_file_field_adapter,
 }
@@ -343,9 +267,7 @@ async def add_resource_field_text_rslug_prefix(
     field_id: FieldIdString,
     field_payload: models.TextField,
 ) -> ResourceFieldAdded:
-    return await add_field_to_resource_by_slug(
-        request, kbid, rslug, field_id, field_payload
-    )
+    return await add_field_to_resource_by_slug(request, kbid, rslug, field_id, field_payload)
 
 
 @api.put(
@@ -383,9 +305,7 @@ async def add_resource_field_link_rslug_prefix(
     field_id: FieldIdString,
     field_payload: models.LinkField,
 ) -> ResourceFieldAdded:
-    return await add_field_to_resource_by_slug(
-        request, kbid, rslug, field_id, field_payload
-    )
+    return await add_field_to_resource_by_slug(request, kbid, rslug, field_id, field_payload)
 
 
 @api.put(
@@ -407,126 +327,6 @@ async def add_resource_field_link_rid_prefix(
     return await add_field_to_resource(request, kbid, rid, field_id, field_payload)
 
 
-@api.put(
-    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/keywordset/{{field_id}}",
-    status_code=201,
-    summary="Add resource keywordset field (by slug)",
-    response_model=ResourceFieldAdded,
-    tags=["Resource fields"],
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def add_resource_field_keywordset_rslug_prefix(
-    request: Request,
-    kbid: str,
-    rslug: str,
-    field_id: FieldIdString,
-    field_payload: models.FieldKeywordset,
-) -> ResourceFieldAdded:
-    return await add_field_to_resource_by_slug(
-        request, kbid, rslug, field_id, field_payload
-    )
-
-
-@api.put(
-    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/keywordset/{{field_id}}",
-    status_code=201,
-    summary="Add resource keywordset field (by id)",
-    response_model=ResourceFieldAdded,
-    tags=["Resource fields"],
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def add_resource_field_keywordset_rid_prefix(
-    request: Request,
-    kbid: str,
-    rid: str,
-    field_id: FieldIdString,
-    field_payload: models.FieldKeywordset,
-) -> ResourceFieldAdded:
-    return await add_field_to_resource(request, kbid, rid, field_id, field_payload)
-
-
-@api.put(
-    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/datetime/{{field_id}}",
-    status_code=201,
-    summary="Add resource datetime field (by slug)",
-    response_model=ResourceFieldAdded,
-    tags=["Resource fields"],
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def add_resource_field_datetime_rslug_prefix(
-    request: Request,
-    kbid: str,
-    rslug: str,
-    field_id: FieldIdString,
-    field_payload: models.FieldDatetime,
-) -> ResourceFieldAdded:
-    return await add_field_to_resource_by_slug(
-        request, kbid, rslug, field_id, field_payload
-    )
-
-
-@api.put(
-    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/datetime/{{field_id}}",
-    status_code=201,
-    summary="Add resource datetime field (by id)",
-    response_model=ResourceFieldAdded,
-    tags=["Resource fields"],
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def add_resource_field_datetime_rid_prefix(
-    request: Request,
-    kbid: str,
-    rid: str,
-    field_id: FieldIdString,
-    field_payload: models.FieldDatetime,
-) -> ResourceFieldAdded:
-    return await add_field_to_resource(request, kbid, rid, field_id, field_payload)
-
-
-@api.put(
-    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/layout/{{field_id}}",
-    status_code=201,
-    summary="Add resource layout field (by slug)",
-    response_model=ResourceFieldAdded,
-    tags=["Resource fields"],
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def add_resource_field_layout_rslug_prefix(
-    request: Request,
-    kbid: str,
-    rslug: str,
-    field_id: FieldIdString,
-    field_payload: models.InputLayoutField,
-) -> ResourceFieldAdded:
-    return await add_field_to_resource_by_slug(
-        request, kbid, rslug, field_id, field_payload
-    )
-
-
-@api.put(
-    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/layout/{{field_id}}",
-    status_code=201,
-    summary="Add resource layout field (by id)",
-    response_model=ResourceFieldAdded,
-    tags=["Resource fields"],
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def add_resource_field_layout_rid_prefix(
-    request: Request,
-    kbid: str,
-    rid: str,
-    field_id: FieldIdString,
-    field_payload: models.InputLayoutField,
-) -> ResourceFieldAdded:
-    return await add_field_to_resource(request, kbid, rid, field_id, field_payload)
-
-
 @api.put(
     f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/conversation/{{field_id}}",
     status_code=201,
@@ -543,9 +343,7 @@ async def add_resource_field_conversation_rslug_prefix(
     field_id: FieldIdString,
     field_payload: models.InputConversationField,
 ) -> ResourceFieldAdded:
-    return await add_field_to_resource_by_slug(
-        request, kbid, rslug, field_id, field_payload
-    )
+    return await add_field_to_resource_by_slug(request, kbid, rslug, field_id, field_payload)
 
 
 @api.put(
@@ -653,48 +451,6 @@ async def append_messages_to_conversation_field_rid_prefix(
     return await add_field_to_resource(request, kbid, rid, field_id, field)
 
 
-@api.put(
-    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/layout/{{field_id}}/blocks",
-    status_code=200,
-    summary="Append blocks to layout field (by slug)",
-    response_model=ResourceFieldAdded,
-    tags=["Resource fields"],
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def append_blocks_to_layout_field_rslug_prefix(
-    request: Request,
-    kbid: str,
-    rslug: str,
-    field_id: FieldIdString,
-    blocks: dict[str, models.InputBlock],
-) -> ResourceFieldAdded:
-    field = models.InputLayoutField(body=models.InputLayoutContent())
-    field.body.blocks.update(blocks)
-    return await add_field_to_resource_by_slug(request, kbid, rslug, field_id, field)
-
-
-@api.put(
-    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/layout/{{field_id}}/blocks",
-    status_code=200,
-    summary="Append blocks to layout field (by id)",
-    response_model=ResourceFieldAdded,
-    tags=["Resource fields"],
-)
-@requires(NucliaDBRoles.WRITER)
-@version(1)
-async def append_blocks_to_layout_field_rid_prefix(
-    request: Request,
-    kbid: str,
-    rid: str,
-    field_id: FieldIdString,
-    blocks: dict[str, models.InputBlock],
-) -> ResourceFieldAdded:
-    field = models.InputLayoutField(body=models.InputLayoutContent())
-    field.body.blocks.update(blocks)
-    return await add_field_to_resource(request, kbid, rid, field_id, field)
-
-
 @api.delete(
     f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/{{field_type}}/{{field_id}}",
     status_code=204,
@@ -711,9 +467,7 @@ async def delete_resource_field_rslug_prefix(
     field_type: models.FieldTypeName,
     field_id: FieldIdString,
 ):
-    return await delete_resource_field_by_slug(
-        request, kbid, rslug, field_type, field_id
-    )
+    return await delete_resource_field_by_slug(request, kbid, rslug, field_type, field_id)
 
 
 @api.delete(
@@ -754,7 +508,6 @@ async def reprocess_file_field(
 ) -> ResourceUpdated:
     await maybe_back_pressure(request, kbid, resource_uuid=rid)
 
-    transaction = get_transaction_utility()
     processing = get_processing()
     partitioning = get_partitioning()
 
@@ -800,13 +553,7 @@ async def reprocess_file_field(
     writer.source = BrokerMessage.MessageSource.WRITER
     writer.basic.metadata.useful = True
     writer.basic.metadata.status = Metadata.Status.PENDING
-    try:
-        await transaction.commit(writer, partition, wait=False)
-    except TransactionCommitTimeoutError:
-        raise HTTPException(
-            status_code=501,
-            detail="Inconsistent write. This resource will not be processed and may not be stored.",
-        )
+    await transaction.commit(writer, partition, wait=False)
     # Send current resource to reprocess.
     try:
         processing_info = await processing.send_to_process(toprocess, partition)