From bcaff00455ebb6b71d74e5078625449a63b6ebcf Mon Sep 17 00:00:00 2001
From: Aakash Singh <mail@singhaakash.dev>
Date: Mon, 12 Sep 2022 17:47:06 +0530
Subject: [PATCH] return file extension from file upload views

---
 care/facility/api/serializers/file_upload.py | 34 ++++++++++++++++----
 care/facility/models/file_upload.py          | 20 +++++++++---
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/care/facility/api/serializers/file_upload.py b/care/facility/api/serializers/file_upload.py
index 25b0592eed..f00d3a8ccc 100644
--- a/care/facility/api/serializers/file_upload.py
+++ b/care/facility/api/serializers/file_upload.py
@@ -49,7 +49,12 @@ def check_permissions(file_type, associating_id, user):
                     if user == sample.consultation.assigned_to:
                         return sample.id
             if sample.testing_facility:
-                if has_facility_permission(user, Facility.objects.get(external_id=sample.testing_facility.external_id)):
+                if has_facility_permission(
+                    user,
+                    Facility.objects.get(
+                        external_id=sample.testing_facility.external_id
+                    ),
+                ):
                     return sample.id
             if not has_facility_permission(user, patient.facility):
                 raise Exception("No Permission")
@@ -86,7 +91,9 @@ class Meta:
 
     def create(self, validated_data):
         user = self.context["request"].user
-        internal_id = check_permissions(validated_data["file_type"], validated_data["associating_id"], user)
+        internal_id = check_permissions(
+            validated_data["file_type"], validated_data["associating_id"], user
+        )
         validated_data["associating_id"] = internal_id
         validated_data["uploaded_by"] = user
         validated_data["internal_name"] = validated_data["original_name"]
@@ -97,12 +104,19 @@ def create(self, validated_data):
 class FileUploadListSerializer(serializers.ModelSerializer):
 
     id = serializers.UUIDField(source="external_id", read_only=True)
-
     uploaded_by = UserBaseMinimumSerializer(read_only=True)
+    extension = serializers.CharField(source="get_extension", read_only=True)
 
     class Meta:
         model = FileUpload
-        fields = ("id", "name", "uploaded_by", "created_date", "file_category")
+        fields = (
+            "id",
+            "name",
+            "uploaded_by",
+            "created_date",
+            "file_category",
+            "extension",
+        )
         read_only_fields = ("associating_id", "name", "created_date")
 
 
@@ -118,11 +132,19 @@ class Meta:
 class FileUploadRetrieveSerializer(serializers.ModelSerializer):
 
     id = serializers.UUIDField(source="external_id", read_only=True)
-
     uploaded_by = UserBaseMinimumSerializer(read_only=True)
     read_signed_url = serializers.CharField(read_only=True)
+    extension = serializers.CharField(source="get_extension", read_only=True)
 
     class Meta:
         model = FileUpload
-        fields = ("id", "name", "uploaded_by", "created_date", "read_signed_url", "file_category")
+        fields = (
+            "id",
+            "name",
+            "uploaded_by",
+            "created_date",
+            "read_signed_url",
+            "file_category",
+            "extension",
+        )
         read_only_fields = ("associating_id", "name", "created_date")
diff --git a/care/facility/models/file_upload.py b/care/facility/models/file_upload.py
index bc8c815356..9dfbc564b3 100644
--- a/care/facility/models/file_upload.py
+++ b/care/facility/models/file_upload.py
@@ -1,13 +1,13 @@
 import enum
 import time
-import boto3
 from uuid import uuid4
+
+import boto3
 from django.conf import settings
 from django.db import models
 
 from care.facility.models import FacilityBaseModel
 from care.users.models import User
-
 from care.utils.csp import config as cs_provider
 
 
@@ -40,12 +40,22 @@ class FileCategory(enum.Enum):
     internal_name = models.CharField(max_length=2000)
     associating_id = models.CharField(max_length=100, blank=False, null=False)
     upload_completed = models.BooleanField(default=False)
-    uploaded_by = models.ForeignKey(User, on_delete=models.PROTECT, null=True, blank=True)
-    file_type = models.IntegerField(choices=FileTypeChoices, default=FileType.PATIENT.value)
+    uploaded_by = models.ForeignKey(
+        User, on_delete=models.PROTECT, null=True, blank=True
+    )
+    file_type = models.IntegerField(
+        choices=FileTypeChoices, default=FileType.PATIENT.value
+    )
     file_category = models.CharField(
-        choices=FileCategoryChoices, default=FileCategory.UNSPECIFIED.value, max_length=100
+        choices=FileCategoryChoices,
+        default=FileCategory.UNSPECIFIED.value,
+        max_length=100,
     )
 
+    def get_extension(self):
+        parts = self.internal_name.split(".")
+        return f".{parts[-1]}" if len(parts) > 1 else ""
+
     def save(self, *args, **kwargs):
         if "force_insert" in kwargs or (not self.internal_name):
             internal_name = str(uuid4()) + str(int(time.time()))