nucliadb-utils 6.1.0.post2520__py3-none-any.whl → 6.1.0.post2533__py3-none-any.whl

nucliadb_utils/storages/azure.py

@@ -215,6 +215,9 @@ class AzureStorage(Storage):
         async for obj in self.object_store.iterate(bucket, prefix):
             yield obj
 
+    async def insert_object(self, bucket_name: str, key: str, data: bytes) -> None:
+        await self.object_store.insert(bucket_name, key, data)
+
 
 class AzureObjectStore(ObjectStore):
     def __init__(self, account_url: str, connection_string: Optional[str] = None):
@@ -336,6 +339,10 @@ class AzureObjectStore(ObjectStore):
             ),
         )
 
+    async def insert(self, bucket: str, key: str, data: bytes) -> None:
+        container_client = self.service_client.get_container_client(bucket)
+        await container_client.upload_blob(name=key, data=data, length=len(data))
+
     async def download(self, bucket: str, key: str) -> bytes:
         container_client = self.service_client.get_container_client(bucket)
         blob_client = container_client.get_blob_client(key)

nucliadb_utils/storages/gcs.py

@@ -260,11 +260,8 @@ class GCSStorageField(StorageField):
                 source=CloudFile.GCS,
             )
             upload_uri = self.key
-
-        init_url = "{}&name={}".format(
-            self.storage._upload_url.format(bucket=self.bucket),
-            quote_plus(upload_uri),
-        )
+        bucket_upload_url = self.storage._upload_url.format(bucket=field.bucket_name)
+        init_url = f"{bucket_upload_url}?uploadType=resumable&name={quote_plus(upload_uri)}"
         metadata = json.dumps(
             {
                 "metadata": {
@@ -284,7 +281,6 @@ class GCSStorageField(StorageField):
                 "Content-Length": str(call_size),
             }
         )
-
         async with self.storage.session.post(
             init_url,
             headers=headers,
@@ -425,6 +421,7 @@ class GCSStorageField(StorageField):
             else:
                 return None
 
+    @storage_ops_observer.wrap({"type": "upload"})
     async def upload(self, iterator: AsyncIterator, origin: CloudFile) -> CloudFile:
         self.field = await self.start(origin)
         if self.field is None:
@@ -493,7 +490,7 @@ class GCSStorage(Storage):
         # https://cloud.google.com/storage/docs/bucket-locations
         self._bucket_labels = labels or {}
         self._executor = executor
-        self._upload_url = url + "/upload/storage/v1/b/{bucket}/o?uploadType=resumable"  # noqa
+        self._upload_url = url + "/upload/storage/v1/b/{bucket}/o"
         self.object_base_url = url + "/storage/v1/b"
         self._client = None
 
@@ -511,7 +508,6 @@ class GCSStorage(Storage):
         self.session = aiohttp.ClientSession(
             loop=loop, connector=aiohttp.TCPConnector(ttl_dns_cache=60 * 5), timeout=TIMEOUT
         )
-
         try:
             if self.deadletter_bucket is not None and self.deadletter_bucket != "":
                 await self.create_bucket(self.deadletter_bucket)
@@ -689,7 +685,7 @@ class GCSStorage(Storage):
                 logger.info(f"Conflict on deleting bucket {bucket_name}: {details}")
                 conflict = True
             elif resp.status == 404:
-                logger.info(f"Does not exit on deleting: {bucket_name}")
+                logger.info(f"Does not exist on deleting: {bucket_name}")
             else:
                 details = await resp.text()
                 msg = f"Delete KB bucket returned an unexpected status {resp.status}: {details}"
@@ -732,6 +728,27 @@ class GCSStorage(Storage):
                     yield ObjectInfo(name=item["name"])
                 page_token = data.get("nextPageToken")
 
+    @storage_ops_observer.wrap({"type": "insert_object"})
+    async def insert_object(self, bucket_name: str, key: str, data: bytes) -> None:
+        """
+        Put an object in the storage without any metadata.
+        """
+        if self.session is None:  # pragma: no cover
+            raise AttributeError()
+        bucket_upload_url = self._upload_url.format(bucket=bucket_name)
+        url = f"{bucket_upload_url}?uploadType=media&name={quote_plus(key)}"
+        headers = await self.get_access_headers()
+        headers.update(
+            {
+                "Content-Length": str(len(data)),
+                "Content-Type": "application/octet-stream",
+            }
+        )
+        async with self.session.post(url, headers=headers, data=data) as resp:
+            if resp.status != 200:
+                text = await resp.text()
+                raise GoogleCloudException(f"{resp.status}: {text}")
+
 
 def parse_object_metadata(object_data: dict[str, Any], key: str) -> ObjectMetadata:
     custom_metadata: dict[str, str] = object_data.get("metadata") or {}

nucliadb_utils/storages/local.py

@@ -278,8 +278,13 @@ class LocalStorage(Storage):
         return deleted
 
     async def iterate_objects(self, bucket: str, prefix: str) -> AsyncGenerator[ObjectInfo, None]:
-        for key in glob.glob(f"{bucket}/{prefix}*"):
-            yield ObjectInfo(name=key)
+        pathname = f"{self.get_file_path(bucket, prefix)}*"
+        for key in glob.glob(pathname):
+            if key.endswith(".metadata"):
+                # Skip metadata files -- they are internal to the local-storage implementation.
+                continue
+            name = key.split("/")[-1]
+            yield ObjectInfo(name=name)
 
     async def download(self, bucket_name: str, key: str, range: Optional[Range] = None):
         key_path = self.get_file_path(bucket_name, key)
@@ -287,3 +292,9 @@ class LocalStorage(Storage):
             return
         async for chunk in super().download(bucket_name, key, range=range):
             yield chunk
+
+    async def insert_object(self, bucket: str, key: str, data: bytes) -> None:
+        path = self.get_file_path(bucket, key)
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, "wb") as file:
+            file.write(data)

nucliadb_utils/storages/object_store.py

@@ -97,6 +97,18 @@ class ObjectStore(abc.ABC, metaclass=abc.ABCMeta):
         metadata: ObjectMetadata,
     ) -> None: ...
 
