nucliadb-utils 4.0.1.post552__py3-none-any.whl → 4.0.1.post553__py3-none-any.whl

nucliadb_utils/storages/gcs.py

@@ -26,7 +26,7 @@ import socket
 from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
 from datetime import datetime
-from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional
+from typing import AsyncGenerator, AsyncIterator, Dict, List, Optional
 from urllib.parse import quote_plus
 
 import aiohttp
@@ -47,7 +47,12 @@ from nucliadb_utils.storages.exceptions import (
     InvalidOffset,
     ResumableUploadGone,
 )
-from nucliadb_utils.storages.storage import Storage, StorageField
+from nucliadb_utils.storages.storage import (
+    ObjectInfo,
+    ObjectMetadata,
+    Storage,
+    StorageField,
+)
 
 storage_ops_observer = metrics.Observer("gcs_ops", labels={"type": ""})
 
@@ -411,7 +416,7 @@ class GCSStorageField(StorageField):
         max_tries=MAX_TRIES,
     )
     @storage_ops_observer.wrap({"type": "exists"})
-    async def exists(self) -> Optional[Dict[str, str]]:
+    async def exists(self) -> Optional[ObjectMetadata]:
         """
         Existence can be checked either with a CloudFile data in the field attribute
         or own StorageField key and bucket. Field takes precendece
@@ -438,16 +443,18 @@ class GCSStorageField(StorageField):
         async with self.storage.session.get(url, headers=headers) as api_resp:
             if api_resp.status == 200:
                 data = await api_resp.json()
-                metadata = data.get("metadata")
-                if metadata is None:
-                    metadata = {}
-                if metadata.get("SIZE") is None:
-                    metadata["SIZE"] = data.get("size")
-                if metadata.get("CONTENT_TYPE") is None:
-                    metadata["CONTENT_TYPE"] = data.get("contentType")
-                if metadata.get("FILENAME") is None:
-                    metadata["FILENAME"] = key.split("/")[-1]
-                return metadata
+                metadata = data.get("metadata") or {}
+                metadata = {k.lower(): v for k, v in metadata.items()}
+                size = metadata.get("size") or data.get("size") or 0
+                content_type = (
+                    metadata.get("content_type") or data.get("contentType") or ""
+                )
+                filename = metadata.get("filename") or key.split("/")[-1]
+                return ObjectMetadata(
+                    filename=filename,
+                    size=int(size),
+                    content_type=content_type,
+                )
             else:
                 return None
 
@@ -718,7 +725,9 @@ class GCSStorage(Storage):
                     errors.capture_message(msg, "error", scope)
         return deleted, conflict
 
-    async def iterate_bucket(self, bucket: str, prefix: str) -> AsyncIterator[Any]:
+    async def iterate_objects(
+        self, bucket: str, prefix: str
+    ) -> AsyncGenerator[ObjectInfo, None]:
         if self.session is None:
             raise AttributeError()
         url = "{}/{}/o".format(self.object_base_url, bucket)
@@ -732,7 +741,7 @@ class GCSStorage(Storage):
             data = await resp.json()
             if "items" in data:
                 for item in data["items"]:
-                    yield item
+                    yield ObjectInfo(name=item["name"])
 
         page_token = data.get("nextPageToken")
         while page_token is not None:
@@ -747,5 +756,5 @@ class GCSStorage(Storage):
                 if len(items) == 0:
                     break
                 for item in items:
-                    yield item
+                    yield ObjectInfo(name=item["name"])
                 page_token = data.get("nextPageToken")

nucliadb_utils/storages/local.py

@@ -24,13 +24,18 @@ import json
 import os
 import shutil
 from datetime import datetime
-from typing import Any, AsyncGenerator, AsyncIterator, Dict, Optional
+from typing import AsyncGenerator, AsyncIterator, Dict, Optional
 
 import aiofiles
 from nucliadb_protos.resources_pb2 import CloudFile
 
 from nucliadb_utils.storages import CHUNK_SIZE
-from nucliadb_utils.storages.storage import Storage, StorageField
+from nucliadb_utils.storages.storage import (
+    ObjectInfo,
+    ObjectMetadata,
+    Storage,
+    StorageField,
+)
 
 
 class LocalStorageField(StorageField):
@@ -142,14 +147,16 @@ class LocalStorageField(StorageField):
 
         init_url = self.storage.get_file_path(self.bucket, upload_uri)
         metadata_init_url = self.metadata_key(init_url)
-        metadata = json.dumps(
-            {"FILENAME": cf.filename, "SIZE": cf.size, "CONTENT_TYPE": cf.content_type}
+        object_metadata = ObjectMetadata(
+            filename=cf.filename,
+            content_type=cf.content_type,
+            size=cf.size,
         )
-
+        raw_metadata = json.dumps(object_metadata.model_dump())
         path_to_create = os.path.dirname(metadata_init_url)
         os.makedirs(path_to_create, exist_ok=True)
         async with aiofiles.open(metadata_init_url, "w+") as resp:
-            await resp.write(metadata)
+            await resp.write(raw_metadata)
 
         self._handler = await aiofiles.threadpool.open(init_url, "wb+")
         field.offset = 0
@@ -190,12 +197,15 @@ class LocalStorageField(StorageField):
         self.field.ClearField("offset")
         self.field.ClearField("upload_uri")
 
-    async def exists(self) -> Optional[Dict[str, str]]:
+    async def exists(self) -> Optional[ObjectMetadata]:
         file_path = self.storage.get_file_path(self.bucket, self.key)
         metadata_path = self.metadata_key(file_path)
         if os.path.exists(metadata_path):
             async with aiofiles.open(metadata_path, "r") as metadata:
-                return json.loads(await metadata.read())
+                raw_metadata = await metadata.read()
+                metadata_dict = json.loads(raw_metadata)
+                metadata_dict = {k.lower(): v for k, v in metadata_dict.items()}
+                return ObjectMetadata.model_validate(metadata_dict)
         return None
 
     async def upload(self, iterator: AsyncIterator, origin: CloudFile) -> CloudFile:
@@ -269,10 +279,11 @@ class LocalStorage(Storage):
             deleted = False
         return deleted
 
-    async def iterate_bucket(self, bucket: str, prefix: str) -> AsyncIterator[Any]:
+    async def iterate_objects(
+        self, bucket: str, prefix: str
+    ) -> AsyncGenerator[ObjectInfo, None]:
         for key in glob.glob(f"{bucket}/{prefix}*"):
-            item = {"name": key}
-            yield item
+            yield ObjectInfo(name=key)
 
     async def download(
         self, bucket_name: str, key: str, headers: Optional[Dict[str, str]] = None

nucliadb_utils/storages/pg.py

@@ -28,7 +28,12 @@ import asyncpg
 from nucliadb_protos.resources_pb2 import CloudFile
 
 from nucliadb_utils.storages import CHUNK_SIZE
-from nucliadb_utils.storages.storage import Storage, StorageField
+from nucliadb_utils.storages.storage import (
+    ObjectInfo,
+    ObjectMetadata,
+    Storage,
+    StorageField,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -288,7 +293,7 @@ where kb_id = $1 and file_id = $2
 
     async def iterate_kb(
         self, bucket: str, prefix: Optional[str] = None
-    ) -> AsyncIterator[FileInfo]:
+    ) -> AsyncGenerator[FileInfo, None]:
         query = """
 SELECT filename, size, content_type, file_id
 FROM kb_files
