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

nucliadb/reader/api/v1/download.py

@@ -28,14 +28,14 @@ from fastapi_versioning import version
 from starlette.datastructures import Headers
 from starlette.responses import StreamingResponse
 
-from nucliadb.ingest.orm.resource import FIELD_TYPE_TO_ID
+from nucliadb.common.ids import FIELD_TYPE_PB_TO_STR
+from nucliadb.common.models_utils import to_proto
 from nucliadb.ingest.serialize import get_resource_uuid_by_slug
 from nucliadb.reader import SERVICE_NAME, logger
-from nucliadb.reader.api.models import FIELD_NAMES_TO_PB_TYPE_MAP
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_utils.authentication import requires_one
-from nucliadb_utils.storages.storage import StorageField  # type: ignore
+from nucliadb_utils.storages.storage import ObjectMetadata, Range, StorageField
 from nucliadb_utils.utilities import get_storage
 
 from .router import KB_PREFIX, RESOURCE_PREFIX, RSLUG_PREFIX, api
@@ -62,9 +62,7 @@ async def download_extract_file_rslug_prefix(
     field_id: str,
     download_field: str,
 ) -> Response:
-    return await _download_extract_file(
-        request, kbid, field_type, field_id, download_field, rslug=rslug
-    )
+    return await _download_extract_file(request, kbid, field_type, field_id, download_field, rslug=rslug)
 
 
 @api.get(
@@ -83,9 +81,7 @@ async def download_extract_file_rid_prefix(
     field_id: str,
     download_field: str,
 ) -> Response:
-    return await _download_extract_file(
-        request, kbid, field_type, field_id, download_field, rid=rid
-    )
+    return await _download_extract_file(request, kbid, field_type, field_id, download_field, rid=rid)
 
 
 async def _download_extract_file(
@@ -101,8 +97,8 @@ async def _download_extract_file(
 
     storage = await get_storage(service_name=SERVICE_NAME)
 
-    pb_field_type = FIELD_NAMES_TO_PB_TYPE_MAP[field_type]
-    field_type_letter = FIELD_TYPE_TO_ID[pb_field_type]
+    pb_field_type = to_proto.field_type_name(field_type)
+    field_type_letter = FIELD_TYPE_PB_TO_STR[pb_field_type]
 
     sf = storage.file_extracted(kbid, rid, field_type_letter, field_id, download_field)
 
@@ -124,9 +120,7 @@ async def download_field_file_rslug_prefix(
     field_id: str,
     inline: bool = False,
 ) -> Response:
-    return await _download_field_file(
-        request, kbid, field_id, rslug=rslug, inline=inline
-    )
+    return await _download_field_file(request, kbid, field_id, rslug=rslug, inline=inline)
 
 
 @api.get(
@@ -164,63 +158,6 @@ async def _download_field_file(
     return await download_api(sf, request.headers, inline=inline)
 
 
-@api.get(
-    f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/layout/{{field_id}}/download/field/{{download_field}}",
-    tags=["Resource fields"],
-    status_code=200,
-    summary="Download layout binary field (by slug)",
-)
-@requires_one([NucliaDBRoles.READER])
-@version(1)
-async def download_field_layout_rslug_prefix(
-    request: Request,
-    kbid: str,
-    rslug: str,
-    field_id: str,
-    download_field: str,
-) -> Response:
-    return await _download_field_layout(
-        request, kbid, field_id, download_field, rslug=rslug
-    )
-
-
-@api.get(
-    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_PREFIX}/{{rid}}/layout/{{field_id}}/download/field/{{download_field}}",
-    tags=["Resource fields"],
-    status_code=200,
-    summary="Download layout binary field (by id)",
-)
-@requires_one([NucliaDBRoles.READER])
-@version(1)
-async def download_field_layout_rid_prefix(
-    request: Request,
-    kbid: str,
-    rid: str,
-    field_id: str,
-    download_field: str,
-) -> Response:
-    return await _download_field_layout(
-        request, kbid, field_id, download_field, rid=rid
-    )
-
-
-async def _download_field_layout(
-    request: Request,
-    kbid: str,
-    field_id: str,
-    download_field: str,
-    rid: Optional[str] = None,
-    rslug: Optional[str] = None,
-) -> Response:
-    rid = await _get_resource_uuid_from_params(kbid, rid, rslug)
-
-    storage = await get_storage(service_name=SERVICE_NAME)
-
-    sf = storage.layout_field(kbid, rid, field_id, download_field)
-
-    return await download_api(sf, request.headers)
-
-
 @api.get(
     f"/{KB_PREFIX}/{{kbid}}/{RSLUG_PREFIX}/{{rslug}}/conversation/{{field_id}}/download/field/{{message_id}}/{{file_num}}",  # noqa
     tags=["Resource fields"],
@@ -237,9 +174,7 @@ async def download_field_conversation_rslug_prefix(
     message_id: str,
     file_num: int,
 ) -> Response:
-    return await _download_field_conversation(
-        request, kbid, field_id, message_id, file_num, rslug=rslug
-    )
+    return await _download_field_conversation(request, kbid, field_id, message_id, file_num, rslug=rslug)
 
 
 @api.get(
@@ -258,9 +193,7 @@ async def download_field_conversation_rid_prefix(
     message_id: str,
     file_num: int,
 ) -> Response:
-    return await _download_field_conversation(
-        request, kbid, field_id, message_id, file_num, rid=rid
-    )
+    return await _download_field_conversation(request, kbid, field_id, message_id, file_num, rid=rid)
 
 
 async def _download_field_conversation(
@@ -282,13 +215,13 @@ async def _download_field_conversation(
 
 
 async def download_api(sf: StorageField, headers: Headers, inline: bool = False):
-    metadata = await sf.exists()
+    metadata: Optional[ObjectMetadata] = await sf.exists()
     if metadata is None:
         raise HTTPException(status_code=404, detail="Specified file doesn't exist")
 
-    file_size = int(metadata.get("SIZE", -1))
-    content_type = metadata.get("CONTENT_TYPE", "application/octet-stream")
-    filename = metadata.get("FILENAME", "file")
+    file_size = metadata.size or -1
+    content_type = metadata.content_type or "application/octet-stream"
+    filename = metadata.filename or "file"
     filename = safe_http_header_encode(filename)
 
     status_code = 200
@@ -299,8 +232,10 @@ async def download_api(sf: StorageField, headers: Headers, inline: bool = False)
         "Content-Type": content_type,
         "Content-Disposition": content_disposition,
     }
-    download_headers = {}
+
+    range = Range()
     if "range" in headers and file_size > -1:
+        status_code = 206
         range_request = headers["range"]
         try:
             start, end, range_size = parse_media_range(range_request, file_size)
@@ -341,23 +276,21 @@ async def download_api(sf: StorageField, headers: Headers, inline: bool = False)
                 headers={"Content-Range": f"bytes */{file_size}"},
                 status_code=416,
             )
-        status_code = 206
         logger.debug(f"Range request: {range_request}")
         extra_headers["Content-Length"] = f"{range_size}"
         extra_headers["Content-Range"] = f"bytes {start}-{end}/{file_size}"
-        download_headers["Range"] = range_request
+        range.start = start
+        range.end = end
 
     return StreamingResponse(
-        sf.storage.download(sf.bucket, sf.key, headers=download_headers),  # type: ignore
+        sf.storage.download(sf.bucket, sf.key, range=range),
         status_code=status_code,
         media_type=content_type,
         headers=extra_headers,
     )
 
 
-async def _get_resource_uuid_from_params(
-    kbid, rid: Optional[str], rslug: Optional[str]
-) -> str:
+async def _get_resource_uuid_from_params(kbid, rid: Optional[str], rslug: Optional[str]) -> str:
     if not any([rid, rslug]):
         raise ValueError("Either rid or slug must be set")
 

nucliadb/reader/api/v1/export_import.py

@@ -160,5 +160,5 @@ async def _get_status(
 
 
 async def exists_kb(kbid: str) -> bool:
-    async with datamanagers.with_transaction(read_only=True) as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         return await datamanagers.kb.exists_kb(txn, kbid=kbid)

nucliadb/reader/api/v1/knowledgebox.py

@@ -18,14 +18,14 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from fastapi import HTTPException
-from fastapi_versioning import version  # type: ignore
+from fastapi_versioning import version
 from starlette.requests import Request
 
 from nucliadb.common import datamanagers
 from nucliadb.common.maindb.utils import get_driver
+from nucliadb.common.models_utils import from_proto
 from nucliadb.reader.api.v1.router import KB_PREFIX, KBS_PREFIX, api
 from nucliadb_models.resource import (
-    KnowledgeBoxConfig,
     KnowledgeBoxList,
     KnowledgeBoxObj,
     KnowledgeBoxObjSummary,
@@ -46,10 +46,10 @@ from nucliadb_utils.authentication import requires, requires_one
 @version(1)
 async def get_kbs(request: Request, prefix: str = "") -> KnowledgeBoxList:
     driver = get_driver()
-    async with driver.transaction() as txn:
+    async with driver.transaction(read_only=True) as txn:
         response = KnowledgeBoxList()
         async for kbid, slug in datamanagers.kb.get_kbs(txn, prefix=prefix):
-            response.kbs.append(KnowledgeBoxObjSummary(slug=slug or None, uuid=kbid))  # type: ignore
+            response.kbs.append(KnowledgeBoxObjSummary(slug=slug or None, uuid=kbid))
         return response
 
 
@@ -64,15 +64,15 @@ async def get_kbs(request: Request, prefix: str = "") -> KnowledgeBoxList:
 @version(1)
 async def get_kb(request: Request, kbid: str) -> KnowledgeBoxObj:
     driver = get_driver()
-    async with driver.transaction() as txn:
+    async with driver.transaction(read_only=True) as txn:
         kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
         if kb_config is None:
             raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
 
         return KnowledgeBoxObj(
             uuid=kbid,
-            slug=kb_config.slug,  # type: ignore
-            config=KnowledgeBoxConfig.from_message(kb_config),
+            slug=kb_config.slug,
+            config=from_proto.knowledgebox_config(kb_config),
         )
 
 
@@ -87,7 +87,7 @@ async def get_kb(request: Request, kbid: str) -> KnowledgeBoxObj:
 @version(1)
 async def get_kb_by_slug(request: Request, slug: str) -> KnowledgeBoxObj:
     driver = get_driver()
-    async with driver.transaction() as txn:
+    async with driver.transaction(read_only=True) as txn:
         kbid = await datamanagers.kb.get_kb_uuid(txn, slug=slug)
         if kbid is None:
             raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
@@ -98,6 +98,6 @@ async def get_kb_by_slug(request: Request, slug: str) -> KnowledgeBoxObj:
 
         return KnowledgeBoxObj(
             uuid=kbid,
-            slug=kb_config.slug,  # type: ignore
-            config=KnowledgeBoxConfig.from_message(kb_config),
+            slug=kb_config.slug,
+            config=from_proto.knowledgebox_config(kb_config),
         )

nucliadb/reader/api/v1/learning_config.py

@@ -44,9 +44,7 @@ async def download_model(
     model_id: str,
     filename: str,
 ):
-    return await learning_config_proxy(
-        request, "GET", f"/download/{kbid}/model/{model_id}/{filename}"
-    )
+    return await learning_config_proxy(request, "GET", f"/download/{kbid}/model/{model_id}/{filename}")
 
 
 @api.get(
@@ -143,7 +141,5 @@ async def get_schema_for_configuration_creation(
     request: Request,
 ):
     if not is_onprem_nucliadb():
-        return HTTPClientError(
-            status_code=404, detail="Endpoint not available for Hosted NucliaDB"
-        )
+        return HTTPClientError(status_code=404, detail="Endpoint not available for Hosted NucliaDB")
     return await learning_config_proxy(request, "GET", f"/schema")

nucliadb/reader/api/v1/resource.py

@@ -22,22 +22,21 @@ from typing import Optional, Union
 from fastapi import Header, HTTPException, Query, Request, Response
 from fastapi_versioning import version
 
-import nucliadb_models as models
+from nucliadb.common.datamanagers.resources import KB_RESOURCE_SLUG_BASE
 from nucliadb.common.maindb.utils import get_driver
+from nucliadb.common.models_utils import from_proto, to_proto
 from nucliadb.ingest.fields.conversation import Conversation
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as ORMKnowledgeBox
-from nucliadb.ingest.orm.resource import KB_RESOURCE_SLUG_BASE
 from nucliadb.ingest.orm.resource import Resource as ORMResource
 from nucliadb.ingest.serialize import (
     managed_serialize,
     serialize,
     set_resource_field_extracted_data,
 )
-from nucliadb.reader import SERVICE_NAME  # type: ignore
+from nucliadb.reader import SERVICE_NAME
 from nucliadb.reader.api import DEFAULT_RESOURCE_LIST_PAGE_SIZE
 from nucliadb.reader.api.models import (
     FIELD_NAME_TO_EXTRACTED_DATA_FIELD_MAP,
-    FIELD_NAMES_TO_PB_TYPE_MAP,
     ResourceField,
 )
 from nucliadb.reader.api.v1.router import KB_PREFIX, RESOURCE_PREFIX, RSLUG_PREFIX, api
@@ -77,60 +76,55 @@ async def list_resources(
 
     # Get counters from maindb
     driver = get_driver()
-    txn = await driver.begin()
-
-    # Filter parameters for serializer
-    show: list[ResourceProperties] = [ResourceProperties.BASIC]
-    field_types: list[FieldTypeName] = []
-    extracted: list[ExtractedDataTypeName] = []
-
-    try:
-        resources: list[Resource] = []
-        max_items_to_iterate = (page + 1) * size
-        first_wanted_item_index = (page * size) + 1  # 1-based index
-        current_key_index = 0
-
-        # ask for one item more than we need, in order to know if it's the last page
-        keys_generator = txn.keys(
-            match=KB_RESOURCE_SLUG_BASE.format(kbid=kbid),
-            count=max_items_to_iterate + 1,
-        )
-        async for key in keys_generator:
-            current_key_index += 1
-
-            # First of all, we need to skip keys, in case we are on a +1 page
-            if page > 0 and current_key_index < first_wanted_item_index:
-                continue
-
-            # Don't fetch keys once we got all items for this
-            if len(resources) == size:
-                await keys_generator.aclose()
-                break
-
-            # Fetch and Add wanted item
-            rid = await txn.get(key)
-            if rid:
-                result = await managed_serialize(
-                    txn,
-                    kbid,
-                    rid.decode(),
-                    show,
-                    field_types,
-                    extracted,
-                    service_name=SERVICE_NAME,
-                )
-                if result is not None:
-                    resources.append(result)
-
-        is_last_page = current_key_index <= max_items_to_iterate
-
-    except Exception as exc:
-        errors.capture_exception(exc)
-        raise HTTPException(
-            status_code=500, detail="Couldn't retrieve list of resources right now"
-        )
-    finally:
-        await txn.abort()
+    async with driver.transaction(read_only=True) as txn:
+        # Filter parameters for serializer
+        show: list[ResourceProperties] = [ResourceProperties.BASIC]
+        field_types: list[FieldTypeName] = []
+        extracted: list[ExtractedDataTypeName] = []
+
+        try:
+            resources: list[Resource] = []
+            max_items_to_iterate = (page + 1) * size
+            first_wanted_item_index = (page * size) + 1  # 1-based index
+            current_key_index = 0
+
+            # ask for one item more than we need, in order to know if it's the last page
+            keys_generator = txn.keys(
+                match=KB_RESOURCE_SLUG_BASE.format(kbid=kbid),
+                count=max_items_to_iterate + 1,
+            )
+            async for key in keys_generator:
+                current_key_index += 1
+
+                # First of all, we need to skip keys, in case we are on a +1 page
+                if page > 0 and current_key_index < first_wanted_item_index:
+                    continue
+
+                # Don't fetch keys once we got all items for this
+                if len(resources) == size:
+                    await keys_generator.aclose()
+                    break
+
+                # Fetch and Add wanted item
+                rid = await txn.get(key, for_update=False)
+                if rid:
+                    result = await managed_serialize(
+                        txn,
+                        kbid,
+                        rid.decode(),
+                        show,
+                        field_types,
+                        extracted,
+                        service_name=SERVICE_NAME,
+                    )
+                    if result is not None:
+                        resources.append(result)
+
+            is_last_page = current_key_index <= max_items_to_iterate
+
+        except Exception as exc:
+            errors.capture_exception(exc)
+            raise HTTPException(status_code=500, detail="Couldn't retrieve list of resources right now")
 
     return ResourceList(
         resources=resources,
@@ -153,9 +147,7 @@ async def get_resource_by_uuid(
     kbid: str,
     rid: str,
     show: list[ResourceProperties] = Query([ResourceProperties.BASIC]),
-    field_type_filter: list[FieldTypeName] = Query(
-        list(FieldTypeName), alias="field_type"
-    ),
+    field_type_filter: list[FieldTypeName] = Query(list(FieldTypeName), alias="field_type"),
     extracted: list[ExtractedDataTypeName] = Query(
         [
             ExtractedDataTypeName.TEXT,
@@ -193,9 +185,7 @@ async def get_resource_by_slug(
     kbid: str,
     rslug: str,
     show: list[ResourceProperties] = Query([ResourceProperties.BASIC]),
-    field_type_filter: list[FieldTypeName] = Query(
-        list(FieldTypeName), alias="field_type"
-    ),
+    field_type_filter: list[FieldTypeName] = Query(list(FieldTypeName), alias="field_type"),
     extracted: list[ExtractedDataTypeName] = Query(
         [
             ExtractedDataTypeName.TEXT,
@@ -235,7 +225,7 @@ async def _get_resource(
     audit = get_audit()
     if audit is not None:
         audit_id = rid if rid else rslug
-        await audit.visited(kbid, audit_id, x_nucliadb_user, x_forwarded_for)  # type: ignore
+        audit.visited(kbid, audit_id, x_nucliadb_user, x_forwarded_for)  # type: ignore
 
     result = await serialize(
         kbid,
@@ -343,9 +333,7 @@ async def _get_resource_field(
 ) -> Response:
     storage = await get_storage(service_name=SERVICE_NAME)
     driver = get_driver()
-
-    pb_field_id = FIELD_NAMES_TO_PB_TYPE_MAP[field_type]
-
+    pb_field_id = to_proto.field_type_name(field_type)
     async with driver.transaction() as txn:
         kb = ORMKnowledgeBox(txn, storage, kbid)
 
@@ -360,34 +348,22 @@ async def _get_resource_field(
         if field is None:
             raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
 
-        resource_field = ResourceField(field_id=field_id, field_type=field_type)  # type: ignore
+        resource_field = ResourceField(field_id=field_id, field_type=field_type)
 
         if ResourceFieldProperties.VALUE in show:
             value = await field.get_value()
 
             if isinstance(value, resources_pb2.FieldText):
                 value = await field.get_value()
-                resource_field.value = models.FieldText.from_message(value)
+                resource_field.value = from_proto.field_text(value)
 
             if isinstance(value, resources_pb2.FieldFile):
                 value = await field.get_value()
-                resource_field.value = models.FieldFile.from_message(value)
+                resource_field.value = from_proto.field_file(value)
 
             if isinstance(value, resources_pb2.FieldLink):
                 value = await field.get_value()
-                resource_field.value = models.FieldLink.from_message(value)
-
-            if isinstance(value, resources_pb2.FieldLayout):
-                value = await field.get_value()
-                resource_field.value = models.FieldLayout.from_message(value)
-
-            if isinstance(value, resources_pb2.FieldDatetime):
-                value = await field.get_value()
-                resource_field.value = models.FieldDatetime.from_message(value)
-
-            if isinstance(value, resources_pb2.FieldKeywordset):
-                value = await field.get_value()
-                resource_field.value = models.FieldKeywordset.from_message(value)
+                resource_field.value = from_proto.field_link(value)
 
             if isinstance(field, Conversation):
                 if page == "first":
@@ -400,12 +376,10 @@ async def _get_resource_field(
 
                 value = await field.get_value(page=page_to_fetch)
                 if value is not None:
-                    resource_field.value = models.Conversation.from_message(value)
+                    resource_field.value = from_proto.conversation(value)
 
         if ResourceFieldProperties.EXTRACTED in show and extracted:
-            resource_field.extracted = FIELD_NAME_TO_EXTRACTED_DATA_FIELD_MAP[
-                field_type
-            ]()
+            resource_field.extracted = FIELD_NAME_TO_EXTRACTED_DATA_FIELD_MAP[field_type]()
             await set_resource_field_extracted_data(
                 field,
                 resource_field.extracted,