Merge pull request #59 from nitrictech/feature/presign-urls

Add presigned url capability

Author: jyecusch

examples/storage/sign_url_read.py

@@ -0,0 +1,15 @@
+# [START import]
+from nitric.api import Storage
+from nitric.api.storage import FileMode
+
+# [END import]
+async def storage_sign_url_read():
+    # [START snippet]
+    # Construct a new storage client with default settings
+    storage = Storage()
+
+    # Create a readonly presigned url for the file valid for the next 3600 seconds
+    await storage.bucket("my-bucket").file("path/to/item").sign_url(mode=FileMode.READ, expiry=3600)
+
+
+# [END snippet]

examples/storage/sign_url_write.py

@@ -0,0 +1,15 @@
+# [START import]
+from nitric.api import Storage
+from nitric.api.storage import FileMode
+
+# [END import]
+async def storage_sign_url_write():
+    # [START snippet]
+    # Construct a new storage client with default settings
+    storage = Storage()
+
+    # Create a writable presigned url for the file valid for the next 3600 seconds
+    await storage.bucket("my-bucket").file("path/to/item").sign_url(mode=FileMode.WRITE, expiry=3600)
+
+
+# [END snippet]

nitric/api/storage.py

@@ -20,9 +20,10 @@
 
 from grpclib import GRPCError
 
-from nitric.api.exception import exception_from_grpc_error
+from nitric.api.exception import exception_from_grpc_error, InvalidArgumentException
 from nitric.utils import new_default_channel
-from nitricapi.nitric.storage.v1 import StorageServiceStub
+from nitricapi.nitric.storage.v1 import StorageServiceStub, StoragePreSignUrlRequestOperation
+from enum import Enum
 
 
 class Storage(object):
@@ -59,6 +60,22 @@ def file(self, key: str):
         return File(_storage=self._storage, _bucket=self.name, key=key)
 
 
+class FileMode(Enum):
+    """Definition of available operation modes for file signed URLs."""
+
+    READ = 0
+    WRITE = 1
+
+    def to_request_operation(self) -> StoragePreSignUrlRequestOperation:
+        """Convert FileMode to a StoragePreSignUrlRequestOperation"""
+        if self == FileMode.READ:
+            return StoragePreSignUrlRequestOperation.READ
+        elif self == FileMode.WRITE:
+            return StoragePreSignUrlRequestOperation.WRITE
+        else:
+            raise InvalidArgumentException("Invalid FileMode")
+
+
 @dataclass(frozen=True, order=True)
 class File(object):
     """A reference to a file in a bucket, used to perform operations on that file."""
@@ -92,3 +109,13 @@ async def delete(self):
             await self._storage._storage_stub.delete(bucket_name=self._bucket, key=self.key)
         except GRPCError as grpc_err:
             raise exception_from_grpc_error(grpc_err)
+
+    async def sign_url(self, mode: FileMode = FileMode.READ, expiry: int = 3600):
+        """Generated a signed url for reading or writing to a file"""
+
+        try:
+            await self._storage._storage_stub.pre_sign_url(
+                bucket_name=self._bucket, key=self.key, operation=mode.to_request_operation(), expiry=expiry
+            )
+        except GRPCError as grpc_err:
+            raise exception_from_grpc_error(grpc_err)

tests/api/test_storage.py

@@ -83,6 +83,22 @@ async def test_delete(self):
         assert mock_read.call_args.kwargs["bucket_name"] == "test-bucket"
         assert mock_read.call_args.kwargs["key"] == "test-file"
 
+    async def test_sign_url(self):
+        mock_pre_sign_url = AsyncMock()
+        mock_pre_sign_url.return_value = Object()
+
+        with patch("nitricapi.nitric.storage.v1.StorageServiceStub.pre_sign_url", mock_pre_sign_url):
+            bucket = Storage().bucket("test-bucket")
+            file = bucket.file("test-file")
+            await file.sign_url()
+
+        # Check expected values were passed to Stub
+        mock_pre_sign_url.assert_called_once()
+        assert mock_pre_sign_url.call_args.kwargs["bucket_name"] == "test-bucket"
+        assert mock_pre_sign_url.call_args.kwargs["key"] == "test-file"
+        assert mock_pre_sign_url.call_args.kwargs["operation"] == 0
+        assert mock_pre_sign_url.call_args.kwargs["expiry"] == 3600
+
     async def test_write_error(self):
         mock_write = AsyncMock()
         mock_write.side_effect = GRPCError(Status.UNKNOWN, "test error")
@@ -106,3 +122,11 @@ async def test_delete_error(self):
         with patch("nitricapi.nitric.storage.v1.StorageServiceStub.delete", mock_delete):
             with pytest.raises(UnknownException) as e:
                 await Storage().bucket("test-bucket").file("test-file").delete()
+
+    async def test_sign_url_error(self):
+        mock_pre_sign_url = AsyncMock()
+        mock_pre_sign_url.side_effect = GRPCError(Status.UNKNOWN, "test error")
+
+        with patch("nitricapi.nitric.storage.v1.StorageServiceStub.pre_sign_url", mock_pre_sign_url):
+            with pytest.raises(UnknownException) as e:
+                await Storage().bucket("test-bucket").file("test-file").sign_url()