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"]
|