nucliadb-utils 4.0.0.post544__py3-none-any.whl → 4.0.0.post546__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, AsyncIterator, Dict, List, Optional
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, List, Optional
 from urllib.parse import quote_plus
 
 import aiohttp
@@ -221,7 +221,7 @@ class GCSStorageField(StorageField):
                     break
 
     @storage_ops_observer.wrap({"type": "read_range"})
-    async def read_range(self, start: int, end: int) -> AsyncIterator[bytes]:
+    async def read_range(self, start: int, end: int) -> AsyncGenerator[bytes, None]:
         """
         Iterate through ranges of data
         """

nucliadb_utils/storages/local.py

@@ -24,7 +24,7 @@ import json
 import os
 import shutil
 from datetime import datetime
-from typing import Any, AsyncIterator, Dict, Optional
+from typing import Any, AsyncGenerator, AsyncIterator, Dict, Optional
 
 import aiofiles
 from nucliadb_protos.resources_pb2 import CloudFile
@@ -87,7 +87,7 @@ class LocalStorageField(StorageField):
                     break
                 yield data
 
-    async def read_range(self, start: int, end: int) -> AsyncIterator[bytes]:
+    async def read_range(self, start: int, end: int) -> AsyncGenerator[bytes, None]:
         """
         Iterate through ranges of data
         """

nucliadb_utils/storages/pg.py

@@ -22,7 +22,7 @@ from __future__ import annotations
 import asyncio
 import logging
 import uuid
-from typing import Any, AsyncIterator, Optional, TypedDict
+from typing import Any, AsyncGenerator, AsyncIterator, Optional, TypedDict
 
 import asyncpg
 from nucliadb_protos.resources_pb2 import CloudFile
@@ -427,7 +427,7 @@ class PostgresStorageField(StorageField):
             async for chunk in dl.iterate_chunks(bucket, key):
                 yield chunk["data"]
 
-    async def read_range(self, start: int, end: int) -> AsyncIterator[bytes]:
+    async def read_range(self, start: int, end: int) -> AsyncGenerator[bytes, None]:
         """
         Iterate through ranges of data
         """

nucliadb_utils/storages/s3.py

@@ -21,12 +21,13 @@ from __future__ import annotations
 
 from contextlib import AsyncExitStack
 from datetime import datetime
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncGenerator, AsyncIterator, Optional
 
 import aiobotocore  # type: ignore
 import aiohttp
 import backoff  # type: ignore
 import botocore  # type: ignore
+from aiobotocore.client import AioBaseClient  # type: ignore
 from aiobotocore.session import AioSession, get_session  # type: ignore
 from nucliadb_protos.resources_pb2 import CloudFile
 
@@ -111,7 +112,7 @@ class S3StorageField(StorageField):
             yield data
             data = await stream.read(CHUNK_SIZE)
 
-    async def read_range(self, start: int, end: int) -> AsyncIterator[bytes]:
+    async def read_range(self, start: int, end: int) -> AsyncGenerator[bytes, None]:
         """
         Iterate through ranges of data
         """
@@ -319,6 +320,18 @@ class S3StorageField(StorageField):
             Key=destination_uri,
         )
 
+    async def move(
+        self,
+        origin_uri: str,
+        destination_uri: str,
+        origin_bucket_name: str,
+        destination_bucket_name: str,
+    ):
+        await self.copy(
+            origin_uri, destination_uri, origin_bucket_name, destination_bucket_name
+        )
+        await self.storage.delete_upload(origin_uri, origin_bucket_name)
+
     async def upload(self, iterator: AsyncIterator, origin: CloudFile) -> CloudFile:
         self.field = await self.start(origin)
         await self.append(origin, iterator)
@@ -384,7 +397,7 @@ class S3Storage(Storage):
 
     async def initialize(self):
         session = AioSession()
