system-backu/backend/app/core/storage/s3.py

"""S3 兼容对象存储（支持 MinIO / Ceph / AWS S3）。"""
import asyncio
import hashlib
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Any, AsyncIterator

import aioboto3

from app.core.storage.base import BaseStorage, StorageObject
from app.utils.logging import get_logger

log = get_logger(__name__)


class S3Storage(BaseStorage):
    type = "s3"

    def __init__(self, config: dict):
        super().__init__(config)
        self.endpoint_url: str | None = config.get("endpoint_url") or None
        self.region: str = config.get("region") or "us-east-1"
        self.bucket: str = config["bucket"]
        self.access_key: str = config["access_key"]
        self.secret_key: str = config["secret_key"]
        self.use_ssl: bool = bool(config.get("use_ssl", True))
        self.path_prefix: str = config.get("path_prefix", "").strip("/")
        self.addressing_style: str = config.get("addressing_style", "auto")

    @asynccontextmanager
    async def _client(self):
        session = aioboto3.Session()
        async with session.client(
            "s3",
            endpoint_url=self.endpoint_url,
            region_name=self.region,
            aws_access_key_id=self.access_key,
            aws_secret_access_key=self.secret_key,
            use_ssl=self.use_ssl,
            config=__import__("botocore.config", fromlist=["Config"]).Config(
                signature_version="s3v4",
                s3={"addressing_style": self.addressing_style},
            ),
        ) as client:
            yield client

    def _full_key(self, key: str) -> str:
        key = key.lstrip("/")
        if self.path_prefix:
            return f"{self.path_prefix}/{key}"
        return key

    async def test_connection(self) -> tuple[bool, str]:
        try:
            async with self._client() as client:
                await client.head_bucket(Bucket=self.bucket)
                return True, f"OK: bucket={self.bucket}"
        except Exception as e:
            return False, f"FAIL: {type(e).__name__}: {e}"

    async def upload_stream(
        self,
        source: AsyncIterator[bytes],
        key: str,
        content_type: str = "application/octet-stream",
    ) -> tuple[int, str]:
        sha = hashlib.sha256()
        size = 0
        full_key = self._full_key(key)

        async with self._client() as client:
            # 用 UploadPart + 内存缓冲整个对象再 put_object（对中小备份足够）
            # 大文件可改为 multipart；这里用 MultipartUploader 模式，兼顾大文件
            from io import BytesIO

            buffer = BytesIO()
            async for chunk in source:
                buffer.write(chunk)
                sha.update(chunk)
                size += len(chunk)

            # 5MB 阈值：超过用 multipart
            PART_SIZE = 5 * 1024 * 1024
            if size <= PART_SIZE:
                buffer.seek(0)
                await client.put_object(
                    Bucket=self.bucket,
                    Key=full_key,
                    Body=buffer.getvalue(),
                    ContentType=content_type,
                )
            else:
                # multipart
                mpu = await client.create_multipart_upload(
                    Bucket=self.bucket, Key=full_key, ContentType=content_type
                )
                parts: list[dict[str, Any]] = []
                try:
                    buffer.seek(0)
                    part_number = 1
                    while True:
                        data = buffer.read(PART_SIZE)
                        if not data:
                            break
                        resp = await client.upload_part(
                            Bucket=self.bucket,
                            Key=full_key,
                            PartNumber=part_number,
                            UploadId=mpu["UploadId"],
                            Body=data,
                        )
                        parts.append({"PartNumber": part_number, "ETag": resp["ETag"]})
                        part_number += 1
                    await client.complete_multipart_upload(
                        Bucket=self.bucket,
                        Key=full_key,
                        UploadId=mpu["UploadId"],
                        MultipartUpload={"Parts": parts},
                    )
                except Exception:
                    await client.abort_multipart_upload(
                        Bucket=self.bucket, Key=full_key, UploadId=mpu["UploadId"]
                    )
                    raise

        return size, sha.hexdigest()

    async def download_stream(self, key: str) -> AsyncIterator[bytes]:
        full_key = self._full_key(key)
        async with self._client() as client:
            obj = await client.get_object(Bucket=self.bucket, Key=full_key)
            body = obj["Body"]
            try:
                while True:
                    chunk = await asyncio.get_event_loop().run_in_executor(
                        None, body.read, 64 * 1024
                    )
                    if not chunk:
                        break
                    yield chunk
            finally:
                body.close()

    async def delete(self, key: str) -> None:
        full_key = self._full_key(key)
        async with self._client() as client:
            await client.delete_object(Bucket=self.bucket, Key=full_key)

    async def exists(self, key: str) -> bool:
        full_key = self._full_key(key)
        async with self._client() as client:
            try:
                await client.head_object(Bucket=self.bucket, Key=full_key)
                return True
            except Exception:
                return False

    async def list(self, prefix: str = "") -> list[StorageObject]:
        full_prefix = self._full_key(prefix)
        out: list[StorageObject] = []
        async with self._client() as client:
            paginator = client.get_paginator("list_objects_v2")
            async for page in paginator.paginate(Bucket=self.bucket, Prefix=full_prefix):
                for obj in page.get("Contents", []):
                    key = obj["Key"]
                    if self.path_prefix and key.startswith(self.path_prefix + "/"):
                        key = key[len(self.path_prefix) + 1 :]
                    last_modified = obj.get("LastModified")
                    out.append(
                        StorageObject(
                            key=key,
                            size=int(obj.get("Size", 0)),
                            last_modified=last_modified.isoformat() if last_modified else None,
                        )
                    )
        return out