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

nucliadb/train/uploader.py

@@ -20,6 +20,14 @@
 from typing import Optional
 
 import aiohttp
+
+from nucliadb.common import datamanagers
+from nucliadb.common.maindb.utils import setup_driver
+from nucliadb.ingest.orm.entities import EntitiesManager
+from nucliadb.ingest.orm.processor import Processor
+from nucliadb.train import SERVICE_NAME
+from nucliadb.train.models import RequestData
+from nucliadb.train.settings import settings
 from nucliadb_protos.knowledgebox_pb2 import Labels
 from nucliadb_protos.train_pb2 import (
     EnabledMetadata,
@@ -34,13 +42,6 @@ from nucliadb_protos.writer_pb2 import (
     GetLabelsRequest,
     GetLabelsResponse,
 )
-
-from nucliadb.common.maindb.utils import setup_driver
-from nucliadb.ingest.orm.entities import EntitiesManager
-from nucliadb.ingest.orm.processor import Processor
-from nucliadb.train import SERVICE_NAME
-from nucliadb.train.models import RequestData
-from nucliadb.train.settings import settings
 from nucliadb_utils.utilities import get_pubsub, get_storage
 
 
@@ -51,8 +52,7 @@ class UploadServicer:
         pubsub = await get_pubsub()
         self.proc = Processor(driver=driver, storage=storage, pubsub=pubsub)
 
-    async def finalize(self):
-        ...
+    async def finalize(self): ...
 
     async def GetSentences(self, request: GetSentencesRequest, context=None):
         async for sentence in self.proc.kb_sentences(request):
@@ -75,9 +75,8 @@ class UploadServicer:
     ) -> GetEntitiesResponse:
         kbid = request.kb.uuid
         response = GetEntitiesResponse()
-        async with self.proc.driver.transaction() as txn:
+        async with self.proc.driver.transaction(read_only=True) as txn:
             kbobj = await self.proc.get_kb_obj(txn, request.kb)
-
             if kbobj is None:
                 response.status = GetEntitiesResponse.Status.NOTFOUND
                 return response
@@ -91,20 +90,16 @@ class UploadServicer:
     async def GetOntology(  # type: ignore
         self, request: GetLabelsRequest, context=None
     ) -> GetLabelsResponse:
-        async with self.proc.driver.transaction() as txn:
-            kbobj = await self.proc.get_kb_obj(txn, request.kb)
-            labels: Optional[Labels] = None
-            if kbobj is not None:
-                labels = await kbobj.get_labels()
-
+        kbid = request.kb.uuid
         response = GetLabelsResponse()
-        if kbobj is None:
+        kb_exists = await datamanagers.atomic.kb.exists_kb(kbid=kbid)
+        if not kb_exists:
             response.status = GetLabelsResponse.Status.NOTFOUND
-        else:
-            response.kb.uuid = kbobj.kbid
-            if labels is not None:
-                response.labels.CopyFrom(labels)
-
+            return response
+        response.kb.uuid = kbid
+        labels: Optional[Labels] = await datamanagers.atomic.labelset.get_all(kbid=kbid)
+        if labels is not None:
+            response.labels.CopyFrom(labels)
         return response
 
 
@@ -124,9 +119,9 @@ async def start_upload(request: str, kb: str):
         }
     ) as sess:
         req = await sess.get(f"{url}/request")
-        request_data = RequestData.parse_raw(await req.read())
+        request_data = RequestData.model_validate_json(await req.read())
 
-        metadata = EnabledMetadata(**request_data.metadata.dict())
+        metadata = EnabledMetadata(**request_data.metadata.model_dump())
 
         if request_data.sentences:
             pbsr = GetSentencesRequest()

nucliadb/train/utils.py

@@ -23,7 +23,7 @@ from grpc import aio
 from grpc_health.v1 import health, health_pb2_grpc
 
 from nucliadb.common.maindb.utils import setup_driver, teardown_driver
-from nucliadb.train.nodes import TrainShardManager  # type: ignore
+from nucliadb.train.nodes import TrainShardManager
 from nucliadb.train.settings import settings
 from nucliadb_protos import train_pb2_grpc
 from nucliadb_telemetry.utils import setup_telemetry

nucliadb/writer/__init__.py

@@ -29,9 +29,7 @@ logger = logging.getLogger(SERVICE_NAME)
 class EndpointFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
         return (
-            record.args is not None
-            and len(record.args) >= 3
-            and record.args[2] not in ("/", "/metrics")  # type: ignore
+            record.args is not None and len(record.args) >= 3 and record.args[2] not in ("/", "/metrics")  # type: ignore
         )
 
 

nucliadb/writer/api/constants.py

@@ -23,7 +23,6 @@ from fastapi.params import Header
 
 if TYPE_CHECKING:  # pragma: no cover
     SKIP_STORE_DEFAULT = False
-    SYNC_CALL = False
     X_NUCLIADB_USER = ""
     X_FILE_PASSWORD = None
 else:
@@ -31,10 +30,6 @@ else:
         False,
         description="If set to true, file fields will not be saved in the blob storage. They will only be sent to process.",  # noqa
     )
