Joey
106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
"""图片/文件存储抽象：本地文件系统（开发兜底）+ 七牛（生产）。"""
|
||
from __future__ import annotations
|
||
|
||
import mimetypes
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
import httpx
|
||
|
||
from config import get_settings
|
||
|
||
# 本地存储根目录（仓库根 data/media/）
|
||
_LOCAL_ROOT = Path(__file__).resolve().parents[2] / "data" / "media"
|
||
|
||
|
||
def _ext_from_url(url: str) -> str:
|
||
ext = mimetypes.guess_extension(url.split("?")[0].lower()) or ".jpg"
|
||
if ext == ".jpe":
|
||
ext = ".jpg"
|
||
return ext
|
||
|
||
|
||
def _ext_from_content_type(content_type: str) -> str:
|
||
ctype = (content_type or "").split(";")[0].strip().lower()
|
||
mapping = {
|
||
"image/jpeg": ".jpg",
|
||
"image/png": ".png",
|
||
"image/webp": ".webp",
|
||
"image/gif": ".gif",
|
||
"image/bmp": ".bmp",
|
||
"image/heic": ".heic",
|
||
"video/mp4": ".mp4",
|
||
}
|
||
return mapping.get(ctype, ".jpg")
|
||
|
||
|
||
async def download_bytes(url: str, referer: str | None = None, timeout: float = 60.0) -> tuple[bytes, str]:
|
||
"""下载远程字节。返回 (bytes, content_type)。"""
|
||
headers = {"Referer": referer} if referer else {}
|
||
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
|
||
resp = await client.get(url, headers=headers)
|
||
resp.raise_for_status()
|
||
ctype = (resp.headers.get("content-type") or "application/octet-stream").split(";")[0].strip()
|
||
return resp.content, ctype
|
||
|
||
|
||
class LocalStorage:
|
||
"""开发期：落 data/media/，由 FastAPI /media 静态托管，返回 app_base_url 可访问 URL。"""
|
||
|
||
async def save_from_url(self, url: str, key_prefix: str = "", referer: str | None = None) -> str:
|
||
data, _ = await download_bytes(url, referer)
|
||
key = self._write(data, key_prefix, url)
|
||
return self.public_url(key)
|
||
|
||
async def save_bytes(self, data: bytes, key: str, content_type: str = "") -> str:
|
||
key = self._write(data, "", key)
|
||
return self.public_url(key)
|
||
|
||
def _write(self, data: bytes, key_prefix: str, hint: str) -> str:
|
||
_LOCAL_ROOT.mkdir(parents=True, exist_ok=True)
|
||
ext = _ext_from_url(hint) if hint and not hint.startswith("data:") else ".jpg"
|
||
key = f"{key_prefix + '/' if key_prefix else ''}{uuid.uuid4().hex}{ext}"
|
||
(_LOCAL_ROOT / key).parent.mkdir(parents=True, exist_ok=True)
|
||
(_LOCAL_ROOT / key).write_bytes(data)
|
||
return key
|
||
|
||
def public_url(self, key: str) -> str:
|
||
settings = get_settings()
|
||
return f"{settings.app_base_url.rstrip('/')}/media/{key}"
|
||
|
||
|
||
class QiniuStorage:
|
||
"""生产：上传七牛，返回绑定域名公网 URL（Ozon 可拉取）。"""
|
||
|
||
def _client(self):
|
||
import qiniu
|
||
|
||
settings = get_settings()
|
||
return qiniu.Auth(settings.qiniu_access_key, settings.qiniu_secret_key), settings
|
||
|
||
async def save_from_url(self, url: str, key_prefix: str = "", referer: str | None = None) -> str:
|
||
data, ctype = await download_bytes(url, referer)
|
||
return await self.save_bytes(data, f"{key_prefix}/{uuid.uuid4().hex}{_ext_from_content_type(ctype)}", ctype)
|
||
|
||
async def save_bytes(self, data: bytes, key: str, content_type: str = "") -> str:
|
||
import qiniu
|
||
|
||
auth, settings = self._client()
|
||
bucket = settings.qiniu_bucket
|
||
token = auth.upload_token(bucket, key, 3600)
|
||
ret, info = qiniu.put_data(token, key, data)
|
||
if info.status_code not in (200,):
|
||
raise RuntimeError(f"七牛上传失败：{info.error or info.text_body or info.status_code}")
|
||
return self.public_url(key)
|
||
|
||
def public_url(self, key: str) -> str:
|
||
settings = get_settings()
|
||
return f"{settings.qiniu_domain.rstrip('/')}/{key}"
|
||
|
||
|
||
def get_storage():
|
||
settings = get_settings()
|
||
if settings.use_qiniu:
|
||
return QiniuStorage()
|
||
return LocalStorage()
|