From d457e60121a42b6637a9cbb4b9cd9bf230463682 Mon Sep 17 00:00:00 2001
From: Martin Wacker <martas@imm.cz>
Date: Fri, 12 Jul 2024 14:55:37 +0200
Subject: [PATCH] fix: project file clean-up

---
 api/activities/activity.py           |  6 +++
 ayon_server/helpers/project_files.py | 69 ++++++++++++++++++++++------
 2 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/api/activities/activity.py b/api/activities/activity.py
index eac5f371..690c6b05 100644
--- a/api/activities/activity.py
+++ b/api/activities/activity.py
@@ -48,6 +48,7 @@ async def post_project_activity(
     entity_id: PathEntityID,
     user: CurrentUser,
     activity: ProjectActivityPostModel,
+    background_tasks: BackgroundTasks,
     x_sender: str | None = Header(default=None),
 ) -> CreateActivityResponseModel:
     """Create an activity.
@@ -78,6 +79,8 @@ async def post_project_activity(
         data=activity.data,
     )
 
+    background_tasks.add_task(delete_unused_files, project_name)
+
     return CreateActivityResponseModel(id=id)
 
 
@@ -86,6 +89,7 @@ async def delete_project_activity(
     project_name: ProjectName,
     activity_id: str,
     user: CurrentUser,
+    background_tasks: BackgroundTasks,
     x_sender: str | None = Header(default=None),
 ) -> EmptyResponse:
     """Delete an activity.
@@ -106,6 +110,8 @@ async def delete_project_activity(
         sender=x_sender,
     )
 
+    background_tasks.add_task(delete_unused_files, project_name)
+
     return EmptyResponse()
 
 
diff --git a/ayon_server/helpers/project_files.py b/ayon_server/helpers/project_files.py
index 31116887..c140ecea 100644
--- a/ayon_server/helpers/project_files.py
+++ b/ayon_server/helpers/project_files.py
@@ -19,6 +19,61 @@ def id_to_path(project_name: str, file_id: str) -> str:
     )
 
 
+async def delete_project_file(project_name: str, file_id: str) -> None:
+    """Delete a project file"""
+
+    query = f"""
+        DELETE FROM project_{project_name}.files
+        WHERE id = $1
+    """
+    await Postgres.execute(query, file_id)
+
+    file_id = str(file_id).replace("-", "")
+    assert len(file_id) == 32
+
+    query = f"""
+        WITH updated_activities AS (
+            SELECT
+                id,
+                jsonb_set(
+                    data,
+                    '{{files}}',
+                    (SELECT jsonb_agg(elem)
+                         FROM jsonb_array_elements(data->'files') elem
+                         WHERE elem->>'id' != '{file_id}')
+                ) AS new_data
+            FROM
+                project_{project_name}.activities
+            WHERE
+                data->'files' @> jsonb_build_array(
+                    jsonb_build_object('id', '{file_id}')
+                )
+        )
+        UPDATE project_{project_name}.activities
+        SET data = updated_activities.new_data
+        FROM updated_activities
+        WHERE activities.id = updated_activities.id;
+    """
+
+    await Postgres.execute(query)
+
+    path = id_to_path(project_name, file_id)
+    if not os.path.exists(path):
+        return
+
+    try:
+        os.remove(path)
+    except Exception as e:
+        logging.error(f"Failed to delete file {path}: {e}")
+
+    directory = os.path.dirname(path)
+    if not os.listdir(directory):
+        try:
+            os.rmdir(directory)
+        except Exception as e:
+            logging.error(f"Failed to delete directory {directory}: {e}")
+
+
 async def delete_unused_files(project_name: str) -> None:
     """Delete files that are not referenced in any activity."""
 
@@ -30,16 +85,4 @@ async def delete_unused_files(project_name: str) -> None:
 
     async for row in Postgres.iterate(query):
         logging.debug(f"Deleting unused file {row['id']}")
-        query = f"""
-            DELETE FROM project_{project_name}.files
-            WHERE id = $1
-        """
-        await Postgres.execute(query, row["id"])
-
-        path = id_to_path(project_name, row["id"])
-        try:
-            os.remove(path)
-        except FileNotFoundError:
-            pass
-        except Exception as e:
-            logging.error(f"Failed to delete file {path}: {e}")
+        await delete_project_file(project_name, row["id"])