+    @abc.abstractmethod
+    async def insert(
+        self,
+        bucket: str,
+        key: str,
+        data: bytes,
+    ) -> None:
+        """
+        Insert data to the object storage without any metadata
+        """
+        ...
+
     @abc.abstractmethod
     async def download(self, bucket: str, key: str) -> bytes: ...
 

nucliadb_utils/storages/s3.py

@@ -330,6 +330,7 @@ class S3StorageField(StorageField):
         await self.copy(origin_uri, destination_uri, origin_bucket_name, destination_bucket_name)
         await self.storage.delete_upload(origin_uri, origin_bucket_name)
 
+    @s3_ops_observer.wrap({"type": "upload"})
     async def upload(self, iterator: AsyncIterator, origin: CloudFile) -> CloudFile:
         self.field = await self.start(origin)
         await self.append(origin, iterator)
@@ -484,6 +485,15 @@ class S3Storage(Storage):
                     deleted = True
         return deleted, conflict
 
+    @s3_ops_observer.wrap({"type": "insert_object"})
+    async def insert_object(self, bucket_name: str, key: str, data: bytes) -> None:
+        await self._s3aioclient.put_object(
+            Bucket=bucket_name,
+            Key=key,
+            Body=data,
+            ContentType="application/octet-stream",
+        )
+
 
 async def bucket_exists(client: AioSession, bucket_name: str) -> bool:
     exists = True

nucliadb_utils/storages/storage.py

@@ -161,7 +161,7 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
             logger.error("No Deadletter Bucket defined will not store the error")
             return
         key = DEADLETTER.format(seqid=seqid, seq=seq, partition=partition)
-        await self.uploadbytes(self.deadletter_bucket, key, message.SerializeToString())
+        await self.upload_object(self.deadletter_bucket, key, message.SerializeToString())
 
     def get_indexing_storage_key(
         self, *, kb: str, logical_shard: str, resource_uid: str, txid: Union[int, str]
@@ -187,7 +187,7 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
             resource_uid=message.resource.uuid,
             txid=txid,
         )
