Add delete methods to storages and fields (#53)

Author: aminalaee

fastapi_storages/base.py

@@ -21,6 +21,9 @@ def open(self, name: str) -> BinaryIO:
     def write(self, file: BinaryIO, name: str) -> str:
         raise NotImplementedError()
 
+    def delete(self, name: str) -> None:
+        raise NotImplementedError()
+
     def generate_new_filename(self, filename: str) -> str:
         raise NotImplementedError()
 
@@ -72,6 +75,13 @@ def write(self, file: BinaryIO) -> str:
 
         return self._storage.write(file=file, name=self._name)
 
+    def delete(self) -> None:
+        """
+        Delete file from the storage
+        """
+
+        return self._storage.delete(self._name)
+
     def __str__(self) -> str:
         return self.path
 

fastapi_storages/filesystem.py

@@ -64,6 +64,13 @@ def write(self, file: BinaryIO, name: str) -> str:
 
         return str(path)
 
+    def delete(self, name: str) -> None:
+        """
+        Delete the file from the filesystem.
+        """
+
+        Path(self.get_path(name)).unlink()
+
     def generate_new_filename(self, filename: str) -> str:
         counter = 0
         path = self._path / filename

fastapi_storages/s3.py

@@ -54,14 +54,13 @@ def __init__(self) -> None:
 
         self._http_scheme = "https" if self.AWS_S3_USE_SSL else "http"
         self._url = f"{self._http_scheme}://{self.AWS_S3_ENDPOINT_URL}"
-        self._s3 = boto3.resource(
+        self._s3 = boto3.client(
             "s3",
             endpoint_url=self._url,
             use_ssl=self.AWS_S3_USE_SSL,
             aws_access_key_id=self.AWS_ACCESS_KEY_ID,
             aws_secret_access_key=self.AWS_SECRET_ACCESS_KEY,
         )
-        self._bucket = self._s3.Bucket(name=self.AWS_S3_BUCKET_NAME)
 
     def get_name(self, name: str) -> str:
         """
@@ -86,10 +85,8 @@ def get_path(self, name: str) -> str:
             )
 
         if self.AWS_QUERYSTRING_AUTH:
-            params = {"Bucket": self._bucket.name, "Key": key}
-            return self._s3.meta.client.generate_presigned_url(
-                "get_object", Params=params
-            )
+            params = {"Bucket": self.AWS_S3_BUCKET_NAME, "Key": key}
+            return self._s3.generate_presigned_url("get_object", Params=params)
 
         return "{}://{}/{}/{}".format(
             self._http_scheme,
@@ -104,7 +101,9 @@ def get_size(self, name: str) -> int:
         """
 
         key = self.get_name(name)
-        return self._bucket.Object(key).content_length
+        return self._s3.head_object(Bucket=self.AWS_S3_BUCKET_NAME, Key=key)[
+            "ContentLength"
+        ]
 
     def write(self, file: BinaryIO, name: str) -> str:
         """
@@ -118,9 +117,16 @@ def write(self, file: BinaryIO, name: str) -> str:
             "ACL": self.AWS_DEFAULT_ACL,
             "ContentType": content_type or self.default_content_type,
         }
-        self._bucket.upload_fileobj(file, key, ExtraArgs=params)
+        self._s3.upload_fileobj(file, self.AWS_S3_BUCKET_NAME, key, ExtraArgs=params)
         return key
 
+    def delete(self, name: str) -> None:
+        """
+        Delete the file from S3
+        """
+
+        self._s3.delete_object(Bucket=self.AWS_S3_BUCKET_NAME, Key=self.get_name(name))
+
     def generate_new_filename(self, filename: str) -> str:
         key = self.get_name(filename)
         stem = Path(filename).stem
@@ -136,7 +142,7 @@ def generate_new_filename(self, filename: str) -> str:
 
     def _check_object_exists(self, key: str) -> bool:
         try:
-            self._bucket.Object(key).load()
+            self._s3.head_object(Bucket=self.AWS_S3_BUCKET_NAME, Key=key)
         except boto3.exceptions.botocore.exceptions.ClientError as e:
             if e.response["Error"]["Code"] == "404":
                 return False

pyproject.toml

@@ -91,7 +91,7 @@ test = "coverage run -m pytest {args}"
 dependencies = [
   "mkdocs-material==9.5.2",
   "mkdocs==1.5.3",
-  "mkdocstrings[python]==0.24.0",
+  "mkdocstrings[python]==0.25.0",
 ]
 
 [tool.hatch.envs.docs.scripts]

tests/test_filesystem_storage.py

@@ -64,3 +64,18 @@ class NonOverwritingFileSystemStorage(FileSystemStorage):
     assert Path(file1.path) == tmp_path / "duplicate.txt"
     assert Path(file2.path) == tmp_path / "duplicate_1.txt"
     assert Path(file3.path) == tmp_path / "duplicate_2.txt"
+
+
+def test_filesystem_storage_delete_file(tmp_path: Path) -> None:
+    input_file = tmp_path / "input.txt"
+    input_file.write_bytes(b"123")
+
+    storage = FileSystemStorage(path=tmp_path)
+    file = StorageFile(name="example.txt", storage=storage)
+    file.write(file=input_file.open("rb"))
+
+    assert (tmp_path / "example.txt").exists() is True
+
+    file.delete()
+
+    assert (tmp_path / "example.txt").exists() is False

tests/test_s3_storage.py

@@ -2,6 +2,8 @@
 from pathlib import Path
 
 import boto3
+import pytest
+from botocore.exceptions import ClientError
 from moto import mock_s3
 
 from fastapi_storages import S3Storage, StorageFile
@@ -111,3 +113,24 @@ class TestStorage(PrivateS3Storage):
     assert file1.path == "http://s3.fastapi.storages/duplicate.txt"
     assert file2.path == "http://s3.fastapi.storages/duplicate_1.txt"
     assert file3.path == "http://s3.fastapi.storages/duplicate_2.txt"
+
+
+@mock_s3
+def test_s3_storage_delete_file(tmp_path: Path) -> None:
+    s3 = boto3.client("s3")
+    s3.create_bucket(Bucket="bucket")
+
+    tmp_file = tmp_path / "example.txt"
+    tmp_file.write_bytes(b"123")
+
+    storage = PrivateS3Storage()
+
+    file = StorageFile(name="file.txt", storage=storage)
+    file.write(file=tmp_file.open("rb"))
+
+    assert s3.head_object(Bucket="bucket", Key="file.txt")
+
+    file.delete()
+
+    with pytest.raises(ClientError):
+        s3.head_object(Bucket="bucket", Key="file.txt")