Merge pull request #897 from kobotoolbox/regroup-attachment-signals

Regroup storage counter calculation in the same DB transaction

onadata/apps/logger/management/commands/move_media_to_s3.py

@@ -3,7 +3,7 @@
 # coding: utf-8
 import sys
 
-from django.core.files.storage import default_storage
+from django.core.files.storage import default_storage, get_storage_class
 from django.core.management.base import BaseCommand
 
 from onadata.apps.logger.models.attachment import Attachment

onadata/apps/logger/signals.py

@@ -1,10 +1,12 @@
 # coding: utf-8
 import logging
 
+from django.core.files.storage import default_storage
+from django.db import transaction
 from django.db.models import F
 from django.db.models.signals import (
-    pre_delete,
     post_save,
+    pre_delete,
 )
 from django.dispatch import receiver
 
@@ -13,62 +15,55 @@
 from onadata.apps.main.models.user_profile import UserProfile
 
 
-@receiver(pre_delete, sender=Attachment)
-def delete_attachment_subtract_user_profile_storage_bytes(instance, **kwargs):
-    """
-    Update the attachment_storage_bytes field in the UserProfile model
-    when an attachment is deleted
-    """
-    attachment = instance
-    file_size = attachment.media_file_size
-    if not file_size:
-        return
-
-    queryset = UserProfile.objects.filter(
-        user_id=attachment.instance.xform.user_id
-    )
-    queryset.update(
-        attachment_storage_bytes=F('attachment_storage_bytes') - file_size
-    )
-
-
-@receiver(pre_delete, sender=Attachment)
-def delete_attachment_subtract_xform_storage_bytes(instance, **kwargs):
-    """
-    Update the attachment_storage_bytes field in the XForm
-    model when an attachment is deleted
-    """
-    attachment = instance
-    file_size = attachment.media_file_size
-    if not file_size:
-        return
-
-    xform_id = instance.instance.xform.pk
-    queryset = XForm.objects.filter(pk=xform_id)
-    queryset.update(
-        attachment_storage_bytes=F('attachment_storage_bytes') - file_size
-    )
-
-
 @receiver(pre_delete, sender=Attachment)
 def pre_delete_attachment(instance, **kwargs):
     # "Model.delete() isn’t called on related models, but the pre_delete and
     # post_delete signals are sent for all deleted objects." See
-    # https://docs.djangoproject.com/en/2.2/ref/models/fields/#django.db.models.CASCADE
-    # We want to delete all files when an Instance (or Attachment) object is
+    # https://docs.djangoproject.com/en/3.2/ref/models/fields/#django.db.models.CASCADE
+    # We want to delete all files when an XForm, an Instance or Attachment object is
     # deleted.
+    # Since the Attachment object is deleted with CASCADE, we must use a
+    # `pre_delete` signal to access its parent Instance and its parent XForm.
+    # Otherwise, with a `post_delete`, they would be gone before reaching the rest
+    # of code below.
 
     # `instance` here means "model instance", and no, it is not allowed to
     # change the name of the parameter
     attachment = instance
+    file_size = attachment.media_file_size
+
+    xform = attachment.instance.xform
+
+    if file_size:
+        with transaction.atomic():
+            """
+            Update both counters at the same time (in a transaction) to avoid 
+            desynchronization as much as possible 
+            """
+            UserProfile.objects.filter(
+                user_id=xform.user_id
+            ).update(
+                attachment_storage_bytes=F('attachment_storage_bytes') - file_size
+            )
+            XForm.all_objects.filter(pk=xform.pk).update(
+                attachment_storage_bytes=F('attachment_storage_bytes') - file_size
+            )
+
+    if not (media_file_name := str(attachment.media_file)):
+        return
+
+    # Clean-up storage
     try:
-        attachment.media_file.delete()
+        # We do not want to call `attachment.media_file.delete()` because it calls
+        # `attachment.save()` behind the scene which would call again the `post_save`
+        # signal below. Bonus: it avoids an extra query 😎.
+        default_storage.delete(media_file_name)
     except Exception as e:
         logging.error('Failed to delete attachment: ' + str(e), exc_info=True)
 
 
 @receiver(post_save, sender=Attachment)
-def update_user_profile_attachment_storage_bytes(instance, created, **kwargs):
+def post_save_attachment(instance, created, **kwargs):
     """
     Update the attachment_storage_bytes field in the UserProfile model
     when an attachment is added
@@ -78,34 +73,17 @@ def update_user_profile_attachment_storage_bytes(instance, created, **kwargs):
     attachment = instance
     if getattr(attachment, 'defer_counting', False):
         return
-    file_size = attachment.media_file_size
-    if not file_size:
-        return
-
-    queryset = UserProfile.objects.filter(user_id=attachment.instance.xform.user_id)
-    queryset.update(
-        attachment_storage_bytes=F('attachment_storage_bytes') + file_size
-    )
-
 
-@receiver(post_save, sender=Attachment)
-def update_xform_attachment_storage_bytes(instance, created, **kwargs):
-    """
-    Update the attachment_storage_bytes field in the XForm model when
-    an attachment is added
-    """
-    if not created:
-        return
-    attachment = instance
-    if getattr(attachment, 'defer_counting', False):
-        return
     file_size = attachment.media_file_size
-
     if not file_size:
         return
 
-    xform_id = attachment.instance.xform_id
-    queryset = XForm.objects.filter(pk=xform_id)
-    queryset.update(
-        attachment_storage_bytes=F('attachment_storage_bytes') + file_size
-    )
+    xform = attachment.instance.xform
+
+    with transaction.atomic():
+        UserProfile.objects.filter(user_id=xform.user_id).update(
+            attachment_storage_bytes=F('attachment_storage_bytes') + file_size
+        )
+        XForm.objects.filter(pk=xform.pk).update(
+            attachment_storage_bytes=F('attachment_storage_bytes') + file_size
+        )

onadata/libs/utils/logger_tools.py

@@ -55,10 +55,7 @@
     update_xform_submission_count,
 )
 from onadata.apps.logger.models.xform import XLSFormError
-from onadata.apps.logger.signals import (
-    update_user_profile_attachment_storage_bytes,
-    update_xform_attachment_storage_bytes,
-)
+from onadata.apps.logger.signals import post_save_attachment
 from onadata.apps.logger.xform_instance_parser import (
     InstanceEmptyError,
     InstanceInvalidUserError,
@@ -676,12 +673,7 @@ def save_submission(
         if getattr(new_attachment, 'defer_counting', False):
             # Remove the Python-only attribute
             del new_attachment.defer_counting
-            update_user_profile_attachment_storage_bytes(
-                instance=new_attachment, created=True
-            )
-            update_xform_attachment_storage_bytes(
-                instance=new_attachment, created=True
-            )
+            post_save_attachment(new_attachment, created=True)
 
     return instance