-        await self.uploadbytes(self.indexing_bucket, key, message.SerializeToString())
+        await self.upload_object(self.indexing_bucket, key, message.SerializeToString())
 
         return key
 
@@ -199,7 +199,7 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
         kb: str,
         logical_shard: str,
     ) -> str:
-        if self.indexing_bucket is None:
+        if self.indexing_bucket is None:  # pragma: no cover
             raise AttributeError()
         key = self.get_indexing_storage_key(
             kb=kb,
@@ -207,10 +207,7 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
             resource_uid=message.resource.uuid,
             txid=reindex_id,
         )
-        message_serialized = message.SerializeToString()
-        logger.debug("Starting to upload bytes")
-        await self.uploadbytes(self.indexing_bucket, key, message_serialized)
-        logger.debug("Finished to upload bytes")
+        await self.upload_object(self.indexing_bucket, key, message.SerializeToString())
         return key
 
     async def get_indexing(self, payload: IndexMessage) -> BrainResource:
@@ -259,7 +256,7 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
             + ".deleted"
         )
 
-        await self.uploadbytes(self.indexing_bucket, key, b"")
+        await self.upload_object(self.indexing_bucket, key, b"")
 
     def needs_move(self, file: CloudFile, kbid: str) -> bool:
         # The cloudfile is valid for our environment
@@ -285,11 +282,9 @@ 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("[Nuclia hosted]")
             return file
         elif file.source == CloudFile.EXPORT:
             # This is for files coming from an export
-            logger.debug(f"[Exported file]: {file.uri}")
             new_cf = CloudFile()
             new_cf.CopyFrom(file)
             new_cf.bucket_name = destination.bucket
@@ -298,18 +293,15 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
         elif file.source == CloudFile.FLAPS:
             # NucliaDB On-Prem: the data is stored in NUA, so we need to
             # download it and upload it to NucliaDB's storage
-            logger.debug(f"[NucliaDB OnPrem]: {file.uri}")
             flaps_storage = await get_nuclia_storage()
             iterator = flaps_storage.download(file)
             new_cf = await self.uploaditerator(iterator, destination, file)
         elif file.source == CloudFile.LOCAL:
             # For testing purposes: protobuffer is stored in a file in the local filesystem
-            logger.debug(f"[Local]: {file.uri}")
             local_storage = get_local_storage()
             iterator = local_storage.download(file.bucket_name, file.uri)
             new_cf = await self.uploaditerator(iterator, destination, file)
         elif file.source == CloudFile.EMPTY:
-            logger.warning(f"[Empty file]: {file.uri}")
             new_cf = CloudFile()
             new_cf.CopyFrom(file)
         else:
@@ -379,7 +371,7 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
         cf = await self.uploaditerator(generator, sf, cf)
         return cf
 
