storage: extract unified AsyncStorageBackend abstraction + migrate from RustFS
Replace the old RustFS-specific storage layer (common.storage.client /
RustFSObjectStore) with a minimal sync/async abstraction:
AsyncStorageBackend: put / get / get_stream / delete / exists / stat /
list / get_url / copy
StorageBackend: same surface, sync implementations
create_storage({"type": "s3" | "local", "mode": "async", ...})
backends/s3.py: S3-compatible (boto3 / aioboto3)
backends/local.py: on-disk filesystem (aiofiles)
Concretely:
- Drop RustFSObjectStore + common.storage.client (deleted).
- Drop the RustFS-specific ensure_bucket / presign_put / move_to_trash /
rewrite_to_public_path / sha256 / put_bytes methods.
- Migrate backend/storage_api.py + backend/main.py + backend/scripts.py
+ schedule/service.py + schedule/worker.py to the new abstraction.
- Migrate backend/storage_client.py + schedule/storage_client.py to
stub status (HTTP wrapper is dead code post-migration; rewrite pending).
- Rename all RUSTFS_* env vars to S3_* across .env.example,
docker-compose.yml, default.conf, scripts/nginx-entrypoint.sh,
common/config.py.
- Replace hardcoded rclone remote name "rustfs" with "s3" in
docker-compose.yml + config.py default.
- Rename "rustfs" SQLAlchemy column comments + table comments to
provider-neutral wording; StorageObjects.storage_backend enum
value moves from "rustfs" to "s3" (DB rows with the old value will
fail the != "s3" check until a one-shot migration is applied).
- Drop unused common/src/common/migrations/{README,env.py,script.py.mako}
(the alembic setup lives in /migrations/, not here).
Migration of the old abstractions has been done in one pass; per-route
method calls (delete / stat / put / get_url) are now direct one-liners
against AsyncStorageBackend.
After this commit:
- All Python imports resolve; routes compile (compileall green).
- s3 mode is fully wired.
- Routes that depended on removed methods (presign_put, move_to_trash,
rewrite_to_public_path, head() metadata) raise NotImplementedError
with a one-line TODO; rewriting these route handlers is the next step.
This commit is contained in:
tao.chen
+35
-33
@@ -59,35 +59,57 @@ class Settings(BaseSettings):
|
||||
description="Backend → rclone RC HTTP endpoint (VFS cache invalidation).",
|
||||
)
|
||||
|
||||
# ── RustFS object storage ────────────────────────────────────
|
||||
rustfs_endpoint: str = Field(
|
||||
default="http://rustfs:9000",
|
||||
description="S3 endpoint for the RustFS upstream.",
|
||||
# ── object storage backend selection ─────────────────────────
|
||||
storage_backend: str = Field(
|
||||
default="s3",
|
||||
description=(
|
||||
"Which storage backend the deployment uses. ``s3`` (default) "
|
||||
"reads the ``s3_*`` settings and connects to an S3-compatible "
|
||||
"service. ``local`` uses on-disk filesystems under "
|
||||
"``local_storage_base_dir`` — useful for dev / single-node / "
|
||||
"air-gapped deployments."
|
||||
),
|
||||
)
|
||||
rustfs_access_key: str = Field(
|
||||
local_storage_base_dir: str = Field(
|
||||
default="/data",
|
||||
description=(
|
||||
"Root directory for the local-filesystem storage backend. The 4 "
|
||||
"buckets become subdirectories: ``<root>/workspace``, "
|
||||
"``<root>/version``, ``<root>/run_log``, ``<root>/trash``. "
|
||||
"Default ``/data``; this directory must be a shared Docker "
|
||||
"volume between the backend and runtime containers in local mode."
|
||||
),
|
||||
)
|
||||
|
||||
# ── S3-compatible object storage ─────────────────────────────
|
||||
s3_endpoint: str = Field(
|
||||
default="http://s3:9000",
|
||||
description="S3 endpoint for the object-storage upstream.",
|
||||
)
|
||||
s3_access_key: str = Field(
|
||||
default="modelplatform",
|
||||
description="boto3 access key for RustFS.",
|
||||
description="boto3 access key for S3-compatible storage.",
|
||||
)
|
||||
rustfs_secret_key: str = Field(
|
||||
s3_secret_key: str = Field(
|
||||
default="modelplatformsecret",
|
||||
description="boto3 secret key for RustFS.",
|
||||
description="boto3 secret key for S3-compatible storage.",
|
||||
)
|
||||
rustfs_workspace_bucket: str = Field(
|
||||
s3_workspace_bucket: str = Field(
|
||||
default="workspaces",
|
||||
description=(
|
||||
"Bucket for workspace files (notebooks / scripts / working "
|
||||
"copies). Layout: s3://<bucket>/<workspace_id>/..."
|
||||
),
|
||||
)
|
||||
rustfs_version_bucket: str = Field(
|
||||
s3_version_bucket: str = Field(
|
||||
default="versions",
|
||||
description="Bucket for immutable script-version artifacts.",
|
||||
)
|
||||
rustfs_run_log_bucket: str = Field(
|
||||
s3_run_log_bucket: str = Field(
|
||||
default="run-logs",
|
||||
description="Bucket for schedule run logs.",
|
||||
)
|
||||
rustfs_trash_bucket: str = Field(
|
||||
s3_trash_bucket: str = Field(
|
||||
default="trash",
|
||||
description=(
|
||||
"Bucket for soft-deleted objects. The source bucket key is "
|
||||
@@ -95,7 +117,7 @@ class Settings(BaseSettings):
|
||||
"Trash is reaped on a schedule out of band."
|
||||
),
|
||||
)
|
||||
rustfs_trash_retention_days: int = Field(
|
||||
s3_trash_retention_days: int = Field(
|
||||
default=30,
|
||||
description=(
|
||||
"How long a trashed object is retained before reaping. "
|
||||
@@ -104,26 +126,6 @@ class Settings(BaseSettings):
|
||||
),
|
||||
)
|
||||
|
||||
# ── local FS roots ────────────────────────────────────────────
|
||||
workspace_root: str = Field(
|
||||
default="/app/workspaces",
|
||||
description=(
|
||||
"Schedule subprocess cwd; staging area for notebook_runner. "
|
||||
"Backend no longer writes here — workspace files live in "
|
||||
"RustFS via rustfs_workspace_bucket."
|
||||
),
|
||||
)
|
||||
workspaces_root: str = Field(
|
||||
default="/app/workspaces",
|
||||
description="Runtime rclone FUSE mount point for the workspace bucket.",
|
||||
)
|
||||
|
||||
# ── rclone remote spec ────────────────────────────────────────
|
||||
remote_bucket: str = Field(
|
||||
default="rustfs:workspaces",
|
||||
description="rclone remote spec for the workspace bucket.",
|
||||
)
|
||||
|
||||
# ── schedule → backend API ────────────────────────────────────
|
||||
backend_api_url: str = Field(
|
||||
default="http://backend:8000",
|
||||
|
||||
@@ -84,7 +84,7 @@ class Versions(Base):
|
||||
CHAR(26), nullable=False, comment="发布时的源对象"
|
||||
)
|
||||
artifact_object_id: Mapped[str] = mapped_column(
|
||||
CHAR(26), nullable=False, comment="RustFS 不可变版本制品"
|
||||
CHAR(26), nullable=False, comment="S3 不可变版本制品"
|
||||
)
|
||||
version_no: Mapped[int] = mapped_column(INTEGER, nullable=False)
|
||||
version_label: Mapped[str] = mapped_column(
|
||||
|
||||
@@ -201,6 +201,25 @@ class UploadSessions(Base):
|
||||
content_type: Mapped[Optional[str]] = mapped_column(String(255))
|
||||
storage_object_id: Mapped[Optional[str]] = mapped_column(CHAR(26))
|
||||
completed_at: Mapped[Optional[datetime.datetime]] = mapped_column(DATETIME(fsp=3))
|
||||
# Metadata persisted at session creation so step 2 (PUT bytes) can build
|
||||
# the StorageObjects row without re-sending them. Replaces the
|
||||
# CompleteUploadRequest payload that lived between presign-PUT and head().
|
||||
file_name: Mapped[str] = mapped_column(String(255), nullable=False, server_default="")
|
||||
usage_type: Mapped[str] = mapped_column(
|
||||
String(32),
|
||||
nullable=False,
|
||||
server_default=text("'working_copy'"),
|
||||
comment="data_resource/version_artifact/snapshot/run_log/run_result/working_copy/public_script",
|
||||
)
|
||||
visibility: Mapped[str] = mapped_column(
|
||||
String(16),
|
||||
nullable=False,
|
||||
server_default=text("'private'"),
|
||||
comment="private/workspace/public",
|
||||
)
|
||||
is_immutable: Mapped[int] = mapped_column(
|
||||
TINYINT(1), nullable=False, server_default=text("0")
|
||||
)
|
||||
is_deleted: Mapped[int] = mapped_column(
|
||||
TINYINT(1), nullable=False, server_default=text("0")
|
||||
)
|
||||
|
||||
@@ -46,10 +46,10 @@ class Workspaces(Base):
|
||||
)
|
||||
description: Mapped[Optional[str]] = mapped_column(String(1000))
|
||||
artifact_bucket: Mapped[Optional[str]] = mapped_column(
|
||||
String(128), comment="RustFS bucket"
|
||||
String(128), comment="S3 bucket"
|
||||
)
|
||||
artifact_prefix: Mapped[Optional[str]] = mapped_column(
|
||||
String(512), comment="RustFS object key prefix"
|
||||
String(512), comment="S3 object key prefix"
|
||||
)
|
||||
is_deleted: Mapped[int] = mapped_column(
|
||||
TINYINT(1), nullable=False, server_default=text("0")
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
Generic single-database configuration.
|
||||
@@ -1,78 +0,0 @@
|
||||
from logging.config import fileConfig
|
||||
|
||||
from sqlalchemy import engine_from_config
|
||||
from sqlalchemy import pool
|
||||
|
||||
from alembic import context
|
||||
from common.db.models import Base
|
||||
# this is the Alembic Config object, which provides
|
||||
# access to the values within the .ini file in use.
|
||||
config = context.config
|
||||
|
||||
# Interpret the config file for Python logging.
|
||||
# This line sets up loggers basically.
|
||||
if config.config_file_name is not None:
|
||||
fileConfig(config.config_file_name)
|
||||
|
||||
# add your model's MetaData object here
|
||||
# for 'autogenerate' support
|
||||
# from myapp import mymodel
|
||||
# target_metadata = mymodel.Base.metadata
|
||||
target_metadata = Base.metadata
|
||||
|
||||
# other values from the config, defined by the needs of env.py,
|
||||
# can be acquired:
|
||||
# my_important_option = config.get_main_option("my_important_option")
|
||||
# ... etc.
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
"""Run migrations in 'offline' mode.
|
||||
|
||||
This configures the context with just a URL
|
||||
and not an Engine, though an Engine is acceptable
|
||||
here as well. By skipping the Engine creation
|
||||
we don't even need a DBAPI to be available.
|
||||
|
||||
Calls to context.execute() here emit the given string to the
|
||||
script output.
|
||||
|
||||
"""
|
||||
url = config.get_main_option("sqlalchemy.url")
|
||||
context.configure(
|
||||
url=url,
|
||||
target_metadata=target_metadata,
|
||||
literal_binds=True,
|
||||
dialect_opts={"paramstyle": "named"},
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
def run_migrations_online() -> None:
|
||||
"""Run migrations in 'online' mode.
|
||||
|
||||
In this scenario we need to create an Engine
|
||||
and associate a connection with the context.
|
||||
|
||||
"""
|
||||
connectable = engine_from_config(
|
||||
config.get_section(config.config_ini_section, {}),
|
||||
prefix="sqlalchemy.",
|
||||
poolclass=pool.NullPool,
|
||||
)
|
||||
|
||||
with connectable.connect() as connection:
|
||||
context.configure(
|
||||
connection=connection, target_metadata=target_metadata
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
|
||||
if context.is_offline_mode():
|
||||
run_migrations_offline()
|
||||
else:
|
||||
run_migrations_online()
|
||||
@@ -1,28 +0,0 @@
|
||||
"""${message}
|
||||
|
||||
Revision ID: ${up_revision}
|
||||
Revises: ${down_revision | comma,n}
|
||||
Create Date: ${create_date}
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
${imports if imports else ""}
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)}
|
||||
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Upgrade schema."""
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Downgrade schema."""
|
||||
${downgrades if downgrades else "pass"}
|
||||
@@ -1,6 +1,44 @@
|
||||
"""Storage building blocks shared by backend and schedule services."""
|
||||
"""统一存储层，同时支持同步和异步，通过 config["mode"] 切换。
|
||||
|
||||
from common.storage.client import StorageClient
|
||||
from common.storage.rustfs import RustFSObjectStore
|
||||
对上层暴露的公开 API：
|
||||
|
||||
__all__ = ["RustFSObjectStore", "StorageClient"]
|
||||
from storage import create_storage, StorageBackend, AsyncStorageBackend, ObjectMeta
|
||||
from storage.exceptions import StorageError, StorageNotFoundError, ...
|
||||
|
||||
用法：
|
||||
# 同步（默认 mode="sync"）
|
||||
storage = create_storage({"type": "local", "base_dir": "./data"})
|
||||
storage.put("a/b.txt", b"hello")
|
||||
|
||||
# 异步：加一个 mode 字段
|
||||
storage = create_storage({"type": "local", "mode": "async", "base_dir": "./data"})
|
||||
await storage.put("a/b.txt", b"hello")
|
||||
|
||||
切换本地/S3，或切换同步/异步，业务代码都不用改，只改配置：
|
||||
storage = create_storage({"type": "s3", "mode": "async", "bucket": "my-bucket"})
|
||||
"""
|
||||
|
||||
from .base import AsyncStorageBackend, ObjectMeta, StorageBackend
|
||||
from .factory import (
|
||||
PURPOSE_BUCKETS,
|
||||
RCLONE_REMOTE_NAME,
|
||||
build_storage_config,
|
||||
create_storage,
|
||||
rclone_remote_spec,
|
||||
workspaces_root,
|
||||
)
|
||||
from .registry import register_backend, registered_backends
|
||||
|
||||
__all__ = [
|
||||
"create_storage",
|
||||
"build_storage_config",
|
||||
"workspaces_root",
|
||||
"rclone_remote_spec",
|
||||
"RCLONE_REMOTE_NAME",
|
||||
"PURPOSE_BUCKETS",
|
||||
"StorageBackend",
|
||||
"AsyncStorageBackend",
|
||||
"ObjectMeta",
|
||||
"register_backend",
|
||||
"registered_backends",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
"""导入本模块即可触发所有内置后端的 @register_backend 注册。
|
||||
|
||||
新增内置后端时，在这里加一行 import 即可；
|
||||
如果是第三方/业务自己的后端，不需要改这个文件，
|
||||
只要在使用前 import 一次那个模块（让装饰器执行）就够了。
|
||||
"""
|
||||
|
||||
from . import local # noqa: F401
|
||||
from . import s3 # noqa: F401
|
||||
@@ -0,0 +1,245 @@
|
||||
"""本地文件系统存储后端。
|
||||
|
||||
- 同步实现 `LocalStorageBackend`：标准库文件 I/O
|
||||
- 异步实现 `LocalAsyncStorageBackend`：aiofiles 做实际读写，
|
||||
stat/exists/delete/mkdir/目录遍历这类轻量元数据操作用
|
||||
asyncio.to_thread 包一层，避免阻塞事件循环
|
||||
（只有创建异步实例时才需要装 aiofiles，同步实现零依赖）
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import AsyncIterator, BinaryIO, Iterable, Optional
|
||||
|
||||
from ..base import AsyncData, AsyncStorageBackend, ObjectMeta, StorageBackend, SyncData
|
||||
from ..exceptions import StorageAlreadyExistsError, StorageNotFoundError
|
||||
from ..registry import register_backend
|
||||
|
||||
|
||||
def _resolve(base_dir: Path, key: str) -> Path:
|
||||
key = key.strip("/")
|
||||
path = (base_dir / key).resolve()
|
||||
if base_dir not in path.parents and path != base_dir:
|
||||
raise ValueError(f"非法 key，路径穿越到 base_dir 之外: {key!r}")
|
||||
return path
|
||||
|
||||
|
||||
def _meta(key: str, path: Path) -> ObjectMeta:
|
||||
st = path.stat()
|
||||
return ObjectMeta(key=key, size=st.st_size, last_modified=st.st_mtime)
|
||||
|
||||
|
||||
# ==================== 同步实现 ====================
|
||||
|
||||
|
||||
@register_backend("local", mode="sync")
|
||||
class LocalStorageBackend(StorageBackend):
|
||||
"""配置示例: {"type": "local", "mode": "sync", "base_dir": "/data/storage"}"""
|
||||
|
||||
def __init__(self, base_dir: str, **_ignored):
|
||||
self.base_dir = Path(base_dir).resolve()
|
||||
self.base_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _resolve(self, key: str) -> Path:
|
||||
return _resolve(self.base_dir, key)
|
||||
|
||||
def put(
|
||||
self,
|
||||
key: str,
|
||||
data: SyncData,
|
||||
*,
|
||||
overwrite: bool = True,
|
||||
content_type: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> ObjectMeta:
|
||||
path = self._resolve(key)
|
||||
if path.exists() and not overwrite:
|
||||
raise StorageAlreadyExistsError(f"key 已存在: {key}")
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if isinstance(data, bytes):
|
||||
path.write_bytes(data)
|
||||
else:
|
||||
with open(path, "wb") as f:
|
||||
shutil.copyfileobj(data, f)
|
||||
# local FS 没有对象级 metadata；content_type / metadata 暂存忽略。
|
||||
return _meta(key, path)
|
||||
|
||||
def get(self, key: str) -> bytes:
|
||||
path = self._resolve(key)
|
||||
if not path.is_file():
|
||||
raise StorageNotFoundError(f"key 不存在: {key}")
|
||||
return path.read_bytes()
|
||||
|
||||
def get_stream(self, key: str) -> BinaryIO:
|
||||
path = self._resolve(key)
|
||||
if not path.is_file():
|
||||
raise StorageNotFoundError(f"key 不存在: {key}")
|
||||
return open(path, "rb")
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
try:
|
||||
self._resolve(key).unlink()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
return self._resolve(key).is_file()
|
||||
|
||||
def stat(self, key: str) -> ObjectMeta:
|
||||
path = self._resolve(key)
|
||||
if not path.is_file():
|
||||
raise StorageNotFoundError(f"key 不存在: {key}")
|
||||
return _meta(key, path)
|
||||
|
||||
def list(self, prefix: str = "") -> Iterable[ObjectMeta]:
|
||||
search_root = self._resolve(prefix) if prefix else self.base_dir
|
||||
if search_root.is_dir():
|
||||
candidates = search_root.rglob("*")
|
||||
else:
|
||||
candidates = search_root.parent.glob(f"{search_root.name}*")
|
||||
|
||||
for path in candidates:
|
||||
if path.is_file():
|
||||
key = str(path.relative_to(self.base_dir)).replace(os.sep, "/")
|
||||
yield _meta(key, path)
|
||||
|
||||
def get_url(self, key: str, *, expires_in: Optional[timedelta] = None) -> str:
|
||||
path = self._resolve(key)
|
||||
if not path.is_file():
|
||||
raise StorageNotFoundError(f"key 不存在: {key}")
|
||||
return path.as_uri()
|
||||
|
||||
def copy(self, src_key: str, dst_key: str) -> ObjectMeta:
|
||||
src_path = self._resolve(src_key)
|
||||
if not src_path.is_file():
|
||||
raise StorageNotFoundError(f"key 不存在: {src_key}")
|
||||
dst_path = self._resolve(dst_key)
|
||||
dst_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(src_path, dst_path)
|
||||
return _meta(dst_key, dst_path)
|
||||
|
||||
|
||||
# ==================== 异步实现 ====================
|
||||
|
||||
|
||||
@register_backend("local", mode="async")
|
||||
class LocalAsyncStorageBackend(AsyncStorageBackend):
|
||||
"""配置示例: {"type": "local", "mode": "async", "base_dir": "/data/storage"}
|
||||
|
||||
需要: pip install aiofiles
|
||||
"""
|
||||
|
||||
def __init__(self, base_dir: str, **_ignored):
|
||||
self.base_dir = Path(base_dir).resolve()
|
||||
self.base_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _resolve(self, key: str) -> Path:
|
||||
return _resolve(self.base_dir, key)
|
||||
|
||||
async def put(
|
||||
self,
|
||||
key: str,
|
||||
data: AsyncData,
|
||||
*,
|
||||
overwrite: bool = True,
|
||||
content_type: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> ObjectMeta:
|
||||
import aiofiles
|
||||
|
||||
path = self._resolve(key)
|
||||
if not overwrite and await asyncio.to_thread(path.exists):
|
||||
raise StorageAlreadyExistsError(f"key 已存在: {key}")
|
||||
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
|
||||
|
||||
async with aiofiles.open(path, "wb") as f:
|
||||
if isinstance(data, (bytes, bytearray)):
|
||||
await f.write(data)
|
||||
else:
|
||||
async for chunk in data:
|
||||
await f.write(chunk)
|
||||
|
||||
# local FS 没有对象级 metadata；content_type / metadata 暂存忽略。
|
||||
return await asyncio.to_thread(_meta, key, path)
|
||||
|
||||
async def get(self, key: str) -> bytes:
|
||||
import aiofiles
|
||||
|
||||
path = self._resolve(key)
|
||||
if not await asyncio.to_thread(path.is_file):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}")
|
||||
async with aiofiles.open(path, "rb") as f:
|
||||
return await f.read()
|
||||
|
||||
def get_stream(self, key: str, chunk_size: int = 64 * 1024) -> AsyncIterator[bytes]:
|
||||
async def _iter():
|
||||
import aiofiles
|
||||
|
||||
path = self._resolve(key)
|
||||
if not await asyncio.to_thread(path.is_file):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}")
|
||||
async with aiofiles.open(path, "rb") as f:
|
||||
while True:
|
||||
chunk = await f.read(chunk_size)
|
||||
if not chunk:
|
||||
break
|
||||
yield chunk
|
||||
|
||||
return _iter()
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
path = self._resolve(key)
|
||||
|
||||
def _unlink():
|
||||
try:
|
||||
path.unlink()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
await asyncio.to_thread(_unlink)
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
return await asyncio.to_thread(self._resolve(key).is_file)
|
||||
|
||||
async def stat(self, key: str) -> ObjectMeta:
|
||||
path = self._resolve(key)
|
||||
if not await asyncio.to_thread(path.is_file):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}")
|
||||
return await asyncio.to_thread(_meta, key, path)
|
||||
|
||||
def list(self, prefix: str = "") -> AsyncIterator[ObjectMeta]:
|
||||
async def _iter():
|
||||
search_root = self._resolve(prefix) if prefix else self.base_dir
|
||||
|
||||
def _collect():
|
||||
if search_root.is_dir():
|
||||
candidates = list(search_root.rglob("*"))
|
||||
else:
|
||||
candidates = list(search_root.parent.glob(f"{search_root.name}*"))
|
||||
return [p for p in candidates if p.is_file()]
|
||||
|
||||
files = await asyncio.to_thread(_collect)
|
||||
for path in files:
|
||||
key = str(path.relative_to(self.base_dir)).replace(os.sep, "/")
|
||||
yield await asyncio.to_thread(_meta, key, path)
|
||||
|
||||
return _iter()
|
||||
|
||||
async def get_url(self, key: str, *, expires_in: Optional[timedelta] = None) -> str:
|
||||
path = self._resolve(key)
|
||||
if not await asyncio.to_thread(path.is_file):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}")
|
||||
return path.as_uri()
|
||||
|
||||
async def copy(self, src_key: str, dst_key: str) -> ObjectMeta:
|
||||
src_path = self._resolve(src_key)
|
||||
if not await asyncio.to_thread(src_path.is_file):
|
||||
raise StorageNotFoundError(f"key 不存在: {src_key}")
|
||||
dst_path = self._resolve(dst_key)
|
||||
await asyncio.to_thread(dst_path.parent.mkdir, parents=True, exist_ok=True)
|
||||
await asyncio.to_thread(shutil.copy2, src_path, dst_path)
|
||||
return await asyncio.to_thread(_meta, dst_key, dst_path)
|
||||
@@ -0,0 +1,420 @@
|
||||
"""S3（及兼容协议）存储后端。
|
||||
|
||||
- 同步实现 `S3StorageBackend`：boto3
|
||||
- 异步实现 `S3AsyncStorageBackend`：aioboto3
|
||||
|
||||
两者只在各自 __init__ 里做 lazy import，互不强制依赖：
|
||||
只用同步模式不需要装 aioboto3，只用异步模式不需要额外装 boto3
|
||||
（aioboto3 本身依赖 botocore，异常类型从它里面拿）。
|
||||
"""
|
||||
|
||||
from datetime import timedelta
|
||||
from typing import AsyncIterator, BinaryIO, Iterable, Optional
|
||||
|
||||
from ..base import AsyncData, AsyncStorageBackend, ObjectMeta, StorageBackend, SyncData
|
||||
from ..exceptions import (
|
||||
StorageAlreadyExistsError,
|
||||
StorageConnectionError,
|
||||
StorageNotFoundError,
|
||||
)
|
||||
from ..registry import register_backend
|
||||
|
||||
|
||||
def _meta_from_head(key: str, head: dict) -> ObjectMeta:
|
||||
return ObjectMeta(
|
||||
key=key,
|
||||
size=head.get("ContentLength", 0),
|
||||
last_modified=head["LastModified"].timestamp() if head.get("LastModified") else None,
|
||||
etag=head.get("ETag"),
|
||||
)
|
||||
|
||||
|
||||
# ==================== 同步实现 ====================
|
||||
|
||||
|
||||
@register_backend("s3", mode="sync")
|
||||
class S3StorageBackend(StorageBackend):
|
||||
"""配置示例:
|
||||
{
|
||||
"type": "s3", "mode": "sync",
|
||||
"bucket": "my-bucket", "prefix": "app1/",
|
||||
"region_name": "cn-north-1", "endpoint_url": "https://s3.example.com",
|
||||
"aws_access_key_id": "...", "aws_secret_access_key": "...",
|
||||
}
|
||||
|
||||
需要: pip install boto3
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bucket: str,
|
||||
prefix: str = "",
|
||||
region_name: Optional[str] = None,
|
||||
endpoint_url: Optional[str] = None,
|
||||
aws_access_key_id: Optional[str] = None,
|
||||
aws_secret_access_key: Optional[str] = None,
|
||||
**_ignored,
|
||||
):
|
||||
try:
|
||||
import boto3
|
||||
from botocore.exceptions import BotoCoreError, ClientError
|
||||
except ImportError as e:
|
||||
raise ImportError("使用同步 S3 存储后端需要先安装 boto3: pip install boto3") from e
|
||||
|
||||
self._ClientError = ClientError
|
||||
self._BotoCoreError = BotoCoreError
|
||||
self.bucket = bucket
|
||||
self.prefix = prefix.strip("/") + "/" if prefix.strip("/") else ""
|
||||
|
||||
try:
|
||||
self.client = boto3.client(
|
||||
"s3",
|
||||
region_name=region_name,
|
||||
endpoint_url=endpoint_url,
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
)
|
||||
except (BotoCoreError, ClientError) as e:
|
||||
raise StorageConnectionError(f"初始化 S3 client 失败: {e}") from e
|
||||
|
||||
def _full_key(self, key: str) -> str:
|
||||
return f"{self.prefix}{key.lstrip('/')}"
|
||||
|
||||
def put(self, key: str, data: SyncData, *, overwrite: bool = True) -> ObjectMeta:
|
||||
full_key = self._full_key(key)
|
||||
if not overwrite and self.exists(key):
|
||||
raise StorageAlreadyExistsError(f"key 已存在: {key}")
|
||||
body = data if isinstance(data, bytes) else data.read()
|
||||
try:
|
||||
self.client.put_object(Bucket=self.bucket, Key=full_key, Body=body)
|
||||
except (self._ClientError, self._BotoCoreError) as e:
|
||||
raise StorageConnectionError(f"上传失败 key={key}: {e}") from e
|
||||
return self.stat(key)
|
||||
|
||||
def get(self, key: str) -> bytes:
|
||||
full_key = self._full_key(key)
|
||||
try:
|
||||
resp = self.client.get_object(Bucket=self.bucket, Key=full_key)
|
||||
return resp["Body"].read()
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("NoSuchKey", "404"):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}") from e
|
||||
raise StorageConnectionError(f"读取失败 key={key}: {e}") from e
|
||||
|
||||
def get_stream(self, key: str) -> BinaryIO:
|
||||
full_key = self._full_key(key)
|
||||
try:
|
||||
resp = self.client.get_object(Bucket=self.bucket, Key=full_key)
|
||||
return resp["Body"]
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("NoSuchKey", "404"):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}") from e
|
||||
raise StorageConnectionError(f"读取失败 key={key}: {e}") from e
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
try:
|
||||
self.client.delete_object(Bucket=self.bucket, Key=self._full_key(key))
|
||||
except (self._ClientError, self._BotoCoreError) as e:
|
||||
raise StorageConnectionError(f"删除失败 key={key}: {e}") from e
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
try:
|
||||
self.client.head_object(Bucket=self.bucket, Key=self._full_key(key))
|
||||
return True
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("404", "NoSuchKey"):
|
||||
return False
|
||||
raise StorageConnectionError(f"检查 exists 失败 key={key}: {e}") from e
|
||||
|
||||
def stat(self, key: str) -> ObjectMeta:
|
||||
try:
|
||||
head = self.client.head_object(Bucket=self.bucket, Key=self._full_key(key))
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("404", "NoSuchKey"):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}") from e
|
||||
raise StorageConnectionError(f"获取元信息失败 key={key}: {e}") from e
|
||||
return _meta_from_head(key, head)
|
||||
|
||||
def list(self, prefix: str = "") -> Iterable[ObjectMeta]:
|
||||
full_prefix = self._full_key(prefix)
|
||||
paginator = self.client.get_paginator("list_objects_v2")
|
||||
try:
|
||||
for page in paginator.paginate(Bucket=self.bucket, Prefix=full_prefix):
|
||||
for obj in page.get("Contents", []):
|
||||
key = obj["Key"][len(self.prefix):] if self.prefix else obj["Key"]
|
||||
yield ObjectMeta(
|
||||
key=key,
|
||||
size=obj["Size"],
|
||||
last_modified=obj["LastModified"].timestamp(),
|
||||
etag=obj.get("ETag"),
|
||||
)
|
||||
except (self._ClientError, self._BotoCoreError) as e:
|
||||
raise StorageConnectionError(f"列举对象失败 prefix={prefix}: {e}") from e
|
||||
|
||||
def get_url(self, key: str, *, expires_in: Optional[timedelta] = None) -> str:
|
||||
expires_seconds = int(expires_in.total_seconds()) if expires_in else 3600
|
||||
try:
|
||||
return self.client.generate_presigned_url(
|
||||
"get_object",
|
||||
Params={"Bucket": self.bucket, "Key": self._full_key(key)},
|
||||
ExpiresIn=expires_seconds,
|
||||
)
|
||||
except (self._ClientError, self._BotoCoreError) as e:
|
||||
raise StorageConnectionError(f"生成预签名 URL 失败 key={key}: {e}") from e
|
||||
|
||||
def copy(self, src_key: str, dst_key: str) -> ObjectMeta:
|
||||
try:
|
||||
self.client.copy_object(
|
||||
Bucket=self.bucket,
|
||||
Key=self._full_key(dst_key),
|
||||
CopySource={"Bucket": self.bucket, "Key": self._full_key(src_key)},
|
||||
)
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("404", "NoSuchKey"):
|
||||
raise StorageNotFoundError(f"key 不存在: {src_key}") from e
|
||||
raise StorageConnectionError(f"复制失败 {src_key} -> {dst_key}: {e}") from e
|
||||
return self.stat(dst_key)
|
||||
|
||||
|
||||
# ==================== 异步实现 ====================
|
||||
|
||||
|
||||
@register_backend("s3", mode="async")
|
||||
class S3AsyncStorageBackend(AsyncStorageBackend):
|
||||
"""配置示例同上，把 "mode" 改成 "async" 即可。
|
||||
|
||||
需要: pip install aioboto3
|
||||
|
||||
每次操作默认通过 `async with session.client(...)` 拿一个短生命周期
|
||||
client；用 `async with create_storage(...) as storage:` 可以复用同一个
|
||||
client（见 __aenter__/__aexit__）。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bucket: str,
|
||||
prefix: str = "",
|
||||
region_name: Optional[str] = None,
|
||||
endpoint_url: Optional[str] = None,
|
||||
aws_access_key_id: Optional[str] = None,
|
||||
aws_secret_access_key: Optional[str] = None,
|
||||
**_ignored,
|
||||
):
|
||||
try:
|
||||
import aioboto3
|
||||
from botocore.exceptions import BotoCoreError, ClientError
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"使用异步 S3 存储后端需要先安装 aioboto3: pip install aioboto3"
|
||||
) from e
|
||||
|
||||
self._ClientError = ClientError
|
||||
self._BotoCoreError = BotoCoreError
|
||||
self.bucket = bucket
|
||||
self.prefix = prefix.strip("/") + "/" if prefix.strip("/") else ""
|
||||
self._client_kwargs = dict(
|
||||
region_name=region_name,
|
||||
endpoint_url=endpoint_url,
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
)
|
||||
self._session = aioboto3.Session()
|
||||
self._persistent_client = None
|
||||
self._persistent_cm = None
|
||||
|
||||
def _client_cm(self):
|
||||
return self._session.client("s3", **self._client_kwargs)
|
||||
|
||||
async def __aenter__(self) -> "S3AsyncStorageBackend":
|
||||
self._persistent_cm = self._client_cm()
|
||||
self._persistent_client = await self._persistent_cm.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
if self._persistent_cm is not None:
|
||||
await self._persistent_cm.__aexit__(exc_type, exc, tb)
|
||||
self._persistent_cm = None
|
||||
self._persistent_client = None
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self.__aexit__(None, None, None)
|
||||
|
||||
def _full_key(self, key: str) -> str:
|
||||
return f"{self.prefix}{key.lstrip('/')}"
|
||||
|
||||
async def _run(self, coro_fn):
|
||||
if self._persistent_client is not None:
|
||||
return await coro_fn(self._persistent_client)
|
||||
async with self._client_cm() as client:
|
||||
return await coro_fn(client)
|
||||
|
||||
async def put(
|
||||
self,
|
||||
key: str,
|
||||
data: AsyncData,
|
||||
*,
|
||||
overwrite: bool = True,
|
||||
content_type: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> ObjectMeta:
|
||||
full_key = self._full_key(key)
|
||||
if not overwrite and await self.exists(key):
|
||||
raise StorageAlreadyExistsError(f"key 已存在: {key}")
|
||||
|
||||
if isinstance(data, (bytes, bytearray)):
|
||||
body = bytes(data)
|
||||
else:
|
||||
chunks = []
|
||||
async for chunk in data:
|
||||
chunks.append(chunk)
|
||||
body = b"".join(chunks)
|
||||
|
||||
# 过滤掉空 dict / None，避免 boto3 报 "parameter must be a non-empty
|
||||
# non-null dictionary of strings" 这种空请求参数错误。
|
||||
meta = {k: str(v) for k, v in (metadata or {}).items() if v is not None} or None
|
||||
|
||||
async def _op(client):
|
||||
kwargs = {"Bucket": self.bucket, "Key": full_key, "Body": body}
|
||||
if content_type:
|
||||
kwargs["ContentType"] = content_type
|
||||
if meta:
|
||||
kwargs["Metadata"] = meta
|
||||
await client.put_object(**kwargs)
|
||||
|
||||
try:
|
||||
await self._run(_op)
|
||||
except (self._ClientError, self._BotoCoreError) as e:
|
||||
raise StorageConnectionError(f"上传失败 key={key}: {e}") from e
|
||||
return await self.stat(key)
|
||||
|
||||
async def get(self, key: str) -> bytes:
|
||||
full_key = self._full_key(key)
|
||||
|
||||
async def _op(client):
|
||||
resp = await client.get_object(Bucket=self.bucket, Key=full_key)
|
||||
async with resp["Body"] as stream:
|
||||
return await stream.read()
|
||||
|
||||
try:
|
||||
return await self._run(_op)
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("NoSuchKey", "404"):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}") from e
|
||||
raise StorageConnectionError(f"读取失败 key={key}: {e}") from e
|
||||
|
||||
def get_stream(self, key: str, chunk_size: int = 64 * 1024) -> AsyncIterator[bytes]:
|
||||
async def _iter():
|
||||
full_key = self._full_key(key)
|
||||
|
||||
async def _op(client):
|
||||
return await client.get_object(Bucket=self.bucket, Key=full_key)
|
||||
|
||||
try:
|
||||
resp = await self._run(_op)
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("NoSuchKey", "404"):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}") from e
|
||||
raise StorageConnectionError(f"读取失败 key={key}: {e}") from e
|
||||
|
||||
async with resp["Body"] as stream:
|
||||
while True:
|
||||
chunk = await stream.read(chunk_size)
|
||||
if not chunk:
|
||||
break
|
||||
yield chunk
|
||||
|
||||
return _iter()
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
async def _op(client):
|
||||
await client.delete_object(Bucket=self.bucket, Key=self._full_key(key))
|
||||
|
||||
try:
|
||||
await self._run(_op)
|
||||
except (self._ClientError, self._BotoCoreError) as e:
|
||||
raise StorageConnectionError(f"删除失败 key={key}: {e}") from e
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
async def _op(client):
|
||||
await client.head_object(Bucket=self.bucket, Key=self._full_key(key))
|
||||
|
||||
try:
|
||||
await self._run(_op)
|
||||
return True
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("404", "NoSuchKey"):
|
||||
return False
|
||||
raise StorageConnectionError(f"检查 exists 失败 key={key}: {e}") from e
|
||||
|
||||
async def stat(self, key: str) -> ObjectMeta:
|
||||
async def _op(client):
|
||||
return await client.head_object(Bucket=self.bucket, Key=self._full_key(key))
|
||||
|
||||
try:
|
||||
head = await self._run(_op)
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("404", "NoSuchKey"):
|
||||
raise StorageNotFoundError(f"key 不存在: {key}") from e
|
||||
raise StorageConnectionError(f"获取元信息失败 key={key}: {e}") from e
|
||||
return _meta_from_head(key, head)
|
||||
|
||||
def list(self, prefix: str = "") -> AsyncIterator[ObjectMeta]:
|
||||
async def _iter():
|
||||
full_prefix = self._full_key(prefix)
|
||||
|
||||
async def _paginate(client):
|
||||
paginator = client.get_paginator("list_objects_v2")
|
||||
results = []
|
||||
async for page in paginator.paginate(Bucket=self.bucket, Prefix=full_prefix):
|
||||
for obj in page.get("Contents", []):
|
||||
key = obj["Key"][len(self.prefix):] if self.prefix else obj["Key"]
|
||||
results.append(
|
||||
ObjectMeta(
|
||||
key=key,
|
||||
size=obj["Size"],
|
||||
last_modified=obj["LastModified"].timestamp(),
|
||||
etag=obj.get("ETag"),
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
try:
|
||||
metas = await self._run(_paginate)
|
||||
except (self._ClientError, self._BotoCoreError) as e:
|
||||
raise StorageConnectionError(f"列举对象失败 prefix={prefix}: {e}") from e
|
||||
|
||||
for meta in metas:
|
||||
yield meta
|
||||
|
||||
return _iter()
|
||||
|
||||
async def get_url(self, key: str, *, expires_in: Optional[timedelta] = None) -> str:
|
||||
expires_seconds = int(expires_in.total_seconds()) if expires_in else 3600
|
||||
|
||||
async def _op(client):
|
||||
return await client.generate_presigned_url(
|
||||
"get_object",
|
||||
Params={"Bucket": self.bucket, "Key": self._full_key(key)},
|
||||
ExpiresIn=expires_seconds,
|
||||
)
|
||||
|
||||
try:
|
||||
return await self._run(_op)
|
||||
except (self._ClientError, self._BotoCoreError) as e:
|
||||
raise StorageConnectionError(f"生成预签名 URL 失败 key={key}: {e}") from e
|
||||
|
||||
async def copy(self, src_key: str, dst_key: str) -> ObjectMeta:
|
||||
async def _op(client):
|
||||
await client.copy_object(
|
||||
Bucket=self.bucket,
|
||||
Key=self._full_key(dst_key),
|
||||
CopySource={"Bucket": self.bucket, "Key": self._full_key(src_key)},
|
||||
)
|
||||
|
||||
try:
|
||||
await self._run(_op)
|
||||
except self._ClientError as e:
|
||||
if e.response.get("Error", {}).get("Code") in ("404", "NoSuchKey"):
|
||||
raise StorageNotFoundError(f"key 不存在: {src_key}") from e
|
||||
raise StorageConnectionError(f"复制失败 {src_key} -> {dst_key}: {e}") from e
|
||||
return await self.stat(dst_key)
|
||||
@@ -0,0 +1,155 @@
|
||||
"""同步 / 异步存储后端统一抽象接口。
|
||||
|
||||
`StorageBackend` 是同步接口，`AsyncStorageBackend` 是异步接口，
|
||||
两者共用同一个 `ObjectMeta` 数据结构，方法签名尽量保持对称
|
||||
（异步版本每个方法多一个 await，get_stream/list 变成异步生成器），
|
||||
这样业务代码从同步切到异步时心智负担最小。
|
||||
|
||||
上层通过 `storage.create_storage(config)` 统一创建实例，
|
||||
用 `config["mode"]` 决定拿到的是同步实现还是异步实现。
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import timedelta
|
||||
from typing import AsyncIterator, BinaryIO, Iterable, Optional, Union
|
||||
|
||||
SyncData = Union[bytes, BinaryIO]
|
||||
AsyncData = Union[bytes, "AsyncIterator[bytes]"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ObjectMeta:
|
||||
"""list/stat 等操作返回的对象元信息，做了跨后端的字段归一化。"""
|
||||
|
||||
key: str
|
||||
size: int
|
||||
last_modified: Optional[float] = None # unix timestamp
|
||||
etag: Optional[str] = None
|
||||
extra: dict = field(default_factory=dict) # 后端特有的额外信息
|
||||
|
||||
|
||||
class StorageBackend(ABC):
|
||||
"""同步存储后端统一抽象基类。"""
|
||||
|
||||
@abstractmethod
|
||||
def put(
|
||||
self,
|
||||
key: str,
|
||||
data: SyncData,
|
||||
*,
|
||||
overwrite: bool = True,
|
||||
content_type: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> ObjectMeta:
|
||||
"""写入对象。overwrite=False 时 key 已存在应抛出 StorageAlreadyExistsError。
|
||||
|
||||
``content_type`` 和 ``metadata`` 是可选的（与异步 put 语义一致）。
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, key: str) -> bytes:
|
||||
"""读取对象内容，不存在时抛出 StorageNotFoundError。"""
|
||||
|
||||
@abstractmethod
|
||||
def get_stream(self, key: str) -> BinaryIO:
|
||||
"""以流方式读取对象，适合大文件。"""
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, key: str) -> None:
|
||||
"""删除对象。删除不存在的 key 不应报错（幂等）。"""
|
||||
|
||||
@abstractmethod
|
||||
def exists(self, key: str) -> bool:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def stat(self, key: str) -> ObjectMeta:
|
||||
"""不存在时抛出 StorageNotFoundError。"""
|
||||
|
||||
@abstractmethod
|
||||
def list(self, prefix: str = "") -> Iterable[ObjectMeta]:
|
||||
"""按前缀列出对象。"""
|
||||
|
||||
@abstractmethod
|
||||
def get_url(self, key: str, *, expires_in: Optional[timedelta] = None) -> str:
|
||||
"""获取可访问 URL；本地存储返回 file://，S3 返回预签名 URL。"""
|
||||
|
||||
def copy(self, src_key: str, dst_key: str) -> ObjectMeta:
|
||||
"""默认实现：读出来再写进去。后端可覆盖为更高效的原生实现。"""
|
||||
data = self.get(src_key)
|
||||
return self.put(dst_key, data)
|
||||
|
||||
def close(self) -> None:
|
||||
"""释放后端持有的资源（连接池等）。不需要的后端可以不覆盖。"""
|
||||
return None
|
||||
|
||||
def __enter__(self) -> "StorageBackend":
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb) -> None:
|
||||
self.close()
|
||||
|
||||
|
||||
class AsyncStorageBackend(ABC):
|
||||
"""异步存储后端统一抽象基类。"""
|
||||
|
||||
@abstractmethod
|
||||
async def put(
|
||||
self,
|
||||
key: str,
|
||||
data: AsyncData,
|
||||
*,
|
||||
overwrite: bool = True,
|
||||
content_type: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> ObjectMeta:
|
||||
"""data 可以是 bytes，也可以是异步字节流（async generator）。
|
||||
|
||||
``content_type`` 和 ``metadata`` 是可选的：S3 后端会把它们分别透传
|
||||
成 ``ContentType`` 请求头和 ``Metadata`` dict；local 后端目前忽略
|
||||
这两个参数（本地 FS 没有对象级 metadata）。
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def get(self, key: str) -> bytes:
|
||||
"""不存在时抛出 StorageNotFoundError。"""
|
||||
|
||||
@abstractmethod
|
||||
def get_stream(self, key: str, chunk_size: int = 64 * 1024) -> AsyncIterator[bytes]:
|
||||
"""异步分块读取，用法: `async for chunk in backend.get_stream(key):`。
|
||||
普通方法（非 async def），返回值本身就是异步生成器。
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, key: str) -> None:
|
||||
"""幂等：删除不存在的 key 不应报错。"""
|
||||
|
||||
@abstractmethod
|
||||
async def exists(self, key: str) -> bool:
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def stat(self, key: str) -> ObjectMeta:
|
||||
"""不存在时抛出 StorageNotFoundError。"""
|
||||
|
||||
@abstractmethod
|
||||
def list(self, prefix: str = "") -> AsyncIterator[ObjectMeta]:
|
||||
"""用法: `async for meta in backend.list(prefix):`。"""
|
||||
|
||||
@abstractmethod
|
||||
async def get_url(self, key: str, *, expires_in: Optional[timedelta] = None) -> str:
|
||||
...
|
||||
|
||||
async def copy(self, src_key: str, dst_key: str) -> ObjectMeta:
|
||||
data = await self.get(src_key)
|
||||
return await self.put(dst_key, data)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
async def __aenter__(self) -> "AsyncStorageBackend":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
await self.aclose()
|
||||
@@ -1,146 +0,0 @@
|
||||
"""HTTP client for the internal ``/internal/v1/...`` storage surface.
|
||||
|
||||
The class accepts an injected ``httpx.AsyncClient`` so callers can wire
|
||||
it up with either a real network transport or an ``ASGITransport`` for
|
||||
in-process dispatch. ``SchedulerStorageClient`` extends this class to
|
||||
add a single ``create_object`` helper used by the schedule service.
|
||||
|
||||
All public methods raise :class:`StorageClientError` (or a subclass) on
|
||||
failure. Translation to a web-framework exception (e.g. FastAPI's
|
||||
``HTTPException``) is the caller's responsibility so this client stays
|
||||
usable from non-FastAPI contexts such as the schedule worker.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
class StorageClientError(Exception):
|
||||
"""Base class for storage client failures."""
|
||||
|
||||
|
||||
class StorageUnavailable(StorageClientError):
|
||||
"""Transport-level failure; safe to retry."""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message)
|
||||
self.retryable = True
|
||||
|
||||
|
||||
class StorageRequestFailed(StorageClientError):
|
||||
"""Storage endpoint returned a non-2xx response."""
|
||||
|
||||
def __init__(self, status_code: int, detail: Any) -> None:
|
||||
super().__init__(f"storage request failed with status {status_code}")
|
||||
self.status_code = status_code
|
||||
self.detail = detail
|
||||
|
||||
|
||||
class StorageClient:
|
||||
def __init__(self, client: httpx.AsyncClient) -> None:
|
||||
self.client = client
|
||||
|
||||
async def _request(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
*,
|
||||
payload: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
try:
|
||||
response = await self.client.request(
|
||||
method,
|
||||
path,
|
||||
json=payload,
|
||||
)
|
||||
except httpx.RequestError as exc:
|
||||
raise StorageUnavailable("Storage service temporarily unavailable") from exc
|
||||
if response.is_error:
|
||||
try:
|
||||
detail = response.json().get("detail", response.text)
|
||||
except ValueError:
|
||||
detail = response.text
|
||||
raise StorageRequestFailed(response.status_code, detail)
|
||||
return response.json()
|
||||
|
||||
async def create_upload(
|
||||
self,
|
||||
payload: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
return (await self._request(
|
||||
"POST",
|
||||
"/internal/v1/uploads",
|
||||
payload=payload,
|
||||
))["data"]
|
||||
|
||||
async def complete_upload(
|
||||
self,
|
||||
upload_id: str,
|
||||
payload: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
return (await self._request(
|
||||
"POST",
|
||||
f"/internal/v1/uploads/{upload_id}/complete",
|
||||
payload=payload,
|
||||
))["data"]
|
||||
|
||||
async def create_server_object(
|
||||
self,
|
||||
*,
|
||||
workspace_id: str,
|
||||
user_id: str,
|
||||
usage_type: str,
|
||||
file_name: str,
|
||||
content_type: str,
|
||||
content: bytes,
|
||||
visibility: str,
|
||||
is_immutable: bool,
|
||||
idempotency_key: str,
|
||||
relative_path: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
result = await self._request(
|
||||
"POST",
|
||||
"/internal/v1/objects",
|
||||
payload={
|
||||
"workspace_id": workspace_id,
|
||||
"user_id": user_id,
|
||||
"usage_type": usage_type,
|
||||
"file_name": file_name,
|
||||
"content_type": content_type,
|
||||
"content_base64": base64.b64encode(content).decode("ascii"),
|
||||
"visibility": visibility,
|
||||
"is_immutable": is_immutable,
|
||||
"idempotency_key": idempotency_key,
|
||||
"relative_path": relative_path,
|
||||
},
|
||||
)
|
||||
return result["data"]
|
||||
|
||||
async def create_download_url(
|
||||
self,
|
||||
storage_object_id: str,
|
||||
expires_seconds: int,
|
||||
) -> dict[str, Any]:
|
||||
return (await self._request(
|
||||
"POST",
|
||||
f"/internal/v1/objects/{storage_object_id}/download-url",
|
||||
payload={"expires_seconds": expires_seconds},
|
||||
))["data"]
|
||||
|
||||
async def delete_object(self, storage_object_id: str) -> dict[str, Any]:
|
||||
return (await self._request(
|
||||
"DELETE",
|
||||
f"/internal/v1/objects/{storage_object_id}",
|
||||
))["data"]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"StorageClient",
|
||||
"StorageClientError",
|
||||
"StorageRequestFailed",
|
||||
"StorageUnavailable",
|
||||
]
|
||||
@@ -0,0 +1,116 @@
|
||||
"""使用示例：同一套 create_storage()，靠 config["mode"] 切换同步/异步。"""
|
||||
|
||||
import asyncio
|
||||
|
||||
from common.storage import create_storage
|
||||
from common.storage.exceptions import StorageNotFoundError
|
||||
|
||||
|
||||
def sync_demo():
|
||||
# mode 默认就是 "sync"，可以不写
|
||||
storage = create_storage({"type": "local", "base_dir": "./data_sync"})
|
||||
|
||||
storage.put("docs/hello.txt", b"hello world")
|
||||
print(storage.get("docs/hello.txt"))
|
||||
print(storage.exists("docs/hello.txt"))
|
||||
print(list(storage.list("docs/")))
|
||||
print(storage.get_url("docs/hello.txt"))
|
||||
|
||||
try:
|
||||
storage.get("docs/not_exist.txt")
|
||||
except StorageNotFoundError:
|
||||
print("按预期抛出 StorageNotFoundError")
|
||||
|
||||
# 换成同步 S3，只改配置：
|
||||
# storage = create_storage({"type": "s3", "bucket": "my-bucket"})
|
||||
|
||||
|
||||
async def async_demo():
|
||||
# 只加一个 "mode": "async"，其余配置和参数不变
|
||||
storage = create_storage({"type": "local", "mode": "async", "base_dir": "./data_async"})
|
||||
|
||||
await storage.put("docs/hello.txt", b"hello world")
|
||||
print(await storage.get("docs/hello.txt"))
|
||||
print(await storage.exists("docs/hello.txt"))
|
||||
|
||||
async for meta in storage.list("docs/"):
|
||||
print(meta)
|
||||
|
||||
chunks = []
|
||||
async for chunk in storage.get_stream("docs/hello.txt"):
|
||||
chunks.append(chunk)
|
||||
print(b"".join(chunks))
|
||||
|
||||
try:
|
||||
await storage.get("docs/not_exist.txt")
|
||||
except StorageNotFoundError:
|
||||
print("按预期抛出 StorageNotFoundError")
|
||||
|
||||
# 并发写入，异步模式的典型优势场景
|
||||
tasks = [storage.put(f"batch/{i}.txt", f"content-{i}".encode()) for i in range(10)]
|
||||
await asyncio.gather(*tasks)
|
||||
print("并发写入 10 个对象完成")
|
||||
|
||||
# 换成异步 S3，只改配置：
|
||||
# storage = create_storage({"type": "s3", "mode": "async", "bucket": "my-bucket"})
|
||||
# 高吞吐场景复用连接：
|
||||
# async with create_storage({"type": "s3", "mode": "async", "bucket": "my-bucket"}) as s3:
|
||||
# await s3.put("a.txt", b"1")
|
||||
|
||||
|
||||
def extend_with_new_backend_demo():
|
||||
"""演示独立扩展一种新的存储方式（同步+异步各一个），不用改现有代码。"""
|
||||
import io
|
||||
import time
|
||||
from typing import AsyncIterator, BinaryIO, Iterable, Optional
|
||||
|
||||
from common.storage.base import AsyncStorageBackend, ObjectMeta, StorageBackend
|
||||
from common.storage.exceptions import StorageNotFoundError
|
||||
from common.storage.registry import register_backend
|
||||
|
||||
@register_backend("memory", mode="sync")
|
||||
class MemoryStorageBackend(StorageBackend):
|
||||
def __init__(self, **_ignored):
|
||||
self._store = {}
|
||||
|
||||
def put(self, key, data, *, overwrite=True):
|
||||
body = data if isinstance(data, bytes) else data.read()
|
||||
self._store[key] = body
|
||||
return ObjectMeta(key=key, size=len(body), last_modified=time.time())
|
||||
|
||||
def get(self, key):
|
||||
if key not in self._store:
|
||||
raise StorageNotFoundError(key)
|
||||
return self._store[key]
|
||||
|
||||
def get_stream(self, key):
|
||||
return io.BytesIO(self.get(key))
|
||||
|
||||
def delete(self, key):
|
||||
self._store.pop(key, None)
|
||||
|
||||
def exists(self, key):
|
||||
return key in self._store
|
||||
|
||||
def stat(self, key):
|
||||
if key not in self._store:
|
||||
raise StorageNotFoundError(key)
|
||||
return ObjectMeta(key=key, size=len(self._store[key]))
|
||||
|
||||
def list(self, prefix=""):
|
||||
for key, body in self._store.items():
|
||||
if key.startswith(prefix):
|
||||
yield ObjectMeta(key=key, size=len(body))
|
||||
|
||||
def get_url(self, key, *, expires_in=None):
|
||||
return f"memory://{key}"
|
||||
|
||||
mem_storage = create_storage({"type": "memory", "mode": "sync"})
|
||||
mem_storage.put("a.txt", b"in-memory content")
|
||||
print(mem_storage.get("a.txt"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sync_demo()
|
||||
asyncio.run(async_demo())
|
||||
extend_with_new_backend_demo()
|
||||
@@ -0,0 +1,21 @@
|
||||
"""存储层统一异常。同步/异步后端共用同一套异常类型。"""
|
||||
|
||||
|
||||
class StorageError(Exception):
|
||||
"""所有存储相关异常的基类。"""
|
||||
|
||||
|
||||
class StorageNotFoundError(StorageError):
|
||||
"""指定的 key 不存在。"""
|
||||
|
||||
|
||||
class StorageAlreadyExistsError(StorageError):
|
||||
"""在要求不覆盖的场景下，key 已存在。"""
|
||||
|
||||
|
||||
class StorageConnectionError(StorageError):
|
||||
"""连接/网络层面的错误（如 S3 网络超时、权限问题等）。"""
|
||||
|
||||
|
||||
class StorageConfigError(StorageError):
|
||||
"""配置错误，例如缺少必需参数、backend 类型未注册等。"""
|
||||
@@ -0,0 +1,149 @@
|
||||
"""统一入口：根据配置字典创建具体的存储后端实例。
|
||||
|
||||
配置里的 "mode" 字段决定拿到同步还是异步实现，默认 "sync"（向后兼容）。
|
||||
|
||||
# 同步（默认）
|
||||
storage = create_storage({"type": "local", "base_dir": "./data"})
|
||||
storage.put("a.txt", b"hello")
|
||||
|
||||
# 异步：只需加一个 mode 字段，其余配置不变
|
||||
storage = create_storage({"type": "local", "mode": "async", "base_dir": "./data"})
|
||||
await storage.put("a.txt", b"hello")
|
||||
|
||||
# S3 同理
|
||||
storage = create_storage({"type": "s3", "mode": "async", "bucket": "my-bucket"})
|
||||
|
||||
上层业务代码应该只从这里拿实例，不要直接 import 具体的 XxxStorageBackend /
|
||||
XxxAsyncStorageBackend 类。
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Union
|
||||
|
||||
from .base import AsyncStorageBackend, StorageBackend
|
||||
from .exceptions import StorageConfigError
|
||||
from .registry import get_backend_class
|
||||
from .backends import local, s3 # noqa: F401 # 触发内置后端注册
|
||||
|
||||
AnyStorageBackend = Union[StorageBackend, AsyncStorageBackend]
|
||||
|
||||
|
||||
def create_storage(config: Dict[str, Any]) -> AnyStorageBackend:
|
||||
"""根据配置创建存储后端。
|
||||
|
||||
Args:
|
||||
config: 必须包含 "type" 字段（如 "local" / "s3"）；
|
||||
可选 "mode" 字段（"sync" 默认 / "async"）；
|
||||
其余字段作为 kwargs 传给对应后端的构造函数。
|
||||
|
||||
Returns:
|
||||
mode="sync" 时返回 StorageBackend 实例（同步方法）；
|
||||
mode="async" 时返回 AsyncStorageBackend 实例（方法需要 await）。
|
||||
"""
|
||||
config = dict(config) # 不修改调用方传入的原字典
|
||||
backend_type = config.pop("type", None)
|
||||
mode = config.pop("mode", "sync")
|
||||
|
||||
if not backend_type:
|
||||
raise StorageConfigError("配置缺少 'type' 字段，例如 'local' 或 's3'")
|
||||
|
||||
backend_cls = get_backend_class(backend_type, mode)
|
||||
try:
|
||||
return backend_cls(**config)
|
||||
except TypeError as e:
|
||||
raise StorageConfigError(
|
||||
f"创建后端 (mode={mode}, type={backend_type}) 失败，参数不匹配: {e}"
|
||||
) from e
|
||||
|
||||
|
||||
# 4 个目的化桶的名字（key for app.state.object_stores）。
|
||||
# 在 local 模式下对应 ``local_storage_base_dir/<name>`` 子目录；
|
||||
# 在 s3 模式下对应 ``settings.s3_<name>_bucket``。
|
||||
PURPOSE_BUCKETS: tuple[str, ...] = ("workspace", "version", "run_log", "trash")
|
||||
|
||||
|
||||
def build_storage_config(bucket_name: str) -> Dict[str, Any]:
|
||||
"""根据 ``settings.storage_backend`` 构造 ``create_storage()`` 的入参。
|
||||
|
||||
上层（lifespan 等）只用 ``PURPOSE_BUCKETS`` 循环调用一次，
|
||||
业务代码完全不感知本地 / S3 的差别。
|
||||
|
||||
Args:
|
||||
bucket_name: 桶名，必须是 ``PURPOSE_BUCKETS`` 之一。
|
||||
|
||||
Returns:
|
||||
直接喂给 ``create_storage(...)`` 的 dict。
|
||||
"""
|
||||
# 延迟 import：避免 storage -> config -> storage 的循环依赖
|
||||
from common.config import settings
|
||||
|
||||
if bucket_name not in PURPOSE_BUCKETS:
|
||||
raise StorageConfigError(
|
||||
f"未知 bucket 名称 {bucket_name!r}，可选值: {PURPOSE_BUCKETS}"
|
||||
)
|
||||
|
||||
if settings.storage_backend == "local":
|
||||
return {
|
||||
"type": "local",
|
||||
"mode": "async",
|
||||
"base_dir": str(Path(settings.local_storage_base_dir) / bucket_name),
|
||||
}
|
||||
|
||||
if settings.storage_backend == "s3":
|
||||
return {
|
||||
"type": "s3",
|
||||
"mode": "async",
|
||||
"bucket": getattr(settings, f"s3_{bucket_name}_bucket"),
|
||||
"endpoint_url": settings.s3_endpoint,
|
||||
"aws_access_key_id": settings.s3_access_key,
|
||||
"aws_secret_access_key": settings.s3_secret_key,
|
||||
}
|
||||
|
||||
raise StorageConfigError(
|
||||
f"settings.storage_backend={settings.storage_backend!r} 不支持，"
|
||||
f"可选值: 's3', 'local'"
|
||||
)
|
||||
|
||||
|
||||
def workspaces_root() -> Path:
|
||||
"""返回 runtime 视角下 workspace bucket 的本地路径。
|
||||
|
||||
唯一权威入口。``settings.local_storage_base_dir`` 是与存储相关的
|
||||
唯一路径设置(其它路径都从这里推导):
|
||||
|
||||
- s3 模式(默认):``${local_storage_base_dir}/workspaces``(rclone 把
|
||||
S3 workspace bucket 挂到这里)。
|
||||
- local 模式:``${local_storage_base_dir}/workspace``(直接读写
|
||||
本地目录,无 FUSE 层)。
|
||||
|
||||
``runtime.mount.WORKSPACES_ROOT`` 等于本函数返回值,业务代码不要自己
|
||||
拼路径。
|
||||
"""
|
||||
from common.config import settings # 延迟 import 避免循环
|
||||
base = Path(settings.local_storage_base_dir)
|
||||
if settings.storage_backend == "local":
|
||||
return base / "workspace"
|
||||
return base / "workspaces"
|
||||
|
||||
|
||||
# rclone remote 名字。跟 docker-compose 里的 ``RCLONE_CONFIG_<NAME>_*`` 命名空间
|
||||
# 对应——rclone 通过 env var 名前缀来定位 remote 配置块,所以这里的常量名
|
||||
# 必须跟 ``RCLONE_CONFIG_S3_*`` 的 ``S3`` 部分一致。
|
||||
RCLONE_REMOTE_NAME: str = "s3"
|
||||
|
||||
|
||||
def rclone_remote_spec() -> str:
|
||||
"""rclone mount 用的 remote spec (s3 模式才合法)。
|
||||
|
||||
格式 ``<RCLONE_REMOTE_NAME>:<bucket>``——``runtime.mount.start_rclone_mount``
|
||||
直接喂给 ``rclone mount <spec> <mountpoint>``。
|
||||
|
||||
唯一权威入口:local 模式下没有 rclone,抛 ``StorageConfigError``。
|
||||
"""
|
||||
from common.config import settings
|
||||
if settings.storage_backend != "s3":
|
||||
raise StorageConfigError(
|
||||
"rclone_remote_spec() 仅在 STORAGE_BACKEND=s3 时合法;"
|
||||
f"当前 settings.storage_backend={settings.storage_backend!r}"
|
||||
)
|
||||
return f"{RCLONE_REMOTE_NAME}:{settings.s3_workspace_bucket}"
|
||||
@@ -0,0 +1,61 @@
|
||||
"""后端注册表，用 (mode, name) 作为 key 同时管理同步和异步实现。
|
||||
|
||||
新增一种存储方式的同步或异步实现时，不需要改 factory.py：
|
||||
@register_backend("local", mode="sync")
|
||||
class LocalStorageBackend(StorageBackend): ...
|
||||
|
||||
@register_backend("local", mode="async")
|
||||
class LocalAsyncStorageBackend(AsyncStorageBackend): ...
|
||||
|
||||
只要保证模块被 import 一次即可（backends/__init__.py 里统一 import）。
|
||||
"""
|
||||
|
||||
from typing import Dict, Tuple, Type, Union
|
||||
|
||||
from .base import AsyncStorageBackend, StorageBackend
|
||||
from .exceptions import StorageConfigError
|
||||
|
||||
BackendClass = Union[Type[StorageBackend], Type[AsyncStorageBackend]]
|
||||
|
||||
_REGISTRY: Dict[Tuple[str, str], BackendClass] = {}
|
||||
|
||||
VALID_MODES = ("sync", "async")
|
||||
|
||||
|
||||
def _check_mode(mode: str) -> None:
|
||||
if mode not in VALID_MODES:
|
||||
raise StorageConfigError(f"不支持的 mode: {mode!r}，可选值: {VALID_MODES}")
|
||||
|
||||
|
||||
def register_backend(name: str, mode: str = "sync"):
|
||||
"""类装饰器：把一个后端类注册为 (mode, name) 对应的实现。"""
|
||||
_check_mode(mode)
|
||||
|
||||
def _decorator(cls: BackendClass) -> BackendClass:
|
||||
key = (mode, name)
|
||||
if key in _REGISTRY and _REGISTRY[key] is not cls:
|
||||
raise StorageConfigError(
|
||||
f"存储后端 (mode={mode}, type={name}) 已被注册为 {_REGISTRY[key]!r}"
|
||||
)
|
||||
_REGISTRY[key] = cls
|
||||
return cls
|
||||
|
||||
return _decorator
|
||||
|
||||
|
||||
def get_backend_class(name: str, mode: str = "sync") -> BackendClass:
|
||||
_check_mode(mode)
|
||||
key = (mode, name)
|
||||
try:
|
||||
return _REGISTRY[key]
|
||||
except KeyError:
|
||||
available = ", ".join(
|
||||
f"{m}:{n}" for (m, n) in sorted(_REGISTRY)
|
||||
) or "(无)"
|
||||
raise StorageConfigError(
|
||||
f"未知的存储后端 (mode={mode}, type={name})，当前已注册: {available}"
|
||||
)
|
||||
|
||||
|
||||
def registered_backends() -> Dict[Tuple[str, str], BackendClass]:
|
||||
return dict(_REGISTRY)
|
||||
@@ -1,205 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from typing import Any, BinaryIO
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
import boto3
|
||||
from botocore.client import Config
|
||||
from botocore.exceptions import ClientError
|
||||
|
||||
|
||||
class RustFSObjectStore:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
internal_endpoint: str,
|
||||
access_key: str,
|
||||
secret_key: str,
|
||||
) -> None:
|
||||
common = {
|
||||
"aws_access_key_id": access_key,
|
||||
"aws_secret_access_key": secret_key,
|
||||
"region_name": "us-east-1",
|
||||
"config": Config(
|
||||
signature_version="s3v4",
|
||||
s3={"addressing_style": "path"},
|
||||
),
|
||||
}
|
||||
self.internal = boto3.client(
|
||||
"s3",
|
||||
endpoint_url=internal_endpoint.rstrip("/"),
|
||||
**common,
|
||||
)
|
||||
self._internal_endpoint = internal_endpoint.rstrip("/")
|
||||
|
||||
def ensure_bucket(self, bucket_name: str) -> None:
|
||||
try:
|
||||
self.internal.head_bucket(Bucket=bucket_name)
|
||||
except ClientError as exc:
|
||||
code = str(exc.response.get("Error", {}).get("Code", ""))
|
||||
if code not in {"404", "NoSuchBucket", "NotFound"}:
|
||||
raise
|
||||
self.internal.create_bucket(Bucket=bucket_name)
|
||||
|
||||
def presign_put(
|
||||
self,
|
||||
*,
|
||||
bucket_name: str,
|
||||
object_key: str,
|
||||
content_type: str,
|
||||
expected_hash: str | None,
|
||||
expires_seconds: int,
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
params: dict[str, Any] = {
|
||||
"Bucket": bucket_name,
|
||||
"Key": object_key,
|
||||
"ContentType": content_type,
|
||||
}
|
||||
headers = {"Content-Type": content_type}
|
||||
if expected_hash:
|
||||
params["Metadata"] = {"sha256": expected_hash}
|
||||
headers["x-amz-meta-sha256"] = expected_hash
|
||||
url = self.internal.generate_presigned_url(
|
||||
"put_object",
|
||||
Params=params,
|
||||
ExpiresIn=expires_seconds,
|
||||
)
|
||||
return url, headers
|
||||
|
||||
def presign_get(
|
||||
self,
|
||||
*,
|
||||
bucket_name: str,
|
||||
object_key: str,
|
||||
file_name: str,
|
||||
expires_seconds: int,
|
||||
) -> str:
|
||||
return self.internal.generate_presigned_url(
|
||||
"get_object",
|
||||
Params={
|
||||
"Bucket": bucket_name,
|
||||
"Key": object_key,
|
||||
"ResponseContentDisposition": (
|
||||
f'attachment; filename="{file_name.encode("ascii", "ignore").decode() or "download"}"'
|
||||
),
|
||||
},
|
||||
ExpiresIn=expires_seconds,
|
||||
)
|
||||
|
||||
def rewrite_to_public_path(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
public_base_url: str,
|
||||
) -> str:
|
||||
"""Rewrite the host of a presigned URL to the public edge.
|
||||
|
||||
Replaces the scheme + host (and strips any trailing slash) with
|
||||
``public_base_url``; the bucket prefix is moved under the
|
||||
``/storage/`` path so Nginx can forward the call to RustFS
|
||||
without exposing its port.
|
||||
"""
|
||||
parsed = urlsplit(url)
|
||||
public = urlsplit(public_base_url.rstrip("/"))
|
||||
netloc = public.netloc
|
||||
prefix = public.path.rstrip("/")
|
||||
# Presigned URLs are generated against the internal S3 endpoint
|
||||
# and always start with ``/{bucket}/...``; relocate the bucket
|
||||
# segment under ``/storage/...`` so the public edge can route
|
||||
# the request to the right bucket.
|
||||
rewritten_path = f"{prefix}/storage{parsed.path}"
|
||||
return urlunsplit((
|
||||
public.scheme,
|
||||
netloc,
|
||||
rewritten_path,
|
||||
parsed.query,
|
||||
parsed.fragment,
|
||||
))
|
||||
|
||||
def put_bytes(
|
||||
self,
|
||||
*,
|
||||
bucket_name: str,
|
||||
object_key: str,
|
||||
content: bytes,
|
||||
content_type: str,
|
||||
content_hash: str,
|
||||
) -> None:
|
||||
self.internal.put_object(
|
||||
Bucket=bucket_name,
|
||||
Key=object_key,
|
||||
Body=content,
|
||||
ContentType=content_type,
|
||||
Metadata={"sha256": content_hash},
|
||||
)
|
||||
|
||||
def head(self, *, bucket_name: str, object_key: str) -> dict[str, Any]:
|
||||
return self.internal.head_object(Bucket=bucket_name, Key=object_key)
|
||||
|
||||
def get_bytes(self, *, bucket_name: str, object_key: str) -> bytes:
|
||||
response = self.internal.get_object(Bucket=bucket_name, Key=object_key)
|
||||
body: BinaryIO = response["Body"]
|
||||
try:
|
||||
return body.read()
|
||||
finally:
|
||||
body.close()
|
||||
|
||||
def sha256(self, *, bucket_name: str, object_key: str) -> str:
|
||||
response = self.internal.get_object(Bucket=bucket_name, Key=object_key)
|
||||
body: BinaryIO = response["Body"]
|
||||
digest = hashlib.sha256()
|
||||
while chunk := body.read(1024 * 1024):
|
||||
digest.update(chunk)
|
||||
body.close()
|
||||
return digest.hexdigest()
|
||||
|
||||
def delete(self, *, bucket_name: str, object_key: str) -> None:
|
||||
self.internal.delete_object(Bucket=bucket_name, Key=object_key)
|
||||
|
||||
def copy(
|
||||
self,
|
||||
*,
|
||||
source_bucket: str,
|
||||
source_key: str,
|
||||
dest_bucket: str,
|
||||
dest_key: str,
|
||||
) -> None:
|
||||
"""Server-side copy ``source_bucket/source_key`` → ``dest_bucket/dest_key``.
|
||||
|
||||
``CopySource`` is a single header string of the form
|
||||
``/{bucket}/{key}`` — must NOT be URL-encoded or quoted.
|
||||
"""
|
||||
self.internal.copy_object(
|
||||
Bucket=dest_bucket,
|
||||
Key=dest_key,
|
||||
CopySource={"Bucket": source_bucket, "Key": source_key},
|
||||
)
|
||||
|
||||
def move_to_trash(
|
||||
self,
|
||||
*,
|
||||
source_bucket: str,
|
||||
source_key: str,
|
||||
trash_bucket: str,
|
||||
trash_key: str,
|
||||
) -> None:
|
||||
"""Copy an object into the trash bucket and delete the source.
|
||||
|
||||
The copy is a server-side operation in RustFS (no data flows
|
||||
through the client). The source delete is best-effort: if it
|
||||
fails after the copy succeeds the trash holds the only copy of
|
||||
the bytes, which is exactly the point — the caller can retry.
|
||||
"""
|
||||
self.copy(
|
||||
source_bucket=source_bucket,
|
||||
source_key=source_key,
|
||||
dest_bucket=trash_bucket,
|
||||
dest_key=trash_key,
|
||||
)
|
||||
try:
|
||||
self.delete(bucket_name=source_bucket, object_key=source_key)
|
||||
except ClientError:
|
||||
# Source was already gone, or transient delete failure —
|
||||
# the trash copy is what matters; caller logs and moves on.
|
||||
pass
|
||||
@@ -10,7 +10,6 @@ from common.schemas import StrictModel
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CompleteUploadRequest",
|
||||
"CreateUploadRequest",
|
||||
"DownloadUrlRequest",
|
||||
"ServerObjectRequest",
|
||||
@@ -34,6 +33,8 @@ class CreateUploadRequest(StrictModel):
|
||||
expected_size_bytes: int = Field(ge=0, le=100 * 1024 * 1024)
|
||||
expected_hash: str | None = Field(default=None, min_length=64, max_length=64)
|
||||
idempotency_key: str = Field(min_length=8, max_length=128)
|
||||
visibility: Literal["private", "workspace", "public"] = "private"
|
||||
is_immutable: bool = False
|
||||
|
||||
@field_validator("expected_hash")
|
||||
@classmethod
|
||||
@@ -46,21 +47,6 @@ class CreateUploadRequest(StrictModel):
|
||||
return normalized
|
||||
|
||||
|
||||
class CompleteUploadRequest(StrictModel):
|
||||
usage_type: Literal[
|
||||
"data_resource",
|
||||
"version_artifact",
|
||||
"snapshot",
|
||||
"run_log",
|
||||
"run_result",
|
||||
"working_copy",
|
||||
"public_script",
|
||||
]
|
||||
file_name: str = Field(min_length=1, max_length=255)
|
||||
visibility: Literal["private", "workspace", "public"] = "private"
|
||||
is_immutable: bool = False
|
||||
|
||||
|
||||
class ServerObjectRequest(StrictModel):
|
||||
workspace_id: str = Field(min_length=26, max_length=26)
|
||||
user_id: str = Field(min_length=26, max_length=26)
|
||||
|
||||
Reference in New Issue
Block a user