Skip to content

Commit

Permalink
feat: add update existing file function
Browse files Browse the repository at this point in the history
  • Loading branch information
silentworks committed Nov 22, 2023
1 parent fc8cb5d commit da4d785
Show file tree
Hide file tree
Showing 2 changed files with 125 additions and 6 deletions.
43 changes: 37 additions & 6 deletions storage3/_async/file_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from dataclasses import dataclass, field
from io import BufferedReader, FileIO
from pathlib import Path
from typing import Any, Optional, Union, cast
from typing import Any, Literal, Optional, Union, cast

from httpx import HTTPError, Response

Expand Down Expand Up @@ -344,8 +344,9 @@ async def download(self, path: str, options: DownloadOptions = {}) -> bytes:
)
return response.content

async def upload(
async def _upload_or_update(
self,
method: Literal["POST", "PUT"],
path: str,
file: Union[BufferedReader, bytes, FileIO, str, Path],
file_options: Optional[FileOptions] = None,
Expand All @@ -367,9 +368,6 @@ async def upload(
file_options = {}
cache_control = file_options.get("cache-control")
_data = {}
if cache_control:
file_options["cache-control"] = f"max-age={cache_control}"
_data = {"cacheControl": cache_control}

headers = {
**self._client.headers,
Expand All @@ -378,6 +376,10 @@ async def upload(
}
filename = path.rsplit("/", maxsplit=1)[-1]

if cache_control:
headers["cache-control"] = f"max-age={cache_control}"
_data = {"cacheControl": cache_control}

if (
isinstance(file, BufferedReader)
or isinstance(file, bytes)
Expand All @@ -398,9 +400,38 @@ async def upload(
_path = self._get_final_path(path)

return await self._request(
"POST", f"/object/{_path}", files=files, headers=headers, data=_data
method, f"/object/{_path}", files=files, headers=headers, data=_data
)

async def upload(
self,
path: str,
file: Union[BufferedReader, bytes, FileIO, str, Path],
file_options: Optional[FileOptions] = None,
) -> Response:
"""
Uploads a file to an existing bucket.
Parameters
----------
path
The relative file path including the bucket ID. Should be of the format `bucket/folder/subfolder/filename.png`.
The bucket must already exist before attempting to upload.
file
The File object to be stored in the bucket. or a async generator of chunks
file_options
HTTP headers.
"""
return await self._upload_or_update("POST", path, file, file_options)

async def update(
self,
path: str,
file: Union[BufferedReader, bytes, FileIO, str, Path],
file_options: Optional[FileOptions] = None,
) -> Response:
return await self._upload_or_update("PUT", path, file, file_options)

def _get_final_path(self, path: str) -> str:
return f"{self.id}/{path}"

Expand Down
88 changes: 88 additions & 0 deletions tests/_async/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,67 @@ def file(tmp_path: Path, uuid_factory: Callable[[], str]) -> FileForTesting:
)


@pytest.fixture
def two_files(tmp_path: Path, uuid_factory: Callable[[], str]) -> list[FileForTesting]:
"""Creates multiple test files (different content, same bucket/folder path, different file names)"""
file_name_1 = "test_image_1.svg"
file_name_2 = "test_image_2.svg"
file_content = (
b'<svg width="109" height="113" viewBox="0 0 109 113" fill="none" xmlns="http://www.w3.org/2000/svg"> '
b'<path d="M63.7076 110.284C60.8481 113.885 55.0502 111.912 54.9813 107.314L53.9738 40.0627L99.1935 '
b'40.0627C107.384 40.0627 111.952 49.5228 106.859 55.9374L63.7076 110.284Z" fill="url(#paint0_linear)"/> '
b'<path d="M63.7076 110.284C60.8481 113.885 55.0502 111.912 54.9813 107.314L53.9738 40.0627L99.1935 '
b'40.0627C107.384 40.0627 111.952 49.5228 106.859 55.9374L63.7076 110.284Z" fill="url(#paint1_linear)" '
b'fill-opacity="0.2"/> <path d="M45.317 2.07103C48.1765 -1.53037 53.9745 0.442937 54.0434 5.041L54.4849 '
b'72.2922H9.83113C1.64038 72.2922 -2.92775 62.8321 2.1655 56.4175L45.317 2.07103Z" fill="#3ECF8E"/> <defs>'
b'<linearGradient id="paint0_linear" x1="53.9738" y1="54.974" x2="94.1635" y2="71.8295"'
b'gradientUnits="userSpaceOnUse"> <stop stop-color="#249361"/> <stop offset="1" stop-color="#3ECF8E"/> '
b'</linearGradient> <linearGradient id="paint1_linear" x1="36.1558" y1="30.578" x2="54.4844" y2="65.0806" '
b'gradientUnits="userSpaceOnUse"> <stop/> <stop offset="1" stop-opacity="0"/> </linearGradient> </defs> </svg>'
)
file_content_2 = (
b'<svg width="119" height="123" viewBox="0 0 119 123" fill="none" xmlns="http://www.w3.org/2000/svg"> '
b'<path d="M63.7076 110.284C60.8481 113.885 55.0502 111.912 54.9813 107.314L53.9738 40.0627L99.1935 '
b'40.0627C107.384 40.0627 111.952 49.5228 106.859 55.9374L63.7076 110.284Z" fill="url(#paint0_linear)"/> '
b'<path d="M63.7076 110.284C60.8481 113.885 55.0502 111.912 54.9813 107.314L53.9738 40.0627L99.1935 '
b'40.0627C107.384 40.0627 111.952 49.5228 106.859 55.9374L63.7076 110.284Z" fill="url(#paint1_linear)" '
b'fill-opacity="0.2"/> <path d="M45.317 2.07103C48.1765 -1.53037 53.9745 0.442937 54.0434 5.041L54.4849 '
b'72.2922H9.83113C1.64038 72.2922 -2.92775 62.8321 2.1655 56.4175L45.317 2.07103Z" fill="#3FDF8E"/> <defs>'
b'<linearGradient id="paint0_linear" x1="53.9738" y1="54.974" x2="94.1635" y2="71.8295"'
b'gradientUnits="userSpaceOnUse"> <stop stop-color="#249361"/> <stop offset="1" stop-color="#3FDF8E"/> '
b'</linearGradient> <linearGradient id="paint1_linear" x1="36.1558" y1="30.578" x2="54.4844" y2="65.0806" '
b'gradientUnits="userSpaceOnUse"> <stop/> <stop offset="1" stop-opacity="0"/> </linearGradient> </defs> </svg>'
)
bucket_folder = uuid_factory()
bucket_path_1 = f"{bucket_folder}/{file_name_1}"
bucket_path_2 = f"{bucket_folder}/{file_name_2}"
file_path_1 = tmp_path / file_name_1
file_path_2 = tmp_path / file_name_2
with open(file_path_1, "wb") as f:
f.write(file_content)
with open(file_path_2, "wb") as f:
f.write(file_content_2)

return [
FileForTesting(
name=file_name_1,
local_path=str(file_path_1),
bucket_folder=bucket_folder,
bucket_path=bucket_path_1,
mime_type="image/svg+xml",
file_content=file_content,
),
FileForTesting(
name=file_name_2,
local_path=str(file_path_2),
bucket_folder=bucket_folder,
bucket_path=bucket_path_2,
mime_type="image/svg+xml",
file_content=file_content_2,
),
]


@pytest.fixture
def multi_file(tmp_path: Path, uuid_factory: Callable[[], str]) -> list[FileForTesting]:
"""Creates multiple test files (same content, same bucket/folder path, different file names)"""
Expand Down Expand Up @@ -223,6 +284,33 @@ async def test_client_upload(
assert image_info.get("metadata", {}).get("mimetype") == file.mime_type


async def test_client_update(
storage_file_client: AsyncBucketProxy,
two_files: list[FileForTesting],
) -> None:
"""Ensure we can upload files to a bucket"""
await storage_file_client.upload(
two_files[0].bucket_path,
two_files[0].local_path,
{"content-type": two_files[0].mime_type},
)

await storage_file_client.update(
two_files[0].bucket_path,
two_files[1].local_path,
{"content-type": two_files[1].mime_type},
)

image = await storage_file_client.download(two_files[0].bucket_path)
file_list = await storage_file_client.list(two_files[0].bucket_folder)
image_info = next(
(f for f in file_list if f.get("name") == two_files[0].name), None
)

assert image == two_files[1].file_content
assert image_info.get("metadata", {}).get("mimetype") == two_files[1].mime_type


@pytest.mark.parametrize(
"path", ["foobar.txt", "example/nested.jpg", "/leading/slash.png"]
)
Expand Down

0 comments on commit da4d785

Please sign in to comment.