leaudit-platform-backend/fastapi_common/fastapi_common_storage/oss_client.py

"""统一 OSS / MinIO 客户端。"""

from __future__ import annotations

import tempfile
from dataclasses import dataclass
from datetime import timedelta
from io import BytesIO
from pathlib import Path
from urllib.parse import urlparse

import httpx

from fastapi_admin.config import (
    OSS_ACCESS_KEY,
    OSS_BASE_URL,
    OSS_BUCKET,
    OSS_ENDPOINT,
    OSS_PRESIGN_EXPIRE_SECONDS,
    OSS_SECRET_KEY,
    OSS_USE_SSL,
)


@dataclass(frozen=True)
class OssObjectRef:
    """对象存储引用。"""

    bucket: str
    objectKey: str
    source: str
    isDirectUrl: bool = False


class OssClient:
    """统一封装文档、规则与产物的 OSS 访问。"""

    def __init__(
        self,
        *,
        endpoint: str | None = None,
        accessKey: str | None = None,
        secretKey: str | None = None,
        bucket: str | None = None,
        useSsl: bool | None = None,
        baseUrl: str | None = None,
        presignExpireSeconds: int | None = None,
    ) -> None:
        self.endpoint = endpoint or OSS_ENDPOINT
        self.accessKey = accessKey or OSS_ACCESS_KEY
        self.secretKey = secretKey or OSS_SECRET_KEY
        self.bucket = bucket or OSS_BUCKET
        self.useSsl = OSS_USE_SSL if useSsl is None else useSsl
        self.baseUrl = (baseUrl or OSS_BASE_URL).rstrip("/")
        self.presignExpireSeconds = presignExpireSeconds or OSS_PRESIGN_EXPIRE_SECONDS
        self._minioClient = None

    def UploadBytes(
        self,
        ObjectKey: str,
        Content: bytes,
        ContentType: str = "application/octet-stream",
        Bucket: str | None = None,
    ) -> str:
        """上传二进制内容并返回对象引用。"""
        Client = self._GetMinioClient()
        TargetBucket = Bucket or self.bucket
        Data = BytesIO(Content)
        Client.put_object(
            TargetBucket,
            ObjectKey,
            Data,
            length=len(Content),
            content_type=ContentType,
        )
        return self.BuildObjectUrl(ObjectKey=ObjectKey, Bucket=TargetBucket)

    def UploadText(
        self,
        ObjectKey: str,
        Content: str,
        ContentType: str = "text/plain; charset=utf-8",
        Bucket: str | None = None,
    ) -> str:
        """上传文本内容并返回对象引用。"""
        return self.UploadBytes(
            ObjectKey=ObjectKey,
            Content=Content.encode("utf-8"),
            ContentType=ContentType,
            Bucket=Bucket,
        )

    def UploadFile(
        self,
        ObjectKey: str,
        LocalPath: str,
        ContentType: str = "application/octet-stream",
        Bucket: str | None = None,
    ) -> str:
        """上传本地文件并返回对象引用。"""
        FilePath = Path(LocalPath)
        return self.UploadBytes(
            ObjectKey=ObjectKey,
            Content=FilePath.read_bytes(),
            ContentType=ContentType,
            Bucket=Bucket,
        )

    def DownloadBytes(self, Source: str, Bucket: str | None = None) -> bytes:
        """下载对象内容。"""
        Ref = self.ResolveObjectRef(Source=Source, Bucket=Bucket)
        if Ref.isDirectUrl:
            return self._DownloadBytesFromUrl(Ref.source)

        Client = self._GetMinioClient()
        Response = Client.get_object(Ref.bucket, Ref.objectKey)
        try:
            return Response.read()
        finally:
            Response.close()
            Response.release_conn()

    def DownloadToTempFile(
        self,
        Source: str,
        *,
        Suffix: str = "",
        Prefix: str = "oss-",
        Bucket: str | None = None,
    ) -> str:
        """下载对象到本地临时文件。"""
        Content = self.DownloadBytes(Source=Source, Bucket=Bucket)
        return self.WriteTempBytes(Content=Content, Suffix=Suffix, Prefix=Prefix)

    def WriteTempBytes(
        self,
        *,
        Content: bytes,
        Suffix: str = "",
        Prefix: str = "oss-",
    ) -> str:
        """把内存中的对象内容写入本地临时文件。"""
        with tempfile.NamedTemporaryFile(mode="wb", suffix=Suffix, prefix=Prefix, delete=False) as TempFile:
            TempFile.write(Content)
            return TempFile.name

    def EnsureBucket(self, Bucket: str | None = None) -> str:
        """确保 bucket 存在，不存在则创建。返回 bucket 名。"""
        TargetBucket = Bucket or self.bucket
        Client = self._GetMinioClient()
        from minio import Minio
        try:
            if not Client.bucket_exists(TargetBucket):
                Client.make_bucket(TargetBucket)
        except Minio.S3Error:
            pass
        return TargetBucket

    def ObjectExists(self, Source: str, Bucket: str | None = None) -> bool:
        """判断对象是否存在。"""
        Ref = self.ResolveObjectRef(Source=Source, Bucket=Bucket)
        if Ref.isDirectUrl:
            try:
                Response = httpx.head(Ref.source, timeout=30.0, follow_redirects=True)
                return Response.status_code < 400
            except Exception:
                return False

        Client = self._GetMinioClient()
        try:
            Client.stat_object(Ref.bucket, Ref.objectKey)
            return True
        except Exception:
            return False

    def BuildObjectUrl(self, ObjectKey: str, Bucket: str | None = None) -> str:
        """构造对象访问地址。"""
        TargetBucket = Bucket or self.bucket
        if self.baseUrl:
            return f"{self.baseUrl}/{TargetBucket}/{ObjectKey.lstrip('/')}"
        return f"oss://{TargetBucket}/{ObjectKey.lstrip('/')}"

    def PresignGetUrl(self, Source: str, Bucket: str | None = None) -> str:
        """生成对象下载签名 URL。"""
        Ref = self.ResolveObjectRef(Source=Source, Bucket=Bucket)
        if Ref.isDirectUrl:
            return Ref.source

        Client = self._GetMinioClient()
        return Client.presigned_get_object(
            Ref.bucket,
            Ref.objectKey,
            expires=timedelta(seconds=self.presignExpireSeconds),
        )

    def ResolveObjectRef(self, Source: str, Bucket: str | None = None) -> OssObjectRef:
        """把 URL / oss:// / object key 统一解析成对象引用。"""
        Parsed = urlparse(Source)
        if Parsed.scheme in {"http", "https"}:
            return OssObjectRef(
                bucket=Bucket or self.bucket,
                objectKey="",
                source=Source,
                isDirectUrl=True,
            )

        if Parsed.scheme == "oss":
            BucketName = Parsed.netloc or Bucket or self.bucket
            ObjectKey = Parsed.path.lstrip("/")
            return OssObjectRef(
                bucket=BucketName,
                objectKey=ObjectKey,
                source=Source,
            )

        return OssObjectRef(
            bucket=Bucket or self.bucket,
            objectKey=Source.lstrip("/"),
            source=Source,
        )

    def _DownloadBytesFromUrl(self, Url: str) -> bytes:
        """从直链地址下载对象。"""
        with httpx.Client(timeout=60.0, follow_redirects=True) as Client:
            Response = Client.get(Url)
            Response.raise_for_status()
            return Response.content

    def _GetMinioClient(self):
        """获取底层 MinIO 客户端。"""
        if self._minioClient is None:
            from minio import Minio

            Endpoint = self.endpoint
            Parsed = urlparse(Endpoint if "://" in Endpoint else f"http://{Endpoint}")
            Host = Parsed.netloc or Parsed.path
            Secure = self.useSsl if Parsed.scheme not in {"http", "https"} else Parsed.scheme == "https"
            self._minioClient = Minio(
                Host,
                access_key=self.accessKey,
                secret_key=self.secretKey,
                secure=Secure,
            )
        return self._minioClient