From ee5999da91e851865b8fc1aa7242c753b301a99e Mon Sep 17 00:00:00 2001
From: Joel Mathew Thomas <90510078+joelmathewthomas@users.noreply.github.com>
Date: Wed, 26 Feb 2025 20:05:49 +0530
Subject: [PATCH] endpoint: /api/cleanup

- Add new endpoint /api/cleanup, params: file_uuid
- Cleanup leftover files on the server
---
 api/api/tasks.py    | 12 ++++++++++++
 api/api/views.py    | 18 +++++++++++++++++-
 api/backend/urls.py |  5 ++++-
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/api/api/tasks.py b/api/api/tasks.py
index 74de546..716bbec 100644
--- a/api/api/tasks.py
+++ b/api/api/tasks.py
@@ -118,5 +118,17 @@ def noisereduce_task(file_path):
     try:
         noisereduce(file_path, file_path)
         return True
+    except Exception as e:
+        return False
+
+@shared_task
+def cleanup_task(file_path):
+    """Celery task to cleanup files"""
+    file_path = Path(file_path)
+    
+    # Cleanup
+    try:
+        shutil.rmtree(os.path.dirname(file_path))
+        return True
     except Exception as e:
         return False
\ No newline at end of file
diff --git a/api/api/views.py b/api/api/views.py
index e4fe3fa..507cc6f 100644
--- a/api/api/views.py
+++ b/api/api/views.py
@@ -12,6 +12,7 @@ from .tasks import trim_audio_task
 from .tasks import resample_audio_task
 from .tasks import music_separation_task
 from .tasks import noisereduce_task
+from .tasks import cleanup_task
 from freqsplit.input.format_checker import is_supported_format
 
 UPLOAD_DIR = "/tmp/freqsplit"
@@ -180,4 +181,19 @@ def download_audio(request):
                 zipf.write(file_path, arcname)
 
     # Stream the ZIP file
-    return FileResponse(open(zip_file_path, "rb"), as_attachment=True, filename=os.path.basename(zip_file_path))
\ No newline at end of file
+    return FileResponse(open(zip_file_path, "rb"), as_attachment=True, filename=os.path.basename(zip_file_path))
+
+@api_view(['POST'])
+def cleanup(request):
+    """Handles file cleanup after pipeline processing"""
+    stat, result, status_code = get_audio_file_path(request, UPLOAD_DIR)
+    if stat == False:
+        return Response({"error": result}, status=status_code)  
+
+    # Call Celery task synchronously
+    task = cleanup_task.apply(args=(result,))
+
+    if task.get():
+        return Response({"message": f"Successfully cleaned up files on the server"}, status=status.HTTP_200_OK)
+    else:
+        return Response({"error": "Failed to cleanup files on the server"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
\ No newline at end of file
diff --git a/api/backend/urls.py b/api/backend/urls.py
index 63be3f9..2aec782 100644
--- a/api/backend/urls.py
+++ b/api/backend/urls.py
@@ -23,6 +23,7 @@ from api.views import resample_audio
 from api.views import separate_music
 from api.views import noisereduce
 from api.views import download_audio
+from api.views import cleanup
 
 urlpatterns = [
     path('admin/', admin.site.urls),
@@ -32,5 +33,7 @@ urlpatterns = [
     path('api/resample', resample_audio, name='resample_audio'),
     path('api/separate', separate_music, name="separate_music"),
     path('api/noisereduce', noisereduce, name="noisreduce"),
-    path('api/download', download_audio, name="download_audio")
+    path('api/download', download_audio, name="download_audio"),
+    path('api/cleanup', cleanup, name="cleanup"),
+    
 ]