-    async def uploadbytes(
+    async def chunked_upload_object(
         self,
         bucket: str,
         key: str,
@@ -387,6 +379,10 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
         filename: str = "payload",
         content_type: str = "",
     ):
+        """
+        Upload bytes to the storage in chunks.
+        This is useful for large files that are already loaded in memory.
+        """
         destination = self.field_klass(storage=self, bucket=bucket, fullkey=key)
 
         cf = CloudFile()
@@ -405,9 +401,16 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
         generator = splitter(buffer)
         await self.uploaditerator(generator, destination, cf)
 
+    # For backwards compatibility
+    uploadbytes = chunked_upload_object
+
     async def uploaditerator(
         self, iterator: AsyncIterator, destination: StorageField, origin: CloudFile
     ) -> CloudFile:
+        """
+        Upload bytes to the storage in chunks, but the data is coming from an iterator.
+        This is when we want to upload large files without loading them in memory.
+        """
         safe_iterator = iterate_storage_compatible(iterator, self, origin)  # type: ignore
         return await destination.upload(safe_iterator, origin)
 
@@ -474,7 +477,7 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
                     yield data
 
     async def upload_pb(self, sf: StorageField, payload: Any):
-        await self.uploadbytes(sf.bucket, sf.key, payload.SerializeToString())
+        await self.upload_object(sf.bucket, sf.key, payload.SerializeToString())
 
     async def download_pb(self, sf: StorageField, PBKlass: Type):
         payload = await self.downloadbytes(sf.bucket, sf.key)
@@ -520,7 +523,8 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
 
     async def set_stream_message(self, kbid: str, rid: str, data: bytes) -> str:
         key = MESSAGE_KEY.format(kbid=kbid, rid=rid, mid=uuid.uuid4())
-        await self.uploadbytes(cast(str, self.indexing_bucket), key, data)
+        indexing_bucket = cast(str, self.indexing_bucket)
+        await self.upload_object(indexing_bucket, key, data)
         return key
 
     async def get_stream_message(self, key: str) -> bytes:
@@ -532,6 +536,23 @@ class Storage(abc.ABC, metaclass=abc.ABCMeta):
     async def del_stream_message(self, key: str) -> None:
         await self.delete_upload(key, cast(str, self.indexing_bucket))
 
+    @abc.abstractmethod
+    async def insert_object(self, bucket: str, key: str, data: bytes) -> None:
+        """
+        Put some binary data into the object storage without any object metadata.
+        """
+        ...
+
+    async def upload_object(self, bucket: str, key: str, data: bytes) -> None:
+        """
+        Put some binary data into the object storage without any object metadata.
+        The data will be uploaded in a single request or in chunks if the data is too large.
+        """
+        if len(data) > self.chunk_size:
+            await self.chunked_upload_object(bucket, key, data)
+        else:
+            await self.insert_object(bucket, key, data)
+
 
 async def iter_and_add_size(
     stream: AsyncGenerator[bytes, None], cf: CloudFile

nucliadb_utils/tests/gcs.py

@@ -75,12 +75,21 @@ class GCS(BaseImage):
 @pytest.fixture(scope="session")
 def gcs():
     container = GCS()
-    _, port = container.run()
-    public_api_url = f"http://172.17.0.1:{port}"
+    host, port = container.run()
+    if running_in_mac_os():
+        public_api_url = f"http://{host}:{port}"
+    else:
+        public_api_url = f"http://172.17.0.1:{port}"
     yield public_api_url
     container.stop()
 
 
+def running_in_mac_os() -> bool:
+    import os
+
+    return os.uname().sysname == "Darwin"
+
+
 @pytest.fixture(scope="function")
 def gcs_storage_settings(gcs) -> dict[str, Any]:
     settings = {

{nucliadb_utils-6.1.0.post2520.dist-info → nucliadb_utils-6.1.0.post2533.dist-info}/METADATA

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

{nucliadb_utils-6.1.0.post2520.dist-info → nucliadb_utils-6.1.0.post2533.dist-info}/RECORD

@@ -40,27 +40,27 @@ nucliadb_utils/nuclia_usage/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn
 nucliadb_utils/nuclia_usage/utils/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/nuclia_usage/utils/kb_usage_report.py,sha256=F6dhQWhrWzl3qfAdvrqV08hnFKcvFlU8IWN2lzYX0vg,3628
 nucliadb_utils/storages/__init__.py,sha256=5Qc8AUWiJv9_JbGCBpAn88AIJhwDlm0OPQpg2ZdRL4U,872
-nucliadb_utils/storages/azure.py,sha256=egMDwLNIGSQyVevuySt2AswzFdNAcih05BbRg3-p8IU,16015
+nucliadb_utils/storages/azure.py,sha256=FK4c_v9AUAwagScm_F1uDmJeQQq7P4jZswiD2trwb4A,16394
 nucliadb_utils/storages/exceptions.py,sha256=mm_wX4YRtp7u7enkk_4pMSlX5AQQuFbq4xLmupVDt3Y,2502
-nucliadb_utils/storages/gcs.py,sha256=AmGOLM19WOlHqzIjSoKZDx7aNosRHE18XgQ5lpkWqkI,27853
-nucliadb_utils/storages/local.py,sha256=jP1glJE8TTPO-ZpYJ8aTxcbmbbay7XJBfUoVWe5PE5U,10501
+nucliadb_utils/storages/gcs.py,sha256=5Yso0o7z4eEVEsBCD3afZlLwMJe6yWvKu1pyB3pqics,28810
+nucliadb_utils/storages/local.py,sha256=JxlWNtu49JJ04dq6o7bBAqbpbeYpVyvvBM5jq1sGJ-4,11003
 nucliadb_utils/storages/nuclia.py,sha256=vEv94xAT7QM2g80S25QyrOw2pzvP2BAX-ADgZLtuCVc,2097
-nucliadb_utils/storages/object_store.py,sha256=Tw10GmpYfM5TMqJ3Tk9pLQ9wLMBk1-snL_m6uasiZDQ,4257
-nucliadb_utils/storages/s3.py,sha256=pbuukqpce_kqkmI_3eUTo390KbM5rmI7h8wsYAXtTAo,20377
+nucliadb_utils/storages/object_store.py,sha256=HtKjIKhErSBvuqx1SuCOnL0SkiHqgfyekNMP8o2piZU,4492
+nucliadb_utils/storages/s3.py,sha256=zV0NBA74JE39caxKd2H33NhMsAtNKo_65tEqJIR-I88,20743
 nucliadb_utils/storages/settings.py,sha256=ugCPy1zxBOmA2KosT-4tsjpvP002kg5iQyi42yCGCJA,1285
-nucliadb_utils/storages/storage.py,sha256=onwQJ4at-XewEG7dxcWdOqobfCw4w0PyPC7olvFJgjI,20295
+nucliadb_utils/storages/storage.py,sha256=eAwr8acamaC2iru3tJaGnkcS35WfioTrnWxcfgKDQSk,21058
 nucliadb_utils/storages/utils.py,sha256=8g2rIwJeYIumQLOB47Yw1rx3twlhRB_cJxer65QfZmk,1479
 nucliadb_utils/tests/__init__.py,sha256=Oo9CAE7B0eW5VHn8sHd6o30SQzOWUhktLPRXdlDOleA,1456
 nucliadb_utils/tests/asyncbenchmark.py,sha256=x4be2IwCawle9zWgMOJkmwoUwk5p1tv7cLQGmybkEOg,10587
 nucliadb_utils/tests/azure.py,sha256=Dg-Eb4KVScG-O6P9y-bVQZTAKTNUMQ0i-CKEd9IdrWw,4474
 nucliadb_utils/tests/fixtures.py,sha256=i0sqPqe5a5JlKGFdaIvOlHYkZ3pHZ2hTIgTsaIB3vSM,3472
-nucliadb_utils/tests/gcs.py,sha256=wO4xRhl-_wrddxlGyEktqa_bXHeQarQN-l_A72Tzr9A,4499
+nucliadb_utils/tests/gcs.py,sha256=KW_DLet1WRlssSW55eI-IQ-0d94Jo2Oh7Di4xGv4JCc,4685
 nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQFxk,1513
 nucliadb_utils/tests/local.py,sha256=fXIBasrvdaFJM-sw2wk1_oiFzBcm9O10iCyC-OiXwY8,1914
 nucliadb_utils/tests/nats.py,sha256=xqpww4jZjTKY9oPGlJdDJG67L3FIBQsa9qDHxILR8r8,7687
 nucliadb_utils/tests/s3.py,sha256=pl-RJFjA4MH6iXkqhsh5g8gDuEhrYu1nPZ-laxlrMlE,3704
-nucliadb_utils-6.1.0.post2520.dist-info/METADATA,sha256=4g5NM63DI6a-tpIXMDuQ0u__2YaMb8M0l7bEf--JPyY,2055
-nucliadb_utils-6.1.0.post2520.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-nucliadb_utils-6.1.0.post2520.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.1.0.post2520.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-6.1.0.post2520.dist-info/RECORD,,
+nucliadb_utils-6.1.0.post2533.dist-info/METADATA,sha256=DeM7PJS1qyI9O3RpgSt8PLyqYG3O8v0dkGnfksl_G88,2055
+nucliadb_utils-6.1.0.post2533.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+nucliadb_utils-6.1.0.post2533.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.1.0.post2533.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-6.1.0.post2533.dist-info/RECORD,,