134 lines
3.8 KiB
Python
134 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import json
|
|
import tempfile
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from fastapi import HTTPException, UploadFile
|
|
from minio import Minio
|
|
|
|
from .settings import get_settings
|
|
|
|
|
|
@dataclass
|
|
class MinioUploadResult:
|
|
url: str
|
|
object_key: str
|
|
size: int
|
|
|
|
|
|
def _client() -> Minio:
|
|
settings = get_settings()
|
|
endpoint = settings.minio_endpoint
|
|
if settings.minio_port and ":" not in endpoint:
|
|
endpoint = f"{endpoint}:{settings.minio_port}"
|
|
return Minio(
|
|
endpoint,
|
|
access_key=settings.minio_access_key,
|
|
secret_key=settings.minio_secret_key,
|
|
secure=settings.minio_use_ssl,
|
|
region=settings.minio_region,
|
|
)
|
|
|
|
|
|
def _object_key(filename: str, purpose: str) -> str:
|
|
settings = get_settings()
|
|
suffix = Path(filename or "upload.bin").suffix.lower() or ".bin"
|
|
day = datetime.now(timezone.utc).strftime("%Y/%m/%d")
|
|
prefix = settings.minio_upload_prefix or "uploads"
|
|
purpose_part = (purpose or "uploads").replace("/", "-")[:40]
|
|
return f"{prefix}/{purpose_part}/{day}/{uuid.uuid4().hex}{suffix}"
|
|
|
|
|
|
async def _ensure_bucket(client: Minio, bucket: str) -> None:
|
|
exists = await asyncio.to_thread(client.bucket_exists, bucket)
|
|
if not exists:
|
|
await asyncio.to_thread(client.make_bucket, bucket)
|
|
policy = {
|
|
"Version": "2012-10-17",
|
|
"Statement": [
|
|
{
|
|
"Effect": "Allow",
|
|
"Principal": {"AWS": ["*"]},
|
|
"Action": ["s3:GetObject"],
|
|
"Resource": [f"arn:aws:s3:::{bucket}/*"],
|
|
}
|
|
],
|
|
}
|
|
await asyncio.to_thread(client.set_bucket_policy, bucket, json.dumps(policy))
|
|
|
|
|
|
async def upload_file_to_minio(
|
|
file: UploadFile,
|
|
*,
|
|
purpose: str,
|
|
max_size: int,
|
|
) -> MinioUploadResult:
|
|
settings = get_settings()
|
|
if not settings.minio_enabled:
|
|
raise HTTPException(status_code=503, detail="MinIO 存储未启用")
|
|
|
|
total = 0
|
|
with tempfile.SpooledTemporaryFile(max_size=8 * 1024 * 1024) as tmp:
|
|
while True:
|
|
chunk = await file.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
total += len(chunk)
|
|
if total > max_size:
|
|
raise HTTPException(status_code=400, detail=f"文件大小不能超过 {max_size // 1024 // 1024}MB")
|
|
tmp.write(chunk)
|
|
tmp.seek(0)
|
|
|
|
client = _client()
|
|
object_key = _object_key(file.filename or "upload.bin", purpose)
|
|
await _ensure_bucket(client, settings.minio_bucket)
|
|
await asyncio.to_thread(
|
|
client.put_object,
|
|
settings.minio_bucket,
|
|
object_key,
|
|
tmp,
|
|
total,
|
|
content_type=file.content_type or "application/octet-stream",
|
|
metadata={"Cache-Control": "public, max-age=31536000, immutable"},
|
|
)
|
|
|
|
return MinioUploadResult(
|
|
url=f"{settings.minio_public_url}/{object_key}",
|
|
object_key=object_key,
|
|
size=total,
|
|
)
|
|
|
|
|
|
async def upload_bytes_to_minio(
|
|
data: bytes,
|
|
*,
|
|
object_key: str,
|
|
content_type: str,
|
|
) -> MinioUploadResult:
|
|
settings = get_settings()
|
|
if not settings.minio_enabled:
|
|
raise HTTPException(status_code=503, detail="MinIO 存储未启用")
|
|
|
|
client = _client()
|
|
await _ensure_bucket(client, settings.minio_bucket)
|
|
await asyncio.to_thread(
|
|
client.put_object,
|
|
settings.minio_bucket,
|
|
object_key,
|
|
io.BytesIO(data),
|
|
len(data),
|
|
content_type=content_type,
|
|
metadata={"Cache-Control": "public, max-age=31536000, immutable"},
|
|
)
|
|
return MinioUploadResult(
|
|
url=f"{settings.minio_public_url}/{object_key}",
|
|
object_key=object_key,
|
|
size=len(data),
|
|
)
|