From dbfa438bb6982eebf7feb94b997979c020357706 Mon Sep 17 00:00:00 2001
From: sophia <scastellarin@quansight.com>
Date: Wed, 15 Jan 2025 15:01:33 -0800
Subject: [PATCH] Add task to archive builds

---
 .../_internal/server/views/api.py             |  2 +-
 .../_internal/worker/tasks.py                 | 24 +++++++++++++++++++
 .../conda_store_server/conda_store.py         | 14 ++++++-----
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/conda-store-server/conda_store_server/_internal/server/views/api.py b/conda-store-server/conda_store_server/_internal/server/views/api.py
index 41ba4256b..24dc973f8 100644
--- a/conda-store-server/conda_store_server/_internal/server/views/api.py
+++ b/conda-store-server/conda_store_server/_internal/server/views/api.py
@@ -1151,7 +1151,7 @@ async def api_archive_build(
         except CondaStoreError as e:
             raise HTTPException(status_code=400, detail=e.message)
 
-        return {"status": schema.APIStatus.OK}
+    return {"status": schema.APIStatus.OK}
 
 
 @router_api.get(
diff --git a/conda-store-server/conda_store_server/_internal/worker/tasks.py b/conda-store-server/conda_store_server/_internal/worker/tasks.py
index 4959a0cf1..28a82ad5f 100644
--- a/conda-store-server/conda_store_server/_internal/worker/tasks.py
+++ b/conda-store-server/conda_store_server/_internal/worker/tasks.py
@@ -304,6 +304,30 @@ def task_delete_build(self, build_id):
         db.commit()
 
 
+@shared_task(base=WorkerTask, name="task_archive_build", bind=True)
+def task_archive_build(self, build_id):
+    conda_store = self.worker.conda_store
+    with conda_store.session_factory() as db:
+        build = api.get_build(db, build_id)
+        
+        archive_save_artifacts = [schema.BuildArtifactType.LOCKFILE, schema.BuildArtifactType.LOGS]
+        # Deletes build artifacts for this build
+        conda_store.log.info(f"archiving build={build.id}")
+        for build_artifact in api.list_build_artifacts(
+            db,
+            build_id=build_id,
+            excluded_artifact_types=archive_save_artifacts,
+        ).all():
+            conda_store.log.info(f"archiving build={build.id} - deleting artifact {build_artifact.artifact_type.value}")
+            delete_build_artifact(db, conda_store, build_artifact)
+
+        # Updates build size and marks build as deleted
+        build.archived_on = datetime.datetime.utcnow()
+        build.status = schema.BuildStatus.ARCHIVED
+        build.size = 0
+        db.commit()
+
+
 @shared_task(base=WorkerTask, name="task_delete_environment", bind=True)
 def task_delete_environment(self, environment_id):
     conda_store = self.worker.conda_store
diff --git a/conda-store-server/conda_store_server/conda_store.py b/conda-store-server/conda_store_server/conda_store.py
index 0431d345b..0d72215b8 100644
--- a/conda-store-server/conda_store_server/conda_store.py
+++ b/conda-store-server/conda_store_server/conda_store.py
@@ -562,10 +562,12 @@ def archive_build(self, db: Session, build_id: int):
             action=auth_schema.Permissions.BUILD_ARCHIVE,
         )
 
-        if build.status not in [
-            schema.BuildStatus.FAILED,
-            schema.BuildStatus.COMPLETED,
-        ]:
-            raise CondaStoreError("cannot archive build since not finished building")
+        if build.status != schema.BuildStatus.COMPLETED:
+            raise CondaStoreError(f"cannot archive build that is in the {build.status.value} state")
+
+        self.celery_app
+
+        # must import tasks after a celery app has been initialized
+        from conda_store_server._internal.worker import tasks
 
-        # TODO: archive build
+        tasks.task_archive_build.si(build.id).apply_async()