Merge pull request #13 from Querent-ai/use-storage_result-type

add storage result type for cleanly handling the result of a storage …

Author: saraswatpuneet

querent/storage/local/local_file_storage.py

@@ -11,6 +11,7 @@
 from querent.storage.storage_errors import StorageError, StorageErrorKind
 from querent.storage.storage_base import Storage
 from querent.storage.storage_factory import StorageFactory
+from querent.storage.storage_result import StorageResult
 
 class AsyncDebouncer:
     def __init__(self):
@@ -94,8 +95,8 @@ async def get_all(self, path):
     async def file_num_bytes(self, path):
         return await self.underlying.file_num_bytes(path)
 
-    def uri(self):
-        return self.underlying.uri()
+    def get_uri(self):
+        return self.underlying.get_uri()
 
 class LocalFileStorage(Storage):
     def __init__(self, uri: Uri, root=None):
@@ -128,7 +129,7 @@ async def check_connectivity(self):
                     f"Failed to create directories at {self.root}: {e}",
                 )
 
-    async def put(self, path: Path, payload: PutPayload):
+    async def put(self, path: Path, payload: PutPayload)-> StorageResult:
         full_path = await self.full_path(path)
         parent_dir = full_path.parent
         try:
@@ -139,32 +140,35 @@ async def put(self, path: Path, payload: PutPayload):
                     for i in range(0, payload_len, 1024):
                         chunk = await payload.range_byte_stream(i, i + 1024)
                         file.write(chunk)
+            return StorageResult.success(None)
         except Exception as e:
             raise StorageError(
                 StorageErrorKind.Io,
-                f"Failed to write file to {full_path}: {e}",
+                f"Failed to write file {full_path}: {e}",
             )
 
-    async def copy_to(self, path, output):
+    async def copy_to(self, path, output) -> StorageResult:
         full_path = await self.full_path(path)
         with open(full_path, "rb") as file:
             await asyncio.to_thread(shutil.copyfileobj, file, output)
+        return StorageResult.success(None)
 
-    async def get_slice(self, path, start, end):
+    async def get_slice(self, path, start, end)-> StorageResult:
         full_path = await self.full_path(path)
         with open(full_path, "rb") as file:
             file.seek(start)
-            return file.read(end - start)
+            return StorageResult.success(file.read(end - start))
 
-    async def get_all(self, path):
+    async def get_all(self, path)-> StorageResult:
         full_path = await self.full_path(path)
         with open(full_path, "rb") as file:
-            return file.read()
+            return StorageResult.success(file.read())
 
-    async def delete(self, path):
+    async def delete(self, path)-> StorageResult:
         full_path = await self.full_path(path)
         try:
             full_path.unlink()
+            return StorageResult.success(None)
         except FileNotFoundError:
             pass
         except Exception as e:
@@ -177,16 +181,23 @@ async def bulk_delete(self, paths):
         for path in paths:
             await self.delete(path)
 
-    async def exists(self, path):
+    async def exists(self, path)-> StorageResult:
         full_path = await self.full_path(path)
-        return full_path.exists()
+        return StorageResult.success(full_path.exists())
 
-    async def file_num_bytes(self, path):
+    async def file_num_bytes(self, path)-> StorageResult:
         full_path = await self.full_path(path)
-        return full_path.stat().st_size
+        try:
+            return StorageResult.success(full_path.stat().st_size)
+        except FileNotFoundError:
+            raise StorageError(
+                StorageErrorKind.NotFound,
+                f"File {full_path} not found",
+            )
 
-    def uri(self):
-        return str(self.uri)
+    @property
+    def get_uri(self)-> Uri:
+        return self.uri
 
 class LocalStorageFactory(StorageFactory):
     def backend(self) -> StorageBackend:

querent/storage/storage_base.py

@@ -2,50 +2,52 @@
 from pathlib import Path
 from typing import IO
 
+from querent.common.uri import Uri
 from querent.storage.payload import PutPayload
+from querent.storage.storage_result import StorageResult
 
 class Storage(ABC):
     @abstractmethod
     async def check_connectivity(self) -> None:
         pass
 
     @abstractmethod
-    async def put(self, path: Path, payload: PutPayload) -> None:
+    async def put(self, path: Path, payload: PutPayload) -> StorageResult:
         pass
 
     @abstractmethod
-    async def copy_to(self, path: Path, output: IO[bytes]) -> None:
+    async def copy_to(self, path: Path, output: IO[bytes]) -> StorageResult:
         pass
 
-    async def copy_to_file(self, path: Path, output_path: Path) -> None:
+    async def copy_to_file(self, path: Path, output_path: Path) -> StorageResult:
         async with open(output_path, "wb") as output_file:
             await self.copy_to(path, output_file)
 
     @abstractmethod
-    async def get_slice(self, path: Path, start: int, end: int) -> bytes:
+    async def get_slice(self, path: Path, start: int, end: int) -> StorageResult:
         pass
 
     @abstractmethod
-    async def get_all(self, path: Path) -> bytes:
+    async def get_all(self, path: Path) -> StorageResult:
         pass
 
     @abstractmethod
-    async def delete(self, path: Path) -> None:
+    async def delete(self, path: Path) -> StorageResult:
         pass
 
     @abstractmethod
     async def bulk_delete(self, paths: list[Path]) -> None:
         pass
 
     @abstractmethod
-    async def exists(self, path: Path) -> bool:
+    async def exists(self, path: Path) -> StorageResult:
         pass
 
     @abstractmethod
-    async def file_num_bytes(self, path: Path) -> int:
+    async def file_num_bytes(self, path: Path) -> StorageResult:
         pass
 
     @property
     @abstractmethod
-    def uri(self) -> str:
+    def get_uri(self) -> Uri:
         pass

tests/test_local_storage.py

@@ -38,12 +38,12 @@ def test_storage_resolver(temp_dir):
 
     storage = asyncio.run(resolver.resolve(uri))
 
-    payload = querent_payload.BytesPayload(b"ok")
+    payload = querent_payload.BytesPayload(b"ok testing")
     asyncio.run(storage.put(Path(temp_dir + "/test.txt"), payload))
 
     file_path = Path(temp_dir, "test.txt")
     assert file_path.exists()
 
     with open(file_path, "rb") as file:
         content = file.read()
-        assert content == b"ok"
+        assert content == b"ok testing"