62 lines
2.2 KiB
Python

from pathlib import Path
import uuid
from django.db import transaction
from rest_framework import status
from rest_framework.parsers import FormParser, MultiPartParser
from rest_framework.response import Response
from rest_framework.views import APIView
from rest_framework.viewsets import ModelViewSet
from apps.common.api import TeamScopedViewSetMixin, get_current_team
from .models import Asset, AssetFile
from .serializers import AssetSerializer, AssetUploadSerializer
from .storage import TosStorage
class AssetViewSet(TeamScopedViewSetMixin, ModelViewSet):
queryset = Asset.objects.prefetch_related("files").all()
serializer_class = AssetSerializer
search_fields = ["name", "description"]
ordering_fields = ["created_at", "updated_at", "name"]
class AssetUploadView(APIView):
parser_classes = [MultiPartParser, FormParser]
@transaction.atomic
def post(self, request):
serializer = AssetUploadSerializer(data=request.data)
serializer.is_valid(raise_exception=True)
team = get_current_team(request.user)
upload = serializer.validated_data["file"]
suffix = Path(upload.name).suffix.lower()
asset_id = uuid.uuid4()
object_key = f"teams/{team.id}/uploads/{asset_id}{suffix}"
stored = TosStorage().upload_fileobj(
fileobj=upload.file,
object_key=object_key,
content_type=upload.content_type or "application/octet-stream",
)
asset = Asset.objects.create(
id=asset_id,
team=team,
created_by=request.user,
name=serializer.validated_data.get("name") or upload.name,
asset_type=serializer.validated_data["asset_type"],
source=Asset.Source.UPLOAD,
category=serializer.validated_data["category"],
description=serializer.validated_data.get("description", ""),
)
AssetFile.objects.create(
asset=asset,
object_key=stored.object_key,
bucket=stored.bucket,
content_type=stored.content_type,
size_bytes=stored.size_bytes,
is_primary=True,
)
return Response(AssetSerializer(asset).data, status=status.HTTP_201_CREATED)