Add mime_type field to UserFile model

Fix 0037 migration in users app

Author: yakser

chats/consumers/chat.py

@@ -65,7 +65,7 @@ async def connect(self):
         await self.accept()
 
     async def disconnect(self, close_code):
-        """User disconnected from websocket, Don't have to do anything here"""
+        """User disconnected from websocket"""
         pass
 
     async def receive_json(self, content, **kwargs):

files/helpers.py

@@ -1,6 +1,10 @@
 import requests
 import time
+import magic
+from django.core.files.uploadedfile import TemporaryUploadedFile
+
 from files.exceptions import SelectelUploadError
+from files.models import UserFile
 
 from procollab.settings import (
     DEBUG,
@@ -12,9 +16,11 @@
 
 
 class FileAPI:
-    def __init__(self, file, user) -> None:
-        self.file = file
+    def __init__(self, file: TemporaryUploadedFile, user) -> None:
+        self.file = file  # it's TemporaryUploadedFile, and it will be
+        # removed after first .close() call, so we must read this file only once
         self.user = user
+        self.file_object = self.file.open(mode="rb")
 
     @staticmethod
     def delete(url: str) -> int:
@@ -23,21 +29,34 @@ def delete(url: str) -> int:
         response = requests.delete(url, headers={"X-Auth-Token": token})
         return response.status_code
 
-    def upload(self):
-        return self._upload_via_selectel_swift()
+    def upload(self) -> str:
+        url = self._upload_via_selectel_swift()
+        info = get_file_info(self.file)
+        UserFile.objects.create(
+            user=self.user,
+            link=url,
+            name=info["name"],
+            size=info["size"],
+            extension=info["extension"],
+            mime_type=info["mime_type"],
+        )
+        self.file_object.close()
+        return url
 
-    def _upload_via_selectel_swift(self) -> tuple[int, str]:
+    def _upload_via_selectel_swift(self) -> str:
         token = self._get_selectel_swift_token()
         url = self._generate_selectel_swift_file_url()
 
-        with self.file.open(mode="rb") as file_object:
-            response = requests.put(
-                url,
-                headers={"X-Auth-Token": token, "Content-Type": file_object.content_type},
-                data=file_object.read(),
-            )
+        requests.put(
+            url,
+            headers={
+                "X-Auth-Token": token,
+                "Content-Type": self.file_object.content_type,
+            },
+            data=self.file_object.read(),
+        )
 
-        return response.status_code, url
+        return url
 
     def _generate_selectel_swift_link(sefl):
         link = f"https://api.selcdn.ru/v1/SEL_{SELECTEL_ACCOUNT_ID}/{SELECTEL_CONTAINER_NAME}/"
@@ -61,10 +80,10 @@ def _get_selectel_swift_token():
                 }
             }
         }
-        r = requests.post("https://api.selcdn.ru/v3/auth/tokens", json=data)
-        if r.status_code not in [200, 201]:
+        response = requests.post("https://api.selcdn.ru/v3/auth/tokens", json=data)
+        if response.status_code not in [200, 201]:
             raise SelectelUploadError("Couldn't generate a token for selcdn")
-        return r.headers["x-subject-token"]
+        return response.headers["x-subject-token"]
 
     def _get_file_extension(self) -> str:
         if len(self.file.name.split(".")) > 1:
@@ -85,11 +104,16 @@ def _generate_selectel_swift_file_url(self) -> str:
         )
 
 
-def get_file_info(request):
-    name, ext = request.name.split(".")
+def get_file_info(file: TemporaryUploadedFile) -> dict:
+    name, ext = file.name.split(".")
 
     return {
-        "size": request.size,
+        "size": file.size,
         "name": name,
         "extension": ext,
+        "mime_type": get_file_mime_type(file),
     }
+
+
+def get_file_mime_type(file: TemporaryUploadedFile):
+    return magic.from_file(file.temporary_file_path(), mime=True)

