[ext] On annotation deletion alse delete from export storage (HumanSi…

…gnal#1529)

label_studio/core/settings/base.py

@@ -18,7 +18,7 @@
     logging.basicConfig(level=logging.DEBUG, format='%(message)s')
 
 from label_studio.core.utils.io import get_data_dir
-from label_studio.core.utils.params import get_bool_env, get_env
+from label_studio.core.utils.params import get_bool_env, get_env, get_env_list_int
 
 logger = logging.getLogger(__name__)
 
@@ -368,6 +368,8 @@
 TASK_LOCK_DEFAULT_TTL = int(get_env('TASK_LOCK_DEFAULT_TTL', 3600))
 TASK_LOCK_MIN_TTL = int(get_env('TASK_LOCK_MIN_TTL', 120))
 
+DELETION_FROM_S3_ENABLED_FOR_ORGS = get_env_list_int('DELETION_FROM_S3_ENABLED_FOR_ORGS', [])
+
 # Email backend
 FROM_EMAIL = get_env('FROM_EMAIL', 'Label Studio <[email protected]>')
 EMAIL_BACKEND = get_env('EMAIL_BACKEND', 'django.core.mail.backends.dummy.EmailBackend')

label_studio/core/utils/params.py

@@ -90,3 +90,15 @@ def get_env(name, default=None, is_bool=False):
 
 def get_bool_env(key, default):
     return get_env(key, default, is_bool=True)
+
+
+def get_env_list_int(key, default=None):
+    """
+    "1,2,3" in env variable => [1, 2, 3] in python
+    """
+    value = get_env(key)
+    if not value:
+        if default is None:
+            return []
+        return default
+    return [int(el) for el in value.split(',')]

label_studio/io_storages/s3/models.py

@@ -5,12 +5,12 @@
 import json
 import boto3
 
-from botocore.exceptions import NoCredentialsError
+from botocore.exceptions import NoCredentialsError, ClientError
 from django.db import models, transaction
 from django.conf import settings
 from django.utils.translation import gettext_lazy as _
 from django.dispatch import receiver
-from django.db.models.signals import post_save
+from django.db.models.signals import post_save, post_delete
 
 from io_storages.base_models import ImportStorage, ImportStorageLink, ExportStorage, ExportStorageLink
 from io_storages.utils import get_uri_via_regex
@@ -194,6 +194,27 @@ def save_annotation(self, annotation):
         # create link if everything ok
         S3ExportStorageLink.create(annotation, self)
 
+    def delete_annotation(self, annotation):
+        client, s3 = self.get_client_and_resource()
+        logger.debug(f'Deleting object on {self.__class__.__name__} Storage {self} for annotation {annotation}')
+
+        # get key that identifies this object in storage
+        key = S3ExportStorageLink.get_key(annotation)
+        key = str(self.prefix) + '/' + key if self.prefix else key
+
+        # delete object from storage
+        try:
+            s3.Object(self.bucket, key).delete()
+        except ClientError as e:
+            # we ignore access denied errors
+            logger.exception(e)
+            if not 'Access Denied' in str(e):
+                raise
+
+
+        # delete link if everything ok
+        S3ExportStorageLink.objects.filter(storage=self, annotation=annotation).delete()
+
 
 @receiver(post_save, sender=Annotation)
 def export_annotation_to_s3_storages(sender, instance, **kwargs):
@@ -204,6 +225,16 @@ def export_annotation_to_s3_storages(sender, instance, **kwargs):
             storage.save_annotation(instance)
 
 
+@receiver(post_delete, sender=Annotation)
+def delete_annotation_from_s3_storages(sender, instance, **kwargs):
+    project = instance.task.project
+    if project.organization_id in settings.DELETION_FROM_S3_ENABLED_FOR_ORGS:
+        if hasattr(project, 'io_storages_s3exportstorages'):
+            for storage in project.io_storages_s3exportstorages.all():
+                logger.debug(f'Delete {instance} from S3 storage {storage}')
+                storage.delete_annotation(instance)
+
+
 class S3ImportStorageLink(ImportStorageLink):
     storage = models.ForeignKey(S3ImportStorage, on_delete=models.CASCADE, related_name='links')