Skip to content

Commit

Permalink
feat: add custom-metadata to upload
Browse files Browse the repository at this point in the history
  • Loading branch information
silentworks committed Dec 25, 2024
1 parent dad65f0 commit 058fa2c
Show file tree
Hide file tree
Showing 7 changed files with 113 additions and 20 deletions.
24 changes: 19 additions & 5 deletions storage3/_async/file_api.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from __future__ import annotations

import base64
import json
import urllib.parse
from dataclasses import dataclass, field
from io import BufferedReader, FileIO
Expand Down Expand Up @@ -435,18 +437,30 @@ async def _upload_or_update(
"""
if file_options is None:
file_options = {}
cache_control = file_options.get("cache-control")
cache_control = file_options.pop("cache-control", None)
_data = {}
if file_options.get("upsert"):
file_options.update({"x-upsert": file_options.get("upsert")})
del file_options["upsert"]

upsert = file_options.pop("upsert", None)
if upsert:
file_options.update({"x-upsert": upsert})

metadata = file_options.pop("metadata", None)
file_opts_headers = file_options.pop("headers", None)

headers = {
**self._client.headers,
**DEFAULT_FILE_OPTIONS,
**file_options,
}

if metadata:
metadata_str = json.dumps(metadata)
headers["x-metadata"] = base64.b64encode(metadata_str.encode())
_data.update({"metadata": metadata_str})

if file_opts_headers:
headers.update({**file_opts_headers})

# Only include x-upsert on a POST method
if method != "POST":
del headers["x-upsert"]
Expand All @@ -455,7 +469,7 @@ async def _upload_or_update(

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

if (
isinstance(file, BufferedReader)
Expand Down
24 changes: 19 additions & 5 deletions storage3/_sync/file_api.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
from __future__ import annotations

import base64
import json
import urllib.parse
from dataclasses import dataclass, field
from io import BufferedReader, FileIO
Expand Down Expand Up @@ -435,18 +437,30 @@ def _upload_or_update(
"""
if file_options is None:
file_options = {}
cache_control = file_options.get("cache-control")
cache_control = file_options.pop("cache-control", None)
_data = {}
if file_options.get("upsert"):
file_options.update({"x-upsert": file_options.get("upsert")})
del file_options["upsert"]

upsert = file_options.pop("upsert", None)
if upsert:
file_options.update({"x-upsert": upsert})

metadata = file_options.pop("metadata", None)
file_opts_headers = file_options.pop("headers", None)

headers = {
**self._client.headers,
**DEFAULT_FILE_OPTIONS,
**file_options,
}

if metadata:
metadata_str = json.dumps(metadata)
headers["x-metadata"] = base64.b64encode(metadata_str.encode())
_data.update({"metadata": metadata_str})

if file_opts_headers:
headers.update({**file_opts_headers})

# Only include x-upsert on a POST method
if method != "POST":
del headers["x-upsert"]
Expand All @@ -455,7 +469,7 @@ def _upload_or_update(

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

if (
isinstance(file, BufferedReader)
Expand Down
11 changes: 9 additions & 2 deletions storage3/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from dataclasses import asdict, dataclass
from datetime import datetime
from typing import Literal, Optional, TypedDict, Union
from typing import Any, Dict, Literal, Optional, TypedDict, Union

import dateutil.parser

Expand Down Expand Up @@ -77,7 +77,14 @@ class DownloadOptions(TypedDict, total=False):

FileOptions = TypedDict(
"FileOptions",
{"cache-control": str, "content-type": str, "x-upsert": str, "upsert": str},
{
"cache-control": str,
"content-type": str,
"x-upsert": str,
"upsert": str,
"metadata": Dict[str, Any],
"headers": Dict[str, str],
},
total=False,
)

Expand Down
12 changes: 9 additions & 3 deletions tests/_async/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import asyncio
import os
from collections.abc import AsyncGenerator, Generator

import pytest
from dotenv import load_dotenv
Expand All @@ -14,13 +15,18 @@ def pytest_configure(config) -> None:


@pytest.fixture(scope="package")
def event_loop() -> asyncio.AbstractEventLoop:
def event_loop() -> Generator[asyncio.AbstractEventLoop]:
"""Returns an event loop for the current thread"""
return asyncio.get_event_loop_policy().get_event_loop()
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
yield loop
loop.close()


@pytest.fixture(scope="package")
async def storage() -> AsyncStorageClient:
async def storage() -> AsyncGenerator[AsyncStorageClient]:
url = os.environ.get("SUPABASE_TEST_URL")
assert url is not None, "Must provide SUPABASE_TEST_URL environment variable"
key = os.environ.get("SUPABASE_TEST_KEY")
Expand Down
27 changes: 25 additions & 2 deletions tests/_async/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -398,8 +398,31 @@ async def test_client_get_public_url(
assert response.content == file.file_content


async def test_client_upload_with_custom_metadata(
storage_file_client_public: AsyncBucketProxy, file: FileForTesting
) -> None:
"""Ensure we can get the public url of a file in a bucket"""
await storage_file_client_public.upload(
file.bucket_path,
file.local_path,
{
"content-type": file.mime_type,
"metadata": {"custom": "metadata", "second": "second", "third": "third"},
},
)

info = await storage_file_client_public.info(file.bucket_path)
assert "metadata" in info.keys()
assert info["name"] == file.bucket_path
assert info["metadata"] == {
"custom": "metadata",
"second": "second",
"third": "third",
}


async def test_client_info(
storage_file_client_public: SyncBucketProxy, file: FileForTesting
storage_file_client_public: AsyncBucketProxy, file: FileForTesting
) -> None:
"""Ensure we can get the public url of a file in a bucket"""
await storage_file_client_public.upload(
Expand All @@ -413,7 +436,7 @@ async def test_client_info(


async def test_client_exists(
storage_file_client_public: SyncBucketProxy, file: FileForTesting
storage_file_client_public: AsyncBucketProxy, file: FileForTesting
) -> None:
"""Ensure we can get the public url of a file in a bucket"""
await storage_file_client_public.upload(
Expand Down
12 changes: 9 additions & 3 deletions tests/_sync/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import asyncio
import os
from collections.abc import Generator

import pytest
from dotenv import load_dotenv
Expand All @@ -14,13 +15,18 @@ def pytest_configure(config) -> None:


@pytest.fixture(scope="package")
def event_loop() -> asyncio.AbstractEventLoop:
def event_loop() -> Generator[asyncio.AbstractEventLoop]:
"""Returns an event loop for the current thread"""
return asyncio.get_event_loop_policy().get_event_loop()
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
yield loop
loop.close()


@pytest.fixture(scope="package")
def storage() -> SyncStorageClient:
def storage() -> Generator[SyncStorageClient]:
url = os.environ.get("SUPABASE_TEST_URL")
assert url is not None, "Must provide SUPABASE_TEST_URL environment variable"
key = os.environ.get("SUPABASE_TEST_KEY")
Expand Down
23 changes: 23 additions & 0 deletions tests/_sync/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,29 @@ def test_client_get_public_url(
assert response.content == file.file_content


def test_client_upload_with_custom_metadata(
storage_file_client_public: SyncBucketProxy, file: FileForTesting
) -> None:
"""Ensure we can get the public url of a file in a bucket"""
storage_file_client_public.upload(
file.bucket_path,
file.local_path,
{
"content-type": file.mime_type,
"metadata": {"custom": "metadata", "second": "second", "third": "third"},
},
)

info = storage_file_client_public.info(file.bucket_path)
assert "metadata" in info.keys()
assert info["name"] == file.bucket_path
assert info["metadata"] == {
"custom": "metadata",
"second": "second",
"third": "third",
}


def test_client_info(
storage_file_client_public: SyncBucketProxy, file: FileForTesting
) -> None:
Expand Down

0 comments on commit 058fa2c

Please sign in to comment.