files/migrations/0004_userfile_mime_type.py

@@ -0,0 +1,18 @@
+# Generated by Django 4.1.3 on 2023-07-08 17:58
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("files", "0003_userfile_extension_userfile_name_userfile_size"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="userfile",
+            name="mime_type",
+            field=models.CharField(default="", max_length=256),
+        ),
+    ]

files/models.py

@@ -22,6 +22,7 @@ class UserFile(models.Model):
     datetime_uploaded = models.DateTimeField(auto_now_add=True)
     name = models.TextField(blank=False, default="file")
     extension = models.TextField(blank=True, default="")
+    mime_type = models.CharField(max_length=256, default="")
     size = models.PositiveBigIntegerField(null=False, blank=True, default=1)
 
     def __str__(self):

files/serializers.py

@@ -9,6 +9,7 @@ class Meta:
         fields = [
             "name",
             "extension",
+            "mime_type",
             "size",
             "link",
             "user",

files/views.py

@@ -4,7 +4,7 @@
 from rest_framework.generics import get_object_or_404
 from rest_framework.response import Response
 
-from files.helpers import FileAPI, get_file_info
+from files.helpers import FileAPI
 from files.models import UserFile
 from files.serializers import UserFileSerializer
 
@@ -17,40 +17,33 @@ class FileView(generics.GenericAPIView):
     @transaction.atomic
     def post(self, request):
         """creates a UserFile object and uploads the file to selectel"""
-        file_api = FileAPI(request.FILES["file"], request.user)
-        status_code, url = file_api.upload()
-
-        if status_code == 201:
-            info = get_file_info(request.FILES["file"])
-            UserFile.objects.create(
-                user=request.user,
-                link=url,
-                name=info["name"],
-                size=info["size"],
-                extension=info["extension"],
-            )
+        try:
+            file_api = FileAPI(request.FILES["file"], request.user)
+            url = file_api.upload()
             return Response({"url": url}, status=status.HTTP_201_CREATED)
-
-        return Response("Failed to upload file", status=status.HTTP_409_CONFLICT)
-
-    def delete(self, request, *args, **kwargs):
-        """deletes the file (only if the request is sent by the user who owns it!)
-        The link has to be specified in the JSON body, not in the URL arguments.
-        """
-        # get the link from the query
-        if request.query_params and (request.query_params.get("link") is not None):
-            link = request.query_params.get("link")
-        else:
-            return Response(
-                {
-                    "error": "you have to pass the link of the object you want to delete in query parameters"
-                },
-                status=status.HTTP_400_BAD_REQUEST,
-            )
-        instance = get_object_or_404(self.get_queryset(), link=link)
-        if instance.user != request.user:
-            return Response(status=status.HTTP_403_FORBIDDEN)
-        FileAPI.delete(instance.link)  # delete the file via api
-        instance.delete()  # delete the UserFile object
-
-        return Response(status=status.HTTP_204_NO_CONTENT)
+        except Exception as error:
+            print(error)
+            return Response("Failed to upload file", status=status.HTTP_409_CONFLICT)
+
+
+def delete(self, request, *args, **kwargs):
+    """deletes the file (only if the request is sent by the user who owns it!)
+    The link has to be specified in the JSON body, not in the URL arguments.
+    """
+    # get the link from the query
+    if request.query_params and (request.query_params.get("link") is not None):
+        link = request.query_params.get("link")
+    else:
+        return Response(
+            {
+                "error": "you have to pass the link of the object you want to delete in query parameters"
+            },
+            status=status.HTTP_400_BAD_REQUEST,
+        )
+    instance = get_object_or_404(self.get_queryset(), link=link)
+    if instance.user != request.user:
+        return Response(status=status.HTTP_403_FORBIDDEN)
+    FileAPI.delete(instance.link)  # delete the file via api
+    instance.delete()  # delete the UserFile object
+
+    return Response(status=status.HTTP_204_NO_CONTENT)