iTi-Flask-IOT/iti/applications/routes/common/upload.py

"""
通用文件上传模块
支持：
1. 小文件直接上传
2. 大文件分片上传（自定义协议）
3. 秒传（基于文件哈希）
4. 支持多存储类型（local/oss/minio）
"""

from __future__ import annotations

import hashlib
from io import BytesIO

from apiflask import APIBlueprint
from flask import request
from flask_jwt_extended import jwt_required, get_jwt_identity

from iti.applications.common.utils import success
from iti.applications.common.exceptions.biz_exp import BizException
from iti.applications.models import SysFileSchema
from iti.applications.service.sys.sys_file import SysFileService
from .schemas.upload import (
    UploadIn,
    ChunkUploadInitIn,
    ChunkUploadIn,
    ChunkUploadMergeIn,
)

bp = APIBlueprint("common_upload", __name__, url_prefix="/upload", tag="通用.上传")


# ---------------------------------------------------------------------------
# 小文件直接上传
# ---------------------------------------------------------------------------


@bp.post("")
@jwt_required()
@bp.doc(security="JWT")
@bp.input(UploadIn, location="form_and_files")
def upload_file(form_and_files_data):
    """文件直接上传"""
    file = form_and_files_data["file"]
    directory_id = form_and_files_data.get("directoryId")
    storage_type = form_and_files_data.get("storageType")

    # 从 UploadIn Schema 中获取所有已定义的字段名
    schema_fields = set(UploadIn().fields.keys())

    # 收集所有非 Schema 定义的额外字段作为 metadata
    metadata = {}
    for key, value in form_and_files_data.items():
        if key not in schema_fields:
            metadata[key] = value

    result = SysFileService.upload_file(
        file,
        directory_id=directory_id,
        metadata=metadata,
        storage_type=storage_type,
    )

    # result 是一个字典，包含 file_record 和 instantUpload
    file_record = result["file"]
    instant_upload = result["instantUpload"]

    # 使用 Schema 序列化文件记录
    from iti.applications.models import SysFileSchema

    schema = SysFileSchema()
    file_data = schema.dump(file_record)

    # 添加秒传标识
    file_data["instantUpload"] = instant_upload

    return success(file_data)


# ---------------------------------------------------------------------------
# 分片上传
# ---------------------------------------------------------------------------


@bp.post("/chunk/init")
@jwt_required()
@bp.doc(security="JWT")
@bp.input(ChunkUploadInitIn)
def chunk_upload_init(json_data):
    """初始化分片上传"""
    result = SysFileService.init_chunk_upload(
        filename=json_data["filename"],
        file_size=json_data["file_size"],
        file_hash=json_data.get("file_hash"),
        chunk_size=json_data.get("chunk_size", 2 * 1024 * 1024),
        total_chunks=json_data.get("total_chunks"),
        directory_id=json_data.get("directory_id"),
        storage_type=json_data.get("storage_type"),
        metadata=json_data.get("metadata", {}),
    )

    # 如果是秒传，需要序列化文件对象
    if result.get("instantUpload"):
        from iti.applications.models import SysFileSchema

        schema = SysFileSchema()
        file_data = schema.dump(result["file"])
        file_data["instantUpload"] = True
        return success(file_data)

    # 非秒传，返回上传ID和已上传分片列表（已经是驼峰命名）
    return success(result)


@bp.post("/chunk/upload")
@jwt_required()
@bp.doc(security="JWT")
@bp.input(ChunkUploadIn, location="form_and_files")
def chunk_upload(form_and_files_data):
    """上传分片"""
    upload_id = form_and_files_data["upload_id"]
    chunk_index = form_and_files_data["chunk_index"]
    chunk_hash = form_and_files_data.get("chunk_hash")
    chunk_file = form_and_files_data["file"]

    chunk_data = chunk_file.read()

    if chunk_hash:
        actual_hash = hashlib.md5(chunk_data).hexdigest()
        if actual_hash != chunk_hash:
            raise BizException("分片数据校验失败", code=400)

    result = SysFileService.upload_chunk(
        upload_id=upload_id,
        chunk_index=chunk_index,
        chunk_data=chunk_data,
    )

    return success(result)


@bp.post("/chunk/merge")
@jwt_required()
@bp.doc(security="JWT")
@bp.input(ChunkUploadMergeIn)
def chunk_upload_merge(json_data):
    """合并分片"""
    upload_id = json_data["upload_id"]
    file_hash = json_data.get("file_hash")

    file_record = SysFileService.merge_chunks(
        upload_id=upload_id,
        file_hash=file_hash,
    )

    # 使用 Schema 序列化文件记录
    from iti.applications.models import SysFileSchema

    schema = SysFileSchema()
    file_data = schema.dump(file_record)

    # 分片上传合并后不是秒传
    file_data["instantUpload"] = False

    return success(file_data)


@bp.delete("/chunk/<string:upload_id>")
@jwt_required()
@bp.doc(security="JWT")
def chunk_upload_abort(upload_id: str):
    """取消分片上传"""
    SysFileService.abort_chunk_upload(upload_id)
    return success(message="上传已取消")


@bp.get("/chunk/<string:upload_id>/progress")
@jwt_required()
@bp.doc(security="JWT")
def chunk_upload_progress(upload_id: str):
    """查询上传进度"""
    progress = SysFileService.get_chunk_upload_progress(upload_id)
    return success(progress)


@bp.post("/chunk/cleanup")
@jwt_required()
@bp.doc(security="JWT")
def cleanup_expired_chunks():
    """清理过期的分片上传临时文件（管理员接口）"""
    from flask_jwt_extended import get_jwt

    # 检查是否有管理员权限（可选，根据你的权限系统调整）
    # claims = get_jwt()
    # if not claims.get("is_admin"):
    #     raise BizException("需要管理员权限", code=403)

    # 默认清理7天前的文件
    days = request.args.get("days", 7, type=int)
    result = SysFileService.cleanup_expired_chunk_uploads(days)

    return success(
        result,
        message=f"清理完成，删除 {result['cleaned_dirs']} 个目录，释放 {result['cleaned_size']} 字节空间",
    )