142 lines
4.6 KiB
Python
142 lines
4.6 KiB
Python
|
"""对象存储适配器（交付报告 P0-2）。
|
|||
|
|
|
|||
|
|
设计：统一接口 save/read/delete/exists，两个实现：
|
|||
|
|
- local（默认）：磁盘目录，行为与现状一致；
|
|||
|
|
- s3（可选）：通过环境变量启用（YAWEN_STORAGE_BACKEND=s3 + endpoint/key/secret/bucket），
|
|||
|
|
依赖 boto3（未安装时导入报错并回退 local，错误信息可读）。
|
|||
|
|
|
|||
|
|
产物路径统一用 "local://相对路径" / "s3://bucket/key" 形式存库，
|
|||
|
|
读取时由 resolve_path 还原为可读的绝对路径或 s3 引用。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import shutil
|
|||
|
|
from pathlib import Path
|
|||
|
|
from urllib.parse import urlparse
|
|||
|
|
|
|||
|
|
from .config import DATA_DIR
|
|||
|
|
|
|||
|
|
|
|||
|
|
class StorageBackend:
|
|||
|
|
def save(self, key: str, data: bytes) -> str: ...
|
|||
|
|
def read(self, key: str) -> bytes: ...
|
|||
|
|
def delete(self, key: str) -> None: ...
|
|||
|
|
def exists(self, key: str) -> bool: ...
|
|||
|
|
|
|||
|
|
|
|||
|
|
class LocalStorage(StorageBackend):
|
|||
|
|
"""磁盘实现：key 是相对 DATA_DIR/storage 的路径。"""
|
|||
|
|
|
|||
|
|
def __init__(self, root: Path | None = None):
|
|||
|
|
self.root = Path(root) if root else DATA_DIR / "storage"
|
|||
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
def _path(self, key: str) -> Path:
|
|||
|
|
clean = key.lstrip("/").replace("..", "_")
|
|||
|
|
return self.root / clean
|
|||
|
|
|
|||
|
|
def save(self, key: str, data: bytes) -> str:
|
|||
|
|
path = self._path(key)
|
|||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|||
|
|
tmp = path.with_suffix(path.suffix + ".tmp")
|
|||
|
|
tmp.write_bytes(data)
|
|||
|
|
tmp.replace(path)
|
|||
|
|
# 引用里不回传原始 key 的 ..，写库引用使用净化后的路径
|
|||
|
|
clean_key = str(path.relative_to(self.root))
|
|||
|
|
return f"local://{clean_key}"
|
|||
|
|
|
|||
|
|
def read(self, key: str) -> bytes:
|
|||
|
|
return self._path(key).read_bytes()
|
|||
|
|
|
|||
|
|
def delete(self, key: str) -> None:
|
|||
|
|
path = self._path(key)
|
|||
|
|
if path.exists():
|
|||
|
|
path.unlink()
|
|||
|
|
|
|||
|
|
def exists(self, key: str) -> bool:
|
|||
|
|
return self._path(key).exists()
|
|||
|
|
|
|||
|
|
def read_text(self, key: str) -> str:
|
|||
|
|
return self.read(key).decode("utf-8", errors="ignore")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class S3Storage(StorageBackend):
|
|||
|
|
"""S3/MinIO/COS 实现：懒加载 boto3，未配置时抛可读错误。"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
try:
|
|||
|
|
import boto3 # noqa: F401
|
|||
|
|
except ImportError as exc:
|
|||
|
|
raise RuntimeError(
|
|||
|
|
"S3 存储后端需要 boto3：请在 server 环境执行 pip install boto3") from exc
|
|||
|
|
self.bucket = os.environ.get("YAWEN_S3_BUCKET", "")
|
|||
|
|
if not self.bucket:
|
|||
|
|
raise RuntimeError("S3 存储后端需要环境变量 YAWEN_S3_BUCKET")
|
|||
|
|
self.client = boto3.client(
|
|||
|
|
"s3",
|
|||
|
|
endpoint_url=os.environ.get("YAWEN_S3_ENDPOINT") or None,
|
|||
|
|
aws_access_key_id=os.environ.get("YAWEN_S3_ACCESS_KEY") or None,
|
|||
|
|
aws_secret_access_key=os.environ.get("YAWEN_S3_SECRET_KEY") or None,
|
|||
|
|
region_name=os.environ.get("YAWEN_S3_REGION") or None,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def save(self, key: str, data: bytes) -> str:
|
|||
|
|
self.client.put_object(Bucket=self.bucket, Key=key, Body=data)
|
|||
|
|
return f"s3://{self.bucket}/{key}"
|
|||
|
|
|
|||
|
|
def read(self, key: str) -> bytes:
|
|||
|
|
return self.client.get_object(Bucket=self.bucket, Key=key)["Body"].read()
|
|||
|
|
|
|||
|
|
def delete(self, key: str) -> None:
|
|||
|
|
self.client.delete_object(Bucket=self.bucket, Key=key)
|
|||
|
|
|
|||
|
|
def exists(self, key: str) -> bool:
|
|||
|
|
try:
|
|||
|
|
self.client.head_object(Bucket=self.bucket, Key=key)
|
|||
|
|
return True
|
|||
|
|
except Exception:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
|
|||
|
|
_backend: StorageBackend | None = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_storage() -> StorageBackend:
|
|||
|
|
global _backend
|
|||
|
|
if _backend is None:
|
|||
|
|
backend_name = os.environ.get("YAWEN_STORAGE_BACKEND", "local").lower()
|
|||
|
|
if backend_name == "s3":
|
|||
|
|
_backend = S3Storage()
|
|||
|
|
else:
|
|||
|
|
_backend = LocalStorage()
|
|||
|
|
return _backend
|
|||
|
|
|
|||
|
|
|
|||
|
|
def reset_storage() -> None:
|
|||
|
|
"""测试用：清空缓存的 backend 单例。"""
|
|||
|
|
global _backend
|
|||
|
|
_backend = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_ref(ref: str) -> tuple[str, str]:
|
|||
|
|
"""'s3://bucket/key' → ('s3', 'bucket/key')；'local://a/b' → ('local', 'a/b')；
|
|||
|
|
裸路径视为 local。"""
|
|||
|
|
ref = str(ref or "")
|
|||
|
|
if "://" in ref:
|
|||
|
|
scheme, rest = ref.split("://", 1)
|
|||
|
|
return scheme, rest
|
|||
|
|
return "local", ref
|
|||
|
|
|
|||
|
|
|
|||
|
|
def resolve_display_path(ref: str) -> str:
|
|||
|
|
"""把存库引用转成可展示/下载用的路径。"""
|
|||
|
|
scheme, key = parse_ref(ref)
|
|||
|
|
if scheme == "s3":
|
|||
|
|
return ref
|
|||
|
|
return str(DATA_DIR / "storage" / key.lstrip("/").replace("..", "_"))
|
|||
|
|
|
|||
|
|
|
|||
|
|
__all__ = ["StorageBackend", "LocalStorage", "S3Storage", "get_storage",
|
|||
|
|
"reset_storage", "parse_ref", "resolve_display_path"]
|