Files
zkxw/backend/app/core/storage.py
T
2026-08-13 14:31:57 +08:00

425 lines
15 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""文件存储工具模块
支持多种存储方式:
1. 本地存储 (local) - 默认
2. 七牛云存储 (qiniu) - 推荐生产环境使用
3. 阿里云OSS (oss) - 可选
使用方式:
from app.core.storage import storage
# 上传文件
url = storage.upload_file(file_content, "test.jpg", "images")
# 上传 Base64 图片
url = storage.upload_base64(base64_data, "avatar.png", "avatar")
# 删除文件
storage.delete_file("images/test.jpg")
"""
import os
import uuid
import base64
import mimetypes
from io import BytesIO
from datetime import datetime, timedelta
from urllib.parse import quote
from typing import List, Optional, BinaryIO
from pathlib import Path
from app.core.config import settings
# HEIC 格式的 MIME 类型
HEIC_TYPES = {"image/heic", "image/heif"}
# 扩展名 → MIME(补齐部分系统 mimetypes 缺失,尤其 SVG)
_EXTRA_CONTENT_TYPES = {
".svg": "image/svg+xml",
".svgz": "image/svg+xml",
".webp": "image/webp",
".heic": "image/heic",
".heif": "image/heif",
}
def guess_content_type(filename: str) -> str:
"""按文件名推断 Content-Type,供对象存储写入元数据(SVG 回显依赖正确 MIME)。"""
name = (filename or "").strip()
ext = os.path.splitext(name)[1].lower()
if ext in _EXTRA_CONTENT_TYPES:
return _EXTRA_CONTENT_TYPES[ext]
guessed, _ = mimetypes.guess_type(name)
return guessed or "application/octet-stream"
def convert_heic_to_jpeg(content: bytes) -> bytes:
"""将 HEIC/HEIF 格式转换为 JPEG 格式"""
try:
import pillow_heif
from PIL import Image
heif_file = pillow_heif.open(BytesIO(content))
image = heif_file.to_pillow()
output = BytesIO()
image.save(output, format="JPEG", quality=85)
return output.getvalue()
except ImportError:
raise Exception("请安装 pillow-heif: pip install pillow-heif")
def process_image_content(content: bytes, filename: str) -> tuple[bytes, str]:
"""处理图片内容,如果是 HEIC 格式则转换为 JPEG"""
import imghdr
img_type = imghdr.what(None, h=content[:32])
if img_type in ("heic", "heif"):
return convert_heic_to_jpeg(content), filename.rsplit(".", 1)[0] + ".jpg"
return content, filename
class BaseStorage:
"""存储基类"""
def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
"""上传文件"""
raise NotImplementedError
def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
"""上传 Base64 编码的图片"""
raise NotImplementedError
def delete_file(self, file_path: str) -> bool:
"""删除文件"""
raise NotImplementedError
def get_url(self, file_path: str) -> str:
"""获取文件访问URL"""
raise NotImplementedError
def presigned_download_url(
self, file_path: str, download_name: str, expires_seconds: int = 3600
) -> str:
"""带下载文件名的访问 URL;默认实现为普通访问地址。"""
return self.get_url(file_path)
class LocalStorage(BaseStorage):
"""本地存储"""
def __init__(self):
self.base_url = settings.LOCAL_STORAGE_URL.rstrip("/")
self.upload_dir = settings.LOCAL_UPLOAD_DIR
def _ensure_dir(self, folder: str) -> Path:
"""确保目录存在"""
path = Path(self.upload_dir) / folder / datetime.now().strftime("%Y%m")
path.mkdir(parents=True, exist_ok=True)
return path
def _generate_filename(self, filename: str) -> str:
"""生成唯一文件名"""
ext = os.path.splitext(filename)[1] or ""
unique_name = f"{uuid.uuid4().hex}{ext}"
return unique_name
def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
"""上传文件到本地"""
# 处理 HEIC 格式图片
file_content, filename = process_image_content(file_content, filename)
folder_path = self._ensure_dir(folder)
unique_filename = self._generate_filename(filename)
file_path = folder_path / unique_filename
with open(file_path, "wb") as f:
f.write(file_content)
relative_path = f"{folder}/{datetime.now().strftime('%Y%m')}/{unique_filename}" if folder else f"{datetime.now().strftime('%Y%m')}/{unique_filename}"
return self.get_url(relative_path)
def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
"""上传 Base64 图片"""
# 去除 data:image/png;base64, 等前缀
if "," in base64_data:
base64_data = base64_data.split(",")[1]
file_content = base64.b64decode(base64_data)
return self.upload_file(file_content, filename, folder)
def delete_file(self, file_path: str) -> bool:
"""删除本地文件"""
try:
# 移除 URL 前缀获取相对路径
if file_path.startswith(self.base_url):
relative_path = file_path.replace(self.base_url, "").lstrip("/")
else:
relative_path = file_path.lstrip("/")
file_full_path = Path(self.upload_dir) / relative_path
if file_full_path.exists():
file_full_path.unlink()
return True
return False
except Exception:
return False
def get_url(self, file_path: str) -> str:
"""获取文件访问URL"""
return f"{self.base_url}/{file_path}"
class QiniuStorage(BaseStorage):
"""七牛云存储"""
def __init__(self):
try:
from qiniu import Auth, put_data, delete
self.qiniu = Auth(settings.QINIU_ACCESS_KEY, settings.QINIU_SECRET_KEY)
self.bucket = settings.QINIU_BUCKET
self.base_url = settings.QINIU_DOMAIN.rstrip("/")
except ImportError:
raise ImportError("请安装七牛云 SDK: pip install qiniu")
def _generate_filename(self, filename: str) -> str:
"""生成唯一文件名"""
ext = os.path.splitext(filename)[1] or ""
unique_name = f"{datetime.now().strftime('%Y%m/%d')}/{uuid.uuid4().hex}{ext}"
return unique_name
def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
"""上传文件到七牛云"""
from qiniu import put_data
# 处理 HEIC 格式图片
file_content, filename = process_image_content(file_content, filename)
key = self._generate_filename(filename)
token = self.qiniu.upload_token(self.bucket, key, 3600)
mime = guess_content_type(filename)
ret, info = put_data(token, key, file_content, mime_type=mime)
if ret is not None:
return self.get_url(ret["key"])
raise Exception(f"七牛云上传失败: {info}")
def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
"""上传 Base64 图片到七牛云"""
if "," in base64_data:
base64_data = base64_data.split(",")[1]
file_content = base64.b64decode(base64_data)
return self.upload_file(file_content, filename, folder)
def delete_file(self, file_path: str) -> bool:
"""删除七牛云文件"""
from qiniu import delete
try:
key = file_path.replace(f"{self.base_url}/", "")
ret, info = delete(self.bucket, key)
return ret is not None
except Exception:
return False
def get_url(self, file_path: str) -> str:
"""获取文件访问URL"""
return f"{self.base_url}/{file_path}"
class OSSStorage(BaseStorage):
"""阿里云OSS存储"""
def __init__(self):
try:
import oss2
self.auth = oss2.Auth(settings.OSS_ACCESS_KEY_ID, settings.OSS_ACCESS_KEY_SECRET)
self.bucket_name = settings.OSS_BUCKET
self.bucket = oss2.Bucket(self.auth, settings.OSS_ENDPOINT, self.bucket_name)
self.base_url = settings.OSS_DOMAIN.rstrip("/")
except ImportError:
raise ImportError("请安装阿里云 OSS SDK: pip install oss2")
def _generate_filename(self, filename: str) -> str:
"""生成唯一文件名"""
ext = os.path.splitext(filename)[1] or ""
unique_name = f"{datetime.now().strftime('%Y%m/%d')}/{uuid.uuid4().hex}{ext}"
return unique_name
def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
"""上传文件到OSS"""
# 处理 HEIC 格式图片
file_content, filename = process_image_content(file_content, filename)
key = self._generate_filename(filename)
headers = {"Content-Type": guess_content_type(filename)}
result = self.bucket.put_object(key, file_content, headers=headers)
if result.status == 200:
return self.get_url(key)
raise Exception(f"OSS上传失败: {result.status}")
def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
"""上传 Base64 图片到OSS"""
if "," in base64_data:
base64_data = base64_data.split(",")[1]
file_content = base64.b64decode(base64_data)
return self.upload_file(file_content, filename, folder)
def delete_file(self, file_path: str) -> bool:
"""删除OSS文件"""
try:
key = file_path.replace(f"{self.base_url}/", "")
result = self.bucket.delete_object(key)
return result.status == 204
except Exception:
return False
def get_url(self, file_path: str) -> str:
"""获取文件访问URL"""
return f"{self.base_url}/{file_path}"
class MinIOStorage(BaseStorage):
"""MinIO存储"""
def __init__(self):
try:
from minio import Minio
self.client = Minio(
settings.MINIO_ENDPOINT,
access_key=settings.MINIO_ACCESS_KEY,
secret_key=settings.MINIO_SECRET_KEY,
secure=settings.MINIO_SECURE,
)
self.bucket_name = settings.MINIO_BUCKET
if settings.MINIO_BASE_URL:
self.base_url = settings.MINIO_BASE_URL.rstrip("/")
else:
scheme = "https" if settings.MINIO_SECURE else "http"
self.base_url = f"{scheme}://{settings.MINIO_ENDPOINT}/{self.bucket_name}"
except ImportError:
raise ImportError("请安装 MinIO SDK: pip install minio")
def _generate_filename(self, filename: str) -> str:
"""生成唯一文件名"""
ext = os.path.splitext(filename)[1] or ""
unique_name = f"{datetime.now().strftime('%Y%m/%d')}/{uuid.uuid4().hex}{ext}"
return unique_name
def upload_file(self, file_content: bytes, filename: str, folder: str = "") -> str:
"""上传文件到MinIO"""
# 处理 HEIC 格式图片
file_content, filename = process_image_content(file_content, filename)
key = self._generate_filename(filename)
if folder:
key = f"{folder.strip('/')}/{key}"
data = BytesIO(file_content)
# 必须写入正确 MIME:SVG 若为 octet-stream,浏览器 <img> 会拒绝回显
self.client.put_object(
bucket_name=self.bucket_name,
object_name=key,
data=data,
length=len(file_content),
content_type=guess_content_type(filename),
)
return self.get_url(key)
def upload_base64(self, base64_data: str, filename: str, folder: str = "") -> str:
"""上传 Base64 图片到MinIO"""
if "," in base64_data:
base64_data = base64_data.split(",")[1]
file_content = base64.b64decode(base64_data)
return self.upload_file(file_content, filename, folder)
def delete_file(self, file_path: str) -> bool:
"""删除MinIO文件"""
try:
key = file_path.replace(f"{self.base_url}/", "")
self.client.remove_object(self.bucket_name, key)
return True
except Exception:
return False
def get_url(self, file_path: str) -> str:
"""获取文件访问URL"""
return f"{self.base_url}/{file_path}"
def _extract_object_key(self, file_path: str) -> str:
u = (file_path or "").strip()
if not u:
return ""
if "?" in u:
u = u.split("?", 1)[0]
prefix = f"{self.base_url}/"
if u.startswith(prefix):
return u[len(prefix) :]
if u.startswith("http://") or u.startswith("https://"):
marker = f"/{self.bucket_name}/"
idx = u.find(marker)
if idx >= 0:
return u[idx + len(marker) :]
return u.lstrip("/")
def _internal_base_prefixes(self) -> List[str]:
"""SDK 预签名 URL 使用的内网地址前缀(http/https 均覆盖)。"""
endpoint = (settings.MINIO_ENDPOINT or "").strip()
if not endpoint:
return []
bucket = self.bucket_name
return [
f"http://{endpoint}/{bucket}",
f"https://{endpoint}/{bucket}",
]
def _rewrite_presigned_public_url(self, url: str) -> str:
"""将 SDK 生成的内网 MinIO 预签名 URL 改写为 MINIO_BASE_URL 对外路径。"""
if not url:
return url
public_base = (self.base_url or "").rstrip("/")
if not public_base:
return url
for prefix in self._internal_base_prefixes():
if url.startswith(prefix):
return public_base + url[len(prefix) :]
return url
def presigned_download_url(
self, file_path: str, download_name: str, expires_seconds: int = 3600
) -> str:
"""MinIO 预签名下载 URL,携带原始文件名。"""
key = self._extract_object_key(file_path)
if not key:
raise ValueError("无效的文件路径")
name = (download_name or "").strip() or key.rsplit("/", 1)[-1] or "download"
ascii_name = name.encode("ascii", "ignore").decode() or "download"
encoded = quote(name, safe="")
disposition = f'attachment; filename="{ascii_name}"; filename*=UTF-8\'\'{encoded}'
url = self.client.presigned_get_object(
bucket_name=self.bucket_name,
object_name=key,
expires=timedelta(seconds=expires_seconds),
response_headers={"response-content-disposition": disposition},
)
return self._rewrite_presigned_public_url(url)
def get_storage() -> BaseStorage:
"""获取存储实例"""
storage_type = settings.STORAGE_TYPE.lower()
if storage_type == "qiniu":
return QiniuStorage()
elif storage_type == "oss":
return OSSStorage()
elif storage_type == "minio":
return MinIOStorage()
else:
return LocalStorage()
# 全局存储实例
storage = get_storage()