-        self._s3aioclient = await self._exit_stack.enter_async_context(
+        self._s3aioclient: AioBaseClient = await self._exit_stack.enter_async_context(
             session.create_client("s3", **self.opts)
         )
         for bucket in (self.deadletter_bucket, self.indexing_bucket):

nucliadb_utils/storages/storage.py

@@ -60,7 +60,7 @@ INDEXING_KEY = "index/{kb}/{shard}/{resource}/{txid}"
 MESSAGE_KEY = "message/{kbid}/{rid}/{mid}"
 
 
-class StorageField:
+class StorageField(abc.ABC, metaclass=abc.ABCMeta):
     storage: Storage
     bucket: str
     key: str
@@ -78,18 +78,18 @@ class StorageField:
         self.key = fullkey
         self.field = field
 
-    async def upload(self, iterator: AsyncIterator, origin: CloudFile) -> CloudFile:
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def upload(self, iterator: AsyncIterator, origin: CloudFile) -> CloudFile: ...
 
-    async def iter_data(self, headers=None):
+    @abc.abstractmethod
+    async def iter_data(self, headers=None) -> AsyncGenerator[bytes, None]:  # type: ignore
         raise NotImplementedError()
+        yield b""
 
-    async def read_range(self, start: int, end: int) -> AsyncIterator[bytes]:
-        """
-        Iterate through ranges of data
-        """
+    @abc.abstractmethod
+    async def read_range(self, start: int, end: int) -> AsyncGenerator[bytes, None]:
         raise NotImplementedError()
-        yield b""  # pragma: no cover
+        yield b""
 
     async def delete(self) -> bool:
         deleted = False
@@ -98,38 +98,38 @@ class StorageField:
             deleted = True
         return deleted
 
-    async def exists(self) -> Optional[Dict[str, str]]:
-        raise NotImplementedError
+    @abc.abstractmethod
+    async def exists(self) -> Optional[Dict[str, str]]: ...
 
+    @abc.abstractmethod
     async def copy(
         self,
         origin_uri: str,
         destination_uri: str,
         origin_bucket_name: str,
         destination_bucket_name: str,
-    ):
-        raise NotImplementedError()
+    ): ...
 
+    @abc.abstractmethod
     async def move(
         self,
         origin_uri: str,
         destination_uri: str,
         origin_bucket_name: str,
         destination_bucket_name: str,
-    ):
-        raise NotImplementedError()
+    ): ...
 
-    async def start(self, cf: CloudFile) -> CloudFile:
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def start(self, cf: CloudFile) -> CloudFile: ...
 
-    async def append(self, cf: CloudFile, iterable: AsyncIterator) -> int:
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def append(self, cf: CloudFile, iterable: AsyncIterator) -> int: ...
 
-    async def finish(self):
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def finish(self): ...
 
 
-class Storage:
+class Storage(abc.ABC, metaclass=abc.ABCMeta):
     source: int
     field_klass: Type
     deadletter_bucket: Optional[str] = None
@@ -498,40 +498,39 @@ class Storage:
         pb.ParseFromString(payload.read())
         return pb
 
-    async def delete_upload(self, uri: str, bucket_name: str):
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def delete_upload(self, uri: str, bucket_name: str): ...
 
-    def get_bucket_name(self, kbid: str):
-        raise NotImplementedError()
+    @abc.abstractmethod
+    def get_bucket_name(self, kbid: str) -> str: ...
 
-    async def initialize(self):
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def initialize(self) -> None: ...
 
-    async def finalize(self):
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def finalize(self) -> None: ...
 
     @abc.abstractmethod
-    def iterate_bucket(self, bucket: str, prefix: str) -> AsyncIterator[Any]:
-        raise NotImplementedError()
+    def iterate_bucket(self, bucket: str, prefix: str) -> AsyncIterator[Any]: ...
 
-    async def copy(self, file: CloudFile, destination: StorageField):
+    async def copy(self, file: CloudFile, destination: StorageField) -> None:
         await destination.copy(
             file.uri, destination.key, file.bucket_name, destination.bucket
         )
 
-    async def move(self, file: CloudFile, destination: StorageField):
+    async def move(self, file: CloudFile, destination: StorageField) -> None:
         await destination.move(
             file.uri, destination.key, file.bucket_name, destination.bucket
         )
 
-    async def create_kb(self, kbid: str) -> bool:
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def create_kb(self, kbid: str) -> bool: ...
 
-    async def delete_kb(self, kbid: str) -> Tuple[bool, bool]:
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def delete_kb(self, kbid: str) -> Tuple[bool, bool]: ...
 
-    async def schedule_delete_kb(self, kbid: str) -> bool:
-        raise NotImplementedError()
+    @abc.abstractmethod
+    async def schedule_delete_kb(self, kbid: str) -> bool: ...
 
     async def set_stream_message(self, kbid: str, rid: str, data: bytes) -> str:
         key = MESSAGE_KEY.format(kbid=kbid, rid=rid, mid=uuid.uuid4())

nucliadb_utils/tests/unit/storages/test_storage.py

@@ -26,6 +26,7 @@ from nucliadb_protos.noderesources_pb2 import ResourceID
 from nucliadb_protos.nodewriter_pb2 import IndexMessage
 from nucliadb_protos.resources_pb2 import CloudFile
 
