model_train_dm/backend/apps/datasets/views/temp_file.py

import mimetypes
import os
import uuid
from datetime import datetime

from django.conf import settings
from django.shortcuts import get_object_or_404
from django.http import StreamingHttpResponse
from rest_framework import status
from rest_framework.decorators import api_view, parser_classes, permission_classes
from rest_framework.pagination import PageNumberPagination
from rest_framework.parsers import MultiPartParser
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response

from apps.common.minio_client import minio_storage
from apps.common.serializers import AiTempFileSerializer
from apps.core.models import AiTempFile


TEMPFILE_BUCKET = "tempfile"
IMAGE_TYPES = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp", ".tif", ".tiff"}
VIDEO_TYPES = {".mp4", ".avi", ".mov", ".mkv", ".wmv", ".flv", ".mpeg", ".mpg", ".webm"}
AUDIO_TYPES = {".mp3", ".wav", ".aac", ".m4a", ".flac", ".ogg"}
TEXT_TYPES = {".txt", ".csv", ".json", ".xml"}


def _detect_original_type(filename: str) -> str:
    ext = os.path.splitext(filename or "")[1].lower()
    if ext in IMAGE_TYPES:
        return "01"
    if ext in VIDEO_TYPES:
        return "02"
    if ext in AUDIO_TYPES:
        return "03"
    if ext == ".zip":
        return "04"
    if ext in TEXT_TYPES:
        return "05"
    return "99"


def _stream_object(response, chunk_size: int = 1024 * 1024):
    try:
        while True:
            chunk = response.read(chunk_size)
            if not chunk:
                break
            yield chunk
    finally:
        response.close()
        response.release_conn()


@api_view(["POST"])
@parser_classes([MultiPartParser])
@permission_classes([IsAuthenticated])
def upload_temp_file(request):
    uploaded_file = request.FILES.get("file")
    if not uploaded_file:
        return Response({"error": "缺少 file 字段"}, status=status.HTTP_400_BAD_REQUEST)

    file_id = str(uuid.uuid4())
    bucket_name = request.data.get("bucket_name") or TEMPFILE_BUCKET
    prefix = request.data.get("prefix") or datetime.now().strftime("%Y-%m-%d")
    object_name = request.data.get("object_name") or minio_storage.generate_object_name(uploaded_file.name, prefix)
    content_type = uploaded_file.content_type or mimetypes.guess_type(uploaded_file.name)[0] or "application/octet-stream"
    original_type = request.data.get("original_type") or _detect_original_type(uploaded_file.name)

    minio_storage.upload_stream(
        bucket_name=bucket_name,
        object_name=object_name,
        data=uploaded_file.file,
        length=uploaded_file.size,
        content_type=content_type,
    )

    obj = AiTempFile.objects.create(
        id=file_id,
        original_path=request.data.get("original_path") or "",
        original_name=uploaded_file.name,
        original_type=original_type,
        bucket_name=bucket_name,
        object_name=object_name,
        object_size=uploaded_file.size,
        content_type=content_type,
        status=request.data.get("status") or "01",
    )
    serializer = AiTempFileSerializer(obj)
    return Response({"message": "临时文件上传成功", "data": serializer.data})


@api_view(["GET"])
@permission_classes([IsAuthenticated])
def search_temp_files(request):
    filter_kwargs = {}
    original_name = request.GET.get("original_name")
    original_type = request.GET.get("original_type")
    status_code = request.GET.get("status")
    page_size = request.GET.get("page_size") or 10

    if original_name:
        filter_kwargs["original_name__contains"] = original_name
    if original_type:
        filter_kwargs["original_type"] = original_type
    if status_code:
        filter_kwargs["status"] = status_code

    queryset = AiTempFile.objects.filter(**filter_kwargs).order_by("-create_time")
    paginator = PageNumberPagination()
    paginator.page_size = page_size
    page = paginator.paginate_queryset(queryset, request)
    serializer = AiTempFileSerializer(page, many=True)
    return paginator.get_paginated_response(serializer.data)


@api_view(["GET"])
@permission_classes([IsAuthenticated])
def read_temp_file(request):
    temp_id = request.GET.get("id")
    obj = get_object_or_404(AiTempFile, id=temp_id)
    serializer = AiTempFileSerializer(obj)
    return Response(serializer.data)


@api_view(["PUT"])
@permission_classes([IsAuthenticated])
def update_temp_file(request):
    data = request.data
    obj = get_object_or_404(AiTempFile, id=data.get("id"))
    serializer = AiTempFileSerializer(instance=obj, data=data, partial=True)
    if serializer.is_valid():
        serializer.save()
        return Response({"message": "临时文件更新成功"})
    return Response({"error": "无效的数据", "details": serializer.errors}, status=status.HTTP_400_BAD_REQUEST)


@api_view(["DELETE"])
@permission_classes([IsAuthenticated])
def delete_temp_file(request):
    temp_id = request.GET.get("id")
    obj = get_object_or_404(AiTempFile, id=temp_id)
    if obj.bucket_name and obj.object_name:
        minio_storage.delete_object(obj.bucket_name, obj.object_name)
    obj.delete()
    return Response({"message": "临时文件删除成功"})


@api_view(["GET"])
@permission_classes([IsAuthenticated])
def download_temp_file(request):
    temp_id = request.GET.get("id")
    obj = get_object_or_404(AiTempFile, id=temp_id)
    response = minio_storage.get_object(obj.bucket_name, obj.object_name)
    content_type = obj.content_type or "application/octet-stream"
    filename = obj.original_name or os.path.basename(obj.object_name or "")
    stream = StreamingHttpResponse(_stream_object(response), content_type=content_type)
    stream["Content-Disposition"] = f'attachment; filename="{filename}"'
    return stream


@api_view(["GET"])
@permission_classes([IsAuthenticated])
def preview_temp_file(request):
    temp_id = request.GET.get("id")
    obj = get_object_or_404(AiTempFile, id=temp_id)
    response = minio_storage.get_object(obj.bucket_name, obj.object_name)
    content_type = obj.content_type or "application/octet-stream"
    stream = StreamingHttpResponse(_stream_object(response), content_type=content_type)
    stream["Content-Disposition"] = "inline"
    return stream