-    SYNC_CALL = Header(
-        False,
-        description="If set to true, the request will return when the changes to be commited to the database.",
-    )
     X_NUCLIADB_USER = Header("")
     X_FILE_PASSWORD = Header(
         None,

nucliadb/{ingest/fields/keywordset.py → writer/api/utils.py}

@@ -17,18 +17,21 @@
 # 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 FieldKeywordset
+from functools import wraps
 
-from nucliadb.ingest.fields.base import Field
+from fastapi import HTTPException
 
+from nucliadb_utils.settings import is_onprem_nucliadb
 
-class Keywordset(Field):
-    pbklass = FieldKeywordset
-    value: FieldKeywordset
-    type: str = "k"
 
-    async def set_value(self, payload: FieldKeywordset):
-        await self.db_set_value(payload)
+def only_for_onprem(fun):
+    @wraps(fun)
+    async def endpoint_wrapper(*args, **kwargs):
+        if not is_onprem_nucliadb():
+            raise HTTPException(
+                status_code=403,
+                detail="This endpoint is only available for onprem NucliaDB",
+            )
+        return await fun(*args, **kwargs)
 
-    async def get_value(self) -> FieldKeywordset:
-        return await self.db_get_value()
+    return endpoint_wrapper

nucliadb/writer/api/v1/export_import.py

@@ -17,6 +17,8 @@
 # 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
 
 from fastapi_versioning import version
@@ -28,23 +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.models import (
     ExportMetadata,
     ImportMetadata,
     NatsTaskMessage,
 )
 from nucliadb.export_import.tasks import get_exports_producer, get_imports_producer
-from nucliadb.export_import.utils import IteratorExportStream
+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
 
@@ -52,7 +63,7 @@ from nucliadb_utils.authentication import requires_one
 @api.post(
     f"/{KB_PREFIX}/{{kbid}}/export",
     status_code=200,
-    name="Start an export of a Knowledge Box",
+    summary="Start an export of a Knowledge Box",
     tags=["Knowledge Boxes"],
     response_model=CreateExportResponse,
 )
@@ -60,9 +71,8 @@ from nucliadb_utils.authentication import requires_one
 @version(1)
 async def start_kb_export_endpoint(request: Request, kbid: str):
     context = get_app_context(request.app)
-    async with datamanagers.with_transaction(read_only=True) as txn:
-        if not await datamanagers.kb.exists_kb(txn, kbid=kbid):
-            return HTTPClientError(status_code=404, detail="Knowledge Box not found")
+    if not await datamanagers.atomic.kb.exists_kb(kbid=kbid):
+        return HTTPClientError(status_code=404, detail="Knowledge Box not found")
 
     export_id = uuid4().hex
     if in_standalone_mode():
@@ -74,10 +84,60 @@ 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,
-    name="Start an import to a Knowledge Box",
+    summary="Start an import to a Knowledge Box",
     tags=["Knowledge Boxes"],
     response_model=CreateImportResponse,
 )
@@ -85,40 +145,45 @@ async def start_kb_export_endpoint(request: Request, kbid: str):
 @version(1)
 async def start_kb_import_endpoint(request: Request, kbid: str):
     context = get_app_context(request.app)
-    async with datamanagers.with_transaction(read_only=True) as txn:
-        if not await datamanagers.kb.exists_kb(txn, kbid=kbid):
-            return HTTPClientError(status_code=404, detail="Knowledge Box not found")
+    if not await datamanagers.atomic.kb.exists_kb(kbid=kbid):
+        return HTTPClientError(status_code=404, detail="Knowledge Box not found")
 
     await maybe_back_pressure(request, kbid)
 
-    import_id = uuid4().hex
-    if in_standalone_mode():
-        # In standalone mode, we import directly from the request content stream.
-        # Note that we return an import_id simply to keep the API consistent with hosted nucliadb.
-        stream = FastAPIExportStream(request)
-        await importer.import_kb(
-            context=context,
-            kbid=kbid,
-            stream=stream,
-        )
-        return CreateImportResponse(import_id=import_id)
-    else:
-        import_size = await upload_import_to_blob_storage(
-            context=context,
-            request=request,
-            kbid=kbid,
-            import_id=import_id,
-        )
-        await start_import_task(context, kbid, import_id, import_size)
-        return CreateImportResponse(import_id=import_id)
+    stream = stream_compatible_with_kb(kbid, request.stream())
+    try:
+        import_id = uuid4().hex
+        if in_standalone_mode():
+            # In standalone mode, we import directly from the request content stream.
+            # Note that we return an import_id simply to keep the API consistent with hosted nucliadb.
+            await importer.import_kb(
+                context=context,
+                kbid=kbid,
+                stream=stream,
+            )
+            return CreateImportResponse(import_id=import_id)
+        else:
+            import_size = await upload_import_to_blob_storage(
+                context=context,
+                stream=stream,
+                kbid=kbid,
+                import_id=import_id,
+            )
+            await start_import_task(context, kbid, import_id, import_size)
+            return CreateImportResponse(import_id=import_id)
+    except IncompatibleExport as exc:
+        return HTTPClientError(status_code=400, detail=str(exc))
 
 
 async def upload_import_to_blob_storage(
-    context: ApplicationContext, request: Request, kbid: str, import_id: str
+    context: ApplicationContext,
+    stream: AsyncGenerator[bytes, None],
+    kbid: str,
+    import_id: str,
 ) -> int:
     dm = ExportImportDataManager(context.kv_driver, context.blob_storage)
     return await dm.upload_import(
-        import_bytes=request.stream(),
+        import_bytes=stream,
         kbid=kbid,
         import_id=import_id,
     )
@@ -133,37 +198,25 @@ 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)
         raise
-
-
-class FastAPIExportStream(IteratorExportStream):
-    def __init__(self, request: Request):
-        iterator = request.stream().__aiter__()
-        super().__init__(iterator)