+from nucliadb_utils.storages.local import LocalStorageField
 from nucliadb_utils.storages.storage import (
     Storage,
     StorageField,
@@ -45,7 +46,7 @@ class TestStorageField:
 
     @pytest.fixture
     def storage_field(self, storage, field):
-        yield StorageField(storage, "bucket", "fullkey", field)
+        yield LocalStorageField(storage, "bucket", "fullkey", field)
 
     @pytest.mark.asyncio
     async def test_delete(self, storage_field: StorageField, storage):
@@ -73,6 +74,24 @@ class StorageTest(Storage):
         br = BrainResource(labels=["label"])
         yield br.SerializeToString()
 
+    async def create_kb(self, kbid):
+        return True
+
+    async def delete_kb(self, kbid):
+        return True
+
+    async def delete_upload(self, uri, bucket):
+        return True
+
+    async def initialize(self) -> None:
+        pass
+
+    async def finalize(self) -> None:
+        pass
+
+    async def schedule_delete_kb(self, kbid: str) -> bool:
+        return True
+
 
 class TestStorage:
     @pytest.fixture
@@ -136,7 +155,7 @@ class TestStorage:
     async def test_download_pb(self, storage: StorageTest):
         assert isinstance(
             await storage.download_pb(
-                StorageField(storage, "bucket", "fullkey"), BrainResource
+                LocalStorageField(storage, "bucket", "fullkey"), BrainResource
             ),
             BrainResource,
         )

{nucliadb_utils-4.0.0.post544.dist-info → nucliadb_utils-4.0.0.post546.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb-utils
-Version: 4.0.0.post544
+Version: 4.0.0.post546
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -22,8 +22,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.0.post544
-Requires-Dist: nucliadb-telemetry >=4.0.0.post544
+Requires-Dist: nucliadb-protos >=4.0.0.post546
+Requires-Dist: nucliadb-telemetry >=4.0.0.post546
 Provides-Extra: cache
 Requires-Dist: redis >=4.3.4 ; extra == 'cache'
 Requires-Dist: orjson >=3.6.7 ; extra == 'cache'

{nucliadb_utils-4.0.0.post544.dist-info → nucliadb_utils-4.0.0.post546.dist-info}/RECORD

@@ -38,13 +38,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=-M-abojobXGwZaS1kuFXRBxwashsAaTXMStjx3QpX4U,27076
-nucliadb_utils/storages/local.py,sha256=SZrBkxnGg8DoMHlT4GaBP5LqMwi6QAe-d_eLU2MJ0RA,10074
+nucliadb_utils/storages/gcs.py,sha256=3A4eqJe6PoF3oE7c9JSCpRtbcv6dtg6Pd28u4cXKwyE,27099
+nucliadb_utils/storages/local.py,sha256=jgUn3AwaQWn5VBKAvikd97sLAS3jx9PaUE7pwq1iJrk,10097
 nucliadb_utils/storages/nuclia.py,sha256=UfvRu92eqG1v-PE-UWH2x8KEJFqDqATMmUGFmEuqSSs,2097
-nucliadb_utils/storages/pg.py,sha256=vxx8FQIToA_C6CO1uu9W080yZl0SgUVtTEWd6lQhPJc,18610
-nucliadb_utils/storages/s3.py,sha256=EYfn2CjoVL0tPFUqMaHgJYHjrSQaIUzLeNYoRa1qp-A,18809
+nucliadb_utils/storages/pg.py,sha256=yFk6AVgZHPgQq6NwLN_qN7fwD05WgCU5XE7gsFt-B0Q,18633
+nucliadb_utils/storages/s3.py,sha256=8IZoDlTeICZtU1Z0eouaxvafSR6y4GqCtLjFUBxTd1E,19269
 nucliadb_utils/storages/settings.py,sha256=ugCPy1zxBOmA2KosT-4tsjpvP002kg5iQyi42yCGCJA,1285
-nucliadb_utils/storages/storage.py,sha256=igV4CNjv5L09PTY4kqUsGhHzFvDUa_FbfLHxrfMSJDA,20324
+nucliadb_utils/storages/storage.py,sha256=lrXa6eWY7HMCUgrEz9-jTk7PBPwP1Nor_zi3oZMiVZ4,20322
 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
@@ -68,9 +68,9 @@ nucliadb_utils/tests/unit/storages/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS
 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=CwX4wO21og1pdw9IqnK4sFxBfKY3sPhQTd56AS7-Me8,6666
-nucliadb_utils-4.0.0.post544.dist-info/METADATA,sha256=KIbQdGrR0WuExS2FVIhJkm0NuViFu_T-wbBetxN0gok,1979
-nucliadb_utils-4.0.0.post544.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-nucliadb_utils-4.0.0.post544.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-4.0.0.post544.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-4.0.0.post544.dist-info/RECORD,,
+nucliadb_utils/tests/unit/storages/test_storage.py,sha256=54nUtElPwT3GQPTT4638F2awBYRl8nD6me9wWijhtbA,7107
+nucliadb_utils-4.0.0.post546.dist-info/METADATA,sha256=pgruTbhfCt2YIDI8al3iPI6RjTHnTHVcmmm0ZLDDR6c,1979
+nucliadb_utils-4.0.0.post546.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+nucliadb_utils-4.0.0.post546.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-4.0.0.post546.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-4.0.0.post546.dist-info/RECORD,,