62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
from app.core.config import settings
|
|
|
|
_s3_client = None
|
|
|
|
|
|
def get_s3_client():
|
|
global _s3_client
|
|
if _s3_client is None:
|
|
import boto3
|
|
|
|
_s3_client = boto3.client(
|
|
"s3",
|
|
endpoint_url=settings.s3_endpoint,
|
|
aws_access_key_id=settings.s3_access_key,
|
|
aws_secret_access_key=settings.s3_secret_key,
|
|
region_name=settings.s3_region,
|
|
)
|
|
return _s3_client
|
|
|
|
|
|
def ensure_bucket() -> None:
|
|
if settings.storage_mode != "s3":
|
|
return
|
|
client = get_s3_client()
|
|
bucket = settings.s3_bucket
|
|
try:
|
|
client.head_bucket(Bucket=bucket)
|
|
except Exception:
|
|
client.create_bucket(Bucket=bucket)
|
|
|
|
|
|
def upload_object(key: str, body: bytes, content_type: str) -> None:
|
|
if settings.storage_mode == "memory":
|
|
memory_store[key] = (body, content_type)
|
|
return
|
|
client = get_s3_client()
|
|
ensure_bucket()
|
|
client.put_object(
|
|
Bucket=settings.s3_bucket,
|
|
Key=key,
|
|
Body=body,
|
|
ContentType=content_type,
|
|
)
|
|
|
|
|
|
def download_object(key: str) -> tuple[bytes, str] | None:
|
|
if settings.storage_mode == "memory":
|
|
return memory_store.get(key)
|
|
client = get_s3_client()
|
|
try:
|
|
response = client.get_object(Bucket=settings.s3_bucket, Key=key)
|
|
body = response["Body"].read()
|
|
content_type = response.get("ContentType", "application/octet-stream")
|
|
return body, content_type
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
memory_store: dict[str, tuple[bytes, str]] = {}
|