shanxia2.0/backend/app/services/doclib_storage.py

"""私有 MinIO 对象访问；不签发公开 URL，不读取整个文件到内存。"""
import hashlib
import logging
from urllib.parse import quote
from uuid import uuid4

from fastapi import HTTPException, UploadFile
from minio import Minio
from minio.error import S3Error
from starlette.background import BackgroundTask
from starlette.responses import StreamingResponse
from urllib3 import PoolManager, Timeout

from app.core.config import settings

logger = logging.getLogger('sqxp')


class PrivateStorage:
    def __init__(self):
        self.bucket = settings.DOCLIB_MINIO_BUCKET.strip()
        if not settings.MINIO_ENDPOINT or not self.bucket or self.bucket == settings.MINIO_BUCKET:
            raise HTTPException(503, '请配置文档库专用私有 MinIO 桶')
        self.client = Minio(settings.MINIO_ENDPOINT, access_key=settings.MINIO_ACCESS_KEY,
                            secret_key=settings.MINIO_SECRET_KEY, secure=settings.MINIO_SECURE,
                            http_client=PoolManager(timeout=Timeout(connect=5, read=30), retries=1))
        try:
            if not self.client.bucket_exists(self.bucket):
                try:
                    self.client.make_bucket(self.bucket)
                except S3Error as exc:
                    if exc.code not in ('BucketAlreadyOwnedByYou', 'BucketAlreadyExists'):
                        raise
            # 专用桶禁止任何匿名策略，不能因部署误配而让文档绕过 ACL。
            try:
                policy = self.client.get_bucket_policy(self.bucket)
            except S3Error as exc:
                if exc.code != 'NoSuchBucketPolicy':
                    raise
                policy = None
            if policy:
                raise HTTPException(503, '文档库桶必须保持私有，请移除该专用桶的匿名访问策略')
        except HTTPException:
            raise
        except Exception:
            logger.exception('doclib private bucket unavailable')
            raise HTTPException(503, '文档存储服务暂不可用')

    def upload(self, key, file, length):
        self.client.put_object(self.bucket, key, file, length=length, content_type='application/octet-stream')

    def open(self, key):
        return self.client.get_object(self.bucket, key)

    def discard(self, key):
        # 仅用于尚未成功建立数据库记录的上传；业务删除保留原对象。
        self.client.remove_object(self.bucket, key)


def private_storage():
    return PrivateStorage()


def upload_document(service, file: UploadFile, folder_id: str):
    name = (file.filename or '').replace('\\', '/').rsplit('/', 1)[-1].strip()
    if not name or len(name) > 255 or any(ord(c) < 32 or ord(c) == 127 for c in name):
        raise HTTPException(400, '文件名称无效或超过255个字符')
    if len(folder_id) > 64:
        raise HTTPException(400, '文件夹参数无效')
    if folder_id:
        service._folder(folder_id, own=True)
    file.file.seek(0, 2)
    length = file.file.tell()
    file.file.seek(0)
    if length > settings.DOCLIB_MAX_UPLOAD_BYTES:
        raise HTTPException(413, f'单个文档不能超过 {settings.DOCLIB_MAX_UPLOAD_BYTES // 1024 // 1024} MB')
    if length == 0:
        raise HTTPException(400, '不能上传空文件')
    storage = private_storage()
    key = hashlib.sha256(service.tenant.encode()).hexdigest()[:20] + '/' + uuid4().hex
    try:
        storage.upload(key, file.file, length)
    except Exception:
        logger.exception('doclib upload failed')
        try:
            storage.discard(key)
        except Exception:
            logger.exception('doclib failed to discard ambiguous upload')
        raise HTTPException(503, '文档上传失败，请稍后重试')
    try:
        return service.create_file(name=name, folder_id=folder_id, size=length,
                                   content_type=(file.content_type or 'application/octet-stream')[:255], object_key=key)
    except Exception:
        service.db.rollback()
        try:
            storage.discard(key)
        except Exception:
            logger.exception('doclib failed to discard unregistered upload')
        raise


def download_document(service, file_id):
    row = service.download_record(file_id)
    name, key, size = row.name, row.object_key, row.size
    try:
        response = private_storage().open(key)
    except HTTPException:
        raise
    except Exception:
        logger.exception('doclib download failed')
        raise HTTPException(503, '文档暂时无法下载，请稍后重试')
    closed = False

    def close():
        nonlocal closed
        if not closed:
            closed = True
            response.close()
            response.release_conn()

    def chunks():
        try:
            yield from response.stream(amt=64 * 1024)
        finally:
            close()

    return StreamingResponse(chunks(), media_type='application/octet-stream', background=BackgroundTask(close),
                             headers={'Content-Disposition': "attachment; filename*=UTF-8''" + quote(name, safe=''),
                                      'Content-Length': str(size), 'Cache-Control': 'no-store',
                                      'X-Content-Type-Options': 'nosniff'})