@@ -517,10 +522,17 @@ class PostgresStorageField(StorageField):
         self.field.ClearField("offset")
         self.field.ClearField("upload_uri")
 
-    async def exists(self) -> Optional[FileInfo]:  # type:ignore
+    async def exists(self) -> Optional[ObjectMetadata]:
         async with self.storage.pool.acquire() as conn:
             dl = PostgresFileDataLayer(conn)
-            return await dl.get_file_info(self.bucket, self.key)
+            file_info = await dl.get_file_info(self.bucket, self.key)
+            if file_info is None:
+                return None
+            return ObjectMetadata(
+                filename=file_info["filename"],
+                size=file_info["size"],
+                content_type=file_info["content_type"],
+            )
 
     async def upload(self, iterator: AsyncIterator, origin: CloudFile) -> CloudFile:
         self.field = await self.start(origin)
@@ -588,11 +600,13 @@ class PostgresStorage(Storage):
         await self.delete_kb(kbid)
         return True
 
-    async def iterate_bucket(self, bucket: str, prefix: str) -> AsyncIterator[Any]:
+    async def iterate_objects(
+        self, bucket: str, prefix: str
+    ) -> AsyncGenerator[ObjectInfo, None]:
         async with self.pool.acquire() as conn:
             dl = PostgresFileDataLayer(conn)
             async for file_data in dl.iterate_kb(bucket, prefix):
