From 7684798f85274133e6a4b13f64c3cfb446b020de Mon Sep 17 00:00:00 2001
From: Taha Yassine Kraiem <tahayk2@gmail.com>
Date: Mon, 17 Apr 2023 13:19:32 +0100
Subject: [PATCH] feat(chalice): tag mob files for scheduled delete

---
 api/chalicelib/core/sessions_devtool.py    |  2 +-
 api/chalicelib/core/sessions_mobs.py       |  2 +-
 api/chalicelib/utils/s3.py                 | 33 +++++++++++-----------
 ee/api/chalicelib/core/sessions_devtool.py |  2 +-
 ee/api/chalicelib/utils/s3_extra.py        | 25 ++++++----------
 5 files changed, 28 insertions(+), 36 deletions(-)

diff --git a/api/chalicelib/core/sessions_devtool.py b/api/chalicelib/core/sessions_devtool.py
index 714718d1c..fb2a42b14 100644
--- a/api/chalicelib/core/sessions_devtool.py
+++ b/api/chalicelib/core/sessions_devtool.py
@@ -29,4 +29,4 @@ def get_urls(session_id, project_id, check_existence: bool = True):
 def delete_mobs(project_id, session_ids):
     for session_id in session_ids:
         for k in __get_devtools_keys(project_id=project_id, session_id=session_id):
-            s3.schedule_for_deletion(bucket=config("sessions_bucket"), key=k)
+            s3.tag_for_deletion(bucket=config("sessions_bucket"), key=k)
diff --git a/api/chalicelib/core/sessions_mobs.py b/api/chalicelib/core/sessions_mobs.py
index 2ac5c4257..ead03505f 100644
--- a/api/chalicelib/core/sessions_mobs.py
+++ b/api/chalicelib/core/sessions_mobs.py
@@ -59,4 +59,4 @@ def delete_mobs(project_id, session_ids):
     for session_id in session_ids:
         for k in __get_mob_keys(project_id=project_id, session_id=session_id) \
                  + __get_mob_keys_deprecated(session_id=session_id):
-            s3.schedule_for_deletion(bucket=config("sessions_bucket"), key=k)
+            s3.tag_for_deletion(bucket=config("sessions_bucket"), key=k)
diff --git a/api/chalicelib/utils/s3.py b/api/chalicelib/utils/s3.py
index f98b541b4..cd4b84f1b 100644
--- a/api/chalicelib/utils/s3.py
+++ b/api/chalicelib/utils/s3.py
@@ -1,5 +1,4 @@
 import hashlib
-from datetime import datetime, timedelta
 from urllib.parse import urlparse
 
 import boto3
@@ -109,23 +108,10 @@ def rename(source_bucket, source_key, target_bucket, target_key):
     s3.Object(source_bucket, source_key).delete()
 
 
-def schedule_for_deletion(bucket, key):
+def tag_for_deletion(bucket, key):
     if not exists(bucket, key):
         return False
-    s3 = __get_s3_resource()
-
-    s3_origin = s3.Object(bucket, key)
-    target_key = key.split("/")
-    target_key[-1] = "del_" + target_key[-1]
-    target_key = "/".join(target_key)
-    s3_target = s3.Object(bucket, target_key)
-    s3_target.copy_from(
-        CopySource={'Bucket': bucket, 'Key': key},
-        Expires=datetime.utcnow(),
-        MetadataDirective='COPY',
-        TaggingDirective='COPY'
-    )
-    s3_origin.delete()
+    tag_file(bucket=bucket, file_key=key, tag_key='to_delete_in_days', tag_value='7')
 
 
 def generate_file_key(project_id, key):
@@ -136,3 +122,18 @@ def generate_file_key_from_url(project_id, url):
     u = urlparse(url)
     new_url = u.scheme + "://" + u.netloc + u.path
     return generate_file_key(project_id=project_id, key=new_url)
+
+
+def tag_file(file_key, bucket, tag_key, tag_value):
+    return client.put_object_tagging(
+        Bucket=bucket,
+        Key=file_key,
+        Tagging={
+            'TagSet': [
+                {
+                    'Key': tag_key,
+                    'Value': tag_value
+                },
+            ]
+        }
+    )
diff --git a/ee/api/chalicelib/core/sessions_devtool.py b/ee/api/chalicelib/core/sessions_devtool.py
index 9cb393f17..2e442195c 100644
--- a/ee/api/chalicelib/core/sessions_devtool.py
+++ b/ee/api/chalicelib/core/sessions_devtool.py
@@ -36,4 +36,4 @@ def get_urls(session_id, project_id, context: schemas_ee.CurrentContext, check_e
 def delete_mobs(project_id, session_ids):
     for session_id in session_ids:
         for k in __get_devtools_keys(project_id=project_id, session_id=session_id):
-            s3.schedule_for_deletion(bucket=config("sessions_bucket"), key=k)
+            s3.tag_for_deletion(bucket=config("sessions_bucket"), key=k)
diff --git a/ee/api/chalicelib/utils/s3_extra.py b/ee/api/chalicelib/utils/s3_extra.py
index d561aa6bf..1aa0d5b40 100644
--- a/ee/api/chalicelib/utils/s3_extra.py
+++ b/ee/api/chalicelib/utils/s3_extra.py
@@ -4,23 +4,14 @@ from chalicelib.utils import s3
 
 
 def tag_session(file_key, tag_key='retention', tag_value='vault'):
-    return tag_file(file_key=file_key, bucket=config("sessions_bucket"), tag_key=tag_key, tag_value=tag_value)
+    bucket = config("sessions_bucket")
+    if not s3.exists(bucket=bucket, key=file_key):
+        return None
+    return s3.tag_file(file_key=file_key, bucket=bucket, tag_key=tag_key, tag_value=tag_value)
 
 
 def tag_record(file_key, tag_key='retention', tag_value='vault'):
-    return tag_file(file_key=file_key, bucket=config('ASSIST_RECORDS_BUCKET'), tag_key=tag_key, tag_value=tag_value)
-
-
-def tag_file(file_key, bucket, tag_key, tag_value):
-    return s3.client.put_object_tagging(
-        Bucket=bucket,
-        Key=file_key,
-        Tagging={
-            'TagSet': [
-                {
-                    'Key': tag_key,
-                    'Value': tag_value
-                },
-            ]
-        }
-    )
+    bucket = config('ASSIST_RECORDS_BUCKET')
+    if not s3.exists(bucket=bucket, key=file_key):
+        return None
+    return s3.tag_file(file_key=file_key, bucket=bucket, tag_key=tag_key, tag_value=tag_value)