feat(storage): support cloud storage
This commit is contained in:
MingxuanGame
13
app/storage/__init__.py
Normal file
13
app/storage/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .aws_s3 import AWSS3StorageService
|
||||
from .base import StorageService
|
||||
from .cloudflare_r2 import CloudflareR2StorageService
|
||||
from .local import LocalStorageService
|
||||
|
||||
__all__ = [
|
||||
"AWSS3StorageService",
|
||||
"CloudflareR2StorageService",
|
||||
"LocalStorageService",
|
||||
"StorageService",
|
||||
]
|
||||
103
app/storage/aws_s3.py
Normal file
103
app/storage/aws_s3.py
Normal file
@@ -0,0 +1,103 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .base import StorageService
|
||||
|
||||
import aioboto3
|
||||
from botocore.exceptions import ClientError
|
||||
|
||||
|
||||
class AWSS3StorageService(StorageService):
|
||||
def __init__(
|
||||
self,
|
||||
access_key_id: str,
|
||||
secret_access_key: str,
|
||||
bucket_name: str,
|
||||
region_name: str,
|
||||
public_url_base: str | None = None,
|
||||
):
|
||||
self.bucket_name = bucket_name
|
||||
self.public_url_base = public_url_base
|
||||
self.session = aioboto3.Session()
|
||||
self.access_key_id = access_key_id
|
||||
self.secret_access_key = secret_access_key
|
||||
self.region_name = region_name
|
||||
|
||||
@property
|
||||
def endpoint_url(self) -> str | None:
|
||||
return None
|
||||
|
||||
def _get_client(self):
|
||||
return self.session.client(
|
||||
"s3",
|
||||
endpoint_url=self.endpoint_url,
|
||||
aws_access_key_id=self.access_key_id,
|
||||
aws_secret_access_key=self.secret_access_key,
|
||||
region_name=self.region_name,
|
||||
)
|
||||
|
||||
async def write_file(
|
||||
self,
|
||||
file_path: str,
|
||||
content: bytes,
|
||||
content_type: str = "application/octet-stream",
|
||||
cache_control: str = "public, max-age=31536000",
|
||||
) -> None:
|
||||
async with self._get_client() as client:
|
||||
await client.put_object(
|
||||
Bucket=self.bucket_name,
|
||||
Key=file_path,
|
||||
Body=content,
|
||||
ContentType=content_type,
|
||||
CacheControl=cache_control,
|
||||
)
|
||||
|
||||
async def read_file(self, file_path: str) -> bytes:
|
||||
async with self._get_client() as client:
|
||||
try:
|
||||
response = await client.get_object(
|
||||
Bucket=self.bucket_name,
|
||||
Key=file_path,
|
||||
)
|
||||
async with response["Body"] as stream:
|
||||
return await stream.read()
|
||||
except ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") == "404":
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
raise RuntimeError(f"Failed to read file from R2: {e}")
|
||||
|
||||
async def delete_file(self, file_path: str) -> None:
|
||||
async with self._get_client() as client:
|
||||
try:
|
||||
await client.delete_object(
|
||||
Bucket=self.bucket_name,
|
||||
Key=file_path,
|
||||
)
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Failed to delete file from R2: {e}")
|
||||
|
||||
async def is_exists(self, file_path: str) -> bool:
|
||||
async with self._get_client() as client:
|
||||
try:
|
||||
await client.head_object(
|
||||
Bucket=self.bucket_name,
|
||||
Key=file_path,
|
||||
)
|
||||
return True
|
||||
except ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") == "404":
|
||||
return False
|
||||
raise RuntimeError(f"Failed to check file existence in R2: {e}")
|
||||
|
||||
async def get_file_url(self, file_path: str) -> str:
|
||||
if self.public_url_base:
|
||||
return f"{self.public_url_base.rstrip('/')}/{file_path.lstrip('/')}"
|
||||
|
||||
async with self._get_client() as client:
|
||||
try:
|
||||
url = await client.generate_presigned_url(
|
||||
"get_object",
|
||||
Params={"Bucket": self.bucket_name, "Key": file_path},
|
||||
)
|
||||
return url
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Failed to generate file URL: {e}")
|
||||
34
app/storage/base.py
Normal file
34
app/storage/base.py
Normal file
@@ -0,0 +1,34 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
|
||||
|
||||
class StorageService(abc.ABC):
|
||||
@abc.abstractmethod
|
||||
async def write_file(
|
||||
self,
|
||||
file_path: str,
|
||||
content: bytes,
|
||||
content_type: str = "application/octet-stream",
|
||||
cache_control: str = "public, max-age=31536000",
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
async def read_file(self, file_path: str) -> bytes:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
async def delete_file(self, file_path: str) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
async def is_exists(self, file_path: str) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
async def get_file_url(self, file_path: str) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
26
app/storage/cloudflare_r2.py
Normal file
26
app/storage/cloudflare_r2.py
Normal file
@@ -0,0 +1,26 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .aws_s3 import AWSS3StorageService
|
||||
|
||||
|
||||
class CloudflareR2StorageService(AWSS3StorageService):
|
||||
def __init__(
|
||||
self,
|
||||
account_id: str,
|
||||
access_key_id: str,
|
||||
secret_access_key: str,
|
||||
bucket_name: str,
|
||||
public_url_base: str | None = None,
|
||||
):
|
||||
super().__init__(
|
||||
access_key_id=access_key_id,
|
||||
secret_access_key=secret_access_key,
|
||||
bucket_name=bucket_name,
|
||||
public_url_base=public_url_base,
|
||||
region_name="auto",
|
||||
)
|
||||
self.account_id = account_id
|
||||
|
||||
@property
|
||||
def endpoint_url(self) -> str:
|
||||
return f"https://{self.account_id}.r2.cloudflarestorage.com"
|
||||
78
app/storage/local.py
Normal file
78
app/storage/local.py
Normal file
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from .base import StorageService
|
||||
|
||||
import aiofiles
|
||||
|
||||
|
||||
class LocalStorageService(StorageService):
|
||||
def __init__(
|
||||
self,
|
||||
storage_path: str,
|
||||
):
|
||||
self.storage_path = Path(storage_path).resolve()
|
||||
self.storage_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _get_file_path(self, file_path: str) -> Path:
|
||||
clean_path = file_path.lstrip("/")
|
||||
full_path = self.storage_path / clean_path
|
||||
|
||||
try:
|
||||
full_path.resolve().relative_to(self.storage_path)
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid file path: {file_path}")
|
||||
|
||||
return full_path
|
||||
|
||||
async def write_file(
|
||||
self,
|
||||
file_path: str,
|
||||
content: bytes,
|
||||
content_type: str = "application/octet-stream",
|
||||
cache_control: str = "public, max-age=31536000",
|
||||
) -> None:
|
||||
full_path = self._get_file_path(file_path)
|
||||
full_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
async with aiofiles.open(full_path, "wb") as f:
|
||||
await f.write(content)
|
||||
except OSError as e:
|
||||
raise RuntimeError(f"Failed to write file: {e}")
|
||||
|
||||
async def read_file(self, file_path: str) -> bytes:
|
||||
full_path = self._get_file_path(file_path)
|
||||
|
||||
if not full_path.exists():
|
||||
raise FileNotFoundError(f"File not found: {file_path}")
|
||||
|
||||
try:
|
||||
async with aiofiles.open(full_path, "rb") as f:
|
||||
return await f.read()
|
||||
except OSError as e:
|
||||
raise RuntimeError(f"Failed to read file: {e}")
|
||||
|
||||
async def delete_file(self, file_path: str) -> None:
|
||||
full_path = self._get_file_path(file_path)
|
||||
|
||||
if not full_path.exists():
|
||||
return
|
||||
|
||||
try:
|
||||
full_path.unlink()
|
||||
|
||||
parent = full_path.parent
|
||||
while parent != self.storage_path and not any(parent.iterdir()):
|
||||
parent.rmdir()
|
||||
parent = parent.parent
|
||||
except OSError as e:
|
||||
raise RuntimeError(f"Failed to delete file: {e}")
|
||||
|
||||
async def is_exists(self, file_path: str) -> bool:
|
||||
full_path = self._get_file_path(file_path)
|
||||
return full_path.exists() and full_path.is_file()
|
||||
|
||||
async def get_file_url(self, file_path: str) -> str:
|
||||
return str(self.storage_path / file_path)
|
||||
Reference in New Issue
Block a user