-                yield {"name": file_data["key"]}
+                yield ObjectInfo(name=file_data["key"])
 
     async def download(
         self, bucket_name: str, key: str, headers: Optional[dict[str, str]] = None

nucliadb_utils/storages/s3.py

@@ -21,7 +21,7 @@ from __future__ import annotations
 
 from contextlib import AsyncExitStack
 from datetime import datetime
-from typing import Any, AsyncGenerator, AsyncIterator, Optional
+from typing import AsyncGenerator, AsyncIterator, Optional
 
 import aiobotocore  # type: ignore
 import aiohttp
@@ -34,7 +34,12 @@ from nucliadb_protos.resources_pb2 import CloudFile
 from nucliadb_telemetry import errors
 from nucliadb_utils import logger
 from nucliadb_utils.storages.exceptions import UnparsableResponse
-from nucliadb_utils.storages.storage import Storage, StorageField
+from nucliadb_utils.storages.storage import (
+    ObjectInfo,
+    ObjectMetadata,
+    Storage,
+    StorageField,
+)
 
 MB = 1024 * 1024
 MIN_UPLOAD_SIZE = 5 * MB
@@ -272,7 +277,7 @@ class S3StorageField(StorageField):
             MultipartUpload=part_info,
         )
 
-    async def exists(self):
+    async def exists(self) -> Optional[ObjectMetadata]:
         """
         Existence can be checked either with a CloudFile data in the field attribute
         or own StorageField key and bucket. Field takes precendece
@@ -292,13 +297,15 @@ class S3StorageField(StorageField):
         try:
             obj = await self.storage._s3aioclient.head_object(Bucket=bucket, Key=key)
             if obj is not None:
-                metadata = obj.get("Metadata", {})
-                return {
-                    "SIZE": metadata.get("size") or obj.get("ContentLength"),
-                    "CONTENT_TYPE": metadata.get("content_type")
-                    or obj.get("ContentType"),
-                    "FILENAME": metadata.get("filename") or key.split("/")[-1],
-                }
+                metadata = obj.get("Metadata") or {}
+                size = metadata.get("size") or obj.get("ContentLength") or 0
+                content_type = (
+                    metadata.get("content_type") or obj.get("ContentType") or ""
+                )
+                filename = metadata.get("filename") or key.split("/")[-1]
+                return ObjectMetadata(
+                    size=int(size), content_type=content_type, filename=filename
+                )
             else:
                 return None
         except botocore.exceptions.ClientError as e:
@@ -424,14 +431,13 @@ class S3Storage(Storage):
         else:
             raise AttributeError("No valid uri")
 
-    async def iterate_bucket(
+    async def iterate_objects(
         self, bucket: str, prefix: str = "/"
-    ) -> AsyncIterator[Any]:
+    ) -> AsyncGenerator[ObjectInfo, None]:
         paginator = self._s3aioclient.get_paginator("list_objects")
         async for result in paginator.paginate(Bucket=bucket, Prefix=prefix):
             for item in result.get("Contents", []):
-                item["name"] = item["Key"]
-                yield item
+                yield ObjectInfo(name=item["Key"])
 
     async def create_kb(self, kbid: str):
         bucket_name = self.get_bucket_name(kbid)

nucliadb_utils/storages/storage.py

@@ -40,6 +40,7 @@ from nucliadb_protos.noderesources_pb2 import Resource as BrainResource
 from nucliadb_protos.nodewriter_pb2 import IndexMessage
 from nucliadb_protos.resources_pb2 import CloudFile
 from nucliadb_protos.writer_pb2 import BrokerMessage
+from pydantic import BaseModel
 
 from nucliadb_utils import logger
 from nucliadb_utils.helpers import async_gen_lookahead
@@ -60,6 +61,16 @@ INDEXING_KEY = "index/{kb}/{shard}/{resource}/{txid}"
 MESSAGE_KEY = "message/{kbid}/{rid}/{mid}"
 
 
+class ObjectInfo(BaseModel):
+    name: str
+
+
+class ObjectMetadata(BaseModel):
+    filename: str
+    content_type: str
+    size: int
+
+
 class StorageField(abc.ABC, metaclass=abc.ABCMeta):
     storage: Storage
     bucket: str
@@ -99,7 +110,7 @@ class StorageField(abc.ABC, metaclass=abc.ABCMeta):
         return deleted
 
     @abc.abstractmethod
-    async def exists(self) -> Optional[Dict[str, str]]: ...
+    async def exists(self) -> Optional[ObjectMetadata]: ...
 
     @abc.abstractmethod
     async def copy(
@@ -141,10 +152,10 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
         # Delete all keys inside a resource
         bucket = self.get_bucket_name(kbid)
         resource_storage_base_path = STORAGE_RESOURCE.format(kbid=kbid, uuid=uuid)
-        async for bucket_info in self.iterate_bucket(
+        async for object_info in self.iterate_objects(
             bucket, resource_storage_base_path
         ):
-            await self.delete_upload(bucket_info["name"], bucket)
+            await self.delete_upload(object_info.name, bucket)
 
     async def deadletter(
         self, message: BrokerMessage, seq: int, seqid: int, partition: str
@@ -286,7 +297,7 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
         elif file.source == self.source:
             # This is the case for NucliaDB hosted deployment (Nuclia's cloud deployment):
             # The data is already stored in the right place by the processing
-            logger.debug(f"[Nuclia hosted]")
+            logger.debug("[Nuclia hosted]")
             return file
         elif file.source == CloudFile.EXPORT:
             # This is for files coming from an export
@@ -511,7 +522,11 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
     async def finalize(self) -> None: ...
 
     @abc.abstractmethod
-    def iterate_bucket(self, bucket: str, prefix: str) -> AsyncIterator[Any]: ...
+    async def iterate_objects(
+        self, bucket: str, prefix: str
+    ) -> AsyncGenerator[ObjectInfo, None]:
+        raise NotImplementedError()
+        yield ObjectInfo(name="")
 
     async def copy(self, file: CloudFile, destination: StorageField) -> None:
         await destination.copy(

nucliadb_utils/tests/unit/storages/test_pg.py

@@ -523,7 +523,7 @@ class TestPostgresStorage:
         await storage.delete_upload("file_id", "kb_id")
         connection.execute.assert_awaited_with(ANY, "kb_id", "file_id")
 
-    async def test_iterate_bucket(self, storage: pg.PostgresStorage, connection):
+    async def test_iterate_objects(self, storage: pg.PostgresStorage, connection):
         connection.cursor = MagicMock(
             return_value=iter_result(
                 [
@@ -543,11 +543,10 @@ class TestPostgresStorage:
             )
         )
 
-        chunks = []
-        async for chunk in storage.iterate_bucket("kb_id", "file_id"):
-            chunks.append(chunk)
-
-        assert chunks == [{"name": "file_id1"}, {"name": "file_id2"}]
+        object_names = []
+        async for object_info in storage.iterate_objects("kb_id", "file_id"):
+            object_names.append(object_info.name)
+        assert object_names == ["file_id1", "file_id2"]
 
     async def test_download(
         self, storage: pg.PostgresStorage, connection, chunk_info, chunk_data

nucliadb_utils/tests/unit/storages/test_storage.py

@@ -28,6 +28,7 @@ from nucliadb_protos.resources_pb2 import CloudFile
 
 from nucliadb_utils.storages.local import LocalStorageField
 from nucliadb_utils.storages.storage import (
+    ObjectInfo,
     Storage,
     StorageField,
     iter_and_add_size,
@@ -67,8 +68,8 @@ class StorageTest(Storage):
     def get_bucket_name(self, kbid):
         return "bucket"
 
-    async def iterate_bucket(self, bucket_name, prefix):
-        yield {"name": "uri"}
+    async def iterate_objects(self, bucket_name, prefix):
+        yield ObjectInfo(name="uri")
 
     async def download(self, bucket_name, uri):
         br = BrainResource(labels=["label"])

{nucliadb_utils-4.0.1.post552.dist-info → nucliadb_utils-4.0.1.post553.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb_utils
-Version: 4.0.1.post552
+Version: 4.0.1.post553
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -23,8 +23,8 @@ Requires-Dist: PyNaCl
 Requires-Dist: pyjwt >=2.4.0
 Requires-Dist: memorylru >=1.1.2
 Requires-Dist: mrflagly
-Requires-Dist: nucliadb-protos >=4.0.1.post552
-Requires-Dist: nucliadb-telemetry >=4.0.1.post552
+Requires-Dist: nucliadb-protos >=4.0.1.post553
+Requires-Dist: nucliadb-telemetry >=4.0.1.post553
 Provides-Extra: cache
 Requires-Dist: redis >=4.3.4 ; extra == 'cache'
 Requires-Dist: orjson >=3.6.7 ; extra == 'cache'

{nucliadb_utils-4.0.1.post552.dist-info → nucliadb_utils-4.0.1.post553.dist-info}/RECORD

@@ -40,13 +40,13 @@ nucliadb_utils/nuclia_usage/utils/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZ
 nucliadb_utils/nuclia_usage/utils/kb_usage_report.py,sha256=E1eUSFXBVNzQP9Q2rWj9y3koCO5S7iKwckny_AoLKuk,3870
 nucliadb_utils/storages/__init__.py,sha256=5Qc8AUWiJv9_JbGCBpAn88AIJhwDlm0OPQpg2ZdRL4U,872
 nucliadb_utils/storages/exceptions.py,sha256=n6aBOyurWMo8mXd1XY6Psgno4VfXJ9TRbxCy67c08-g,2417
-nucliadb_utils/storages/gcs.py,sha256=3A4eqJe6PoF3oE7c9JSCpRtbcv6dtg6Pd28u4cXKwyE,27099
-nucliadb_utils/storages/local.py,sha256=jgUn3AwaQWn5VBKAvikd97sLAS3jx9PaUE7pwq1iJrk,10097
+nucliadb_utils/storages/gcs.py,sha256=krBkNd7wkHhfIn3T-4QvYu1Rw-envYCa6G4G90oOjvM,27303
+nucliadb_utils/storages/local.py,sha256=JewYQ-fes9iUtUjlbHgWXrG1RsQWh16TJDunJnwfbTg,10447
 nucliadb_utils/storages/nuclia.py,sha256=UfvRu92eqG1v-PE-UWH2x8KEJFqDqATMmUGFmEuqSSs,2097
-nucliadb_utils/storages/pg.py,sha256=yFk6AVgZHPgQq6NwLN_qN7fwD05WgCU5XE7gsFt-B0Q,18633
-nucliadb_utils/storages/s3.py,sha256=8IZoDlTeICZtU1Z0eouaxvafSR6y4GqCtLjFUBxTd1E,19269
+nucliadb_utils/storages/pg.py,sha256=DxXNwcstAFOTC6kaXlWp-b4WrvR8aSSOfgVJNDQ5oDI,18976
+nucliadb_utils/storages/s3.py,sha256=f2bjgmT6JRlUr5DHy3tRUip4kYSA1MzXfYrLNVUp_Cg,19447
 nucliadb_utils/storages/settings.py,sha256=ugCPy1zxBOmA2KosT-4tsjpvP002kg5iQyi42yCGCJA,1285
-nucliadb_utils/storages/storage.py,sha256=lrXa6eWY7HMCUgrEz9-jTk7PBPwP1Nor_zi3oZMiVZ4,20322
+nucliadb_utils/storages/storage.py,sha256=sR2Qvev6eLUvbH1WTXjqXIOnKRy1YMMx6Vsj0wZ2x8A,20585
 nucliadb_utils/tests/__init__.py,sha256=Oo9CAE7B0eW5VHn8sHd6o30SQzOWUhktLPRXdlDOleA,1456
 nucliadb_utils/tests/asyncbenchmark.py,sha256=rN_NNDk4ras0qgFp0QlRyAi9ZU9xITdzxl2s5CigzBo,10698
 nucliadb_utils/tests/conftest.py,sha256=gPYVuVhj_e6Aeanb91wvUerwuxZgaS7d3luIBRQFIU0,1876
@@ -69,10 +69,10 @@ nucliadb_utils/tests/unit/test_utilities.py,sha256=KcHSPp3RZyKAnscJrIwc2M3PCD3l1
 nucliadb_utils/tests/unit/storages/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
 nucliadb_utils/tests/unit/storages/test_aws.py,sha256=GCsB_jwCUNV3Ogt8TZZEmNKAHvOlR0HGU7blrFbtJqs,1924
 nucliadb_utils/tests/unit/storages/test_gcs.py,sha256=2XzJwgNpfjVGjtE-QdZhu3ayuT1EMEXINdM-_SatPCY,3554
-nucliadb_utils/tests/unit/storages/test_pg.py,sha256=sJfUttMSzq8W1XYolAUcMxl_R5HcEzb5fpCklPeMJiY,17000
-nucliadb_utils/tests/unit/storages/test_storage.py,sha256=54nUtElPwT3GQPTT4638F2awBYRl8nD6me9wWijhtbA,7107
-nucliadb_utils-4.0.1.post552.dist-info/METADATA,sha256=vpP4LjOwypL6nmadjZg8BouwRLZGF3gaR1PavmNTRzs,2030
-nucliadb_utils-4.0.1.post552.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-nucliadb_utils-4.0.1.post552.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-4.0.1.post552.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-4.0.1.post552.dist-info/RECORD,,
+nucliadb_utils/tests/unit/storages/test_pg.py,sha256=hOR8WSvnuRJKH_rD2vcSMDOJqs7EndJqaPNrTENBOy8,17016
+nucliadb_utils/tests/unit/storages/test_storage.py,sha256=OT21FUnQTU2oxxaFfVC1LMkJ34jyW7Leit5ZCQgNyZ0,7131
+nucliadb_utils-4.0.1.post553.dist-info/METADATA,sha256=ZtL1febvSlU2s6HaxeRuUpMHxQ7Sn4g7vQjLuaRhSUk,2030
+nucliadb_utils-4.0.1.post553.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+nucliadb_utils-4.0.1.post553.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-4.0.1.post553.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-4.0.1.post553.dist-info/RECORD,,