From 82a75878b3b0be7a729642f629086e8c29bff08e Mon Sep 17 00:00:00 2001
From: Ruge Li <rugeli0605@gmail.com>
Date: Mon, 4 Mar 2024 16:35:44 -0800
Subject: [PATCH] check and delete expired results

---
 cellpack/autopack/DBRecipeHandler.py | 41 ++++++++++++++++++++++++++++
 cellpack/autopack/FirebaseHandler.py |  7 +++++
 2 files changed, 48 insertions(+)

diff --git a/cellpack/autopack/DBRecipeHandler.py b/cellpack/autopack/DBRecipeHandler.py
index ba5992d7..6d132ff8 100644
--- a/cellpack/autopack/DBRecipeHandler.py
+++ b/cellpack/autopack/DBRecipeHandler.py
@@ -2,6 +2,7 @@
 from enum import Enum
 
 from deepdiff import DeepDiff
+import requests
 
 from cellpack.autopack.utils import deep_merge
 
@@ -376,6 +377,30 @@ def should_write(self, db, grad_name):
         return None, None
 
 
+class ResultDoc(DataDoc):
+    def __init__(self, db):
+        self.db = db
+
+    def check_expired_results(self):
+        """
+        Check if the results in the database are expired and delete them if the linked object expired.
+        """
+        current_utc = self.db.create_timestamp()
+        results = self.db.collection("results").stream()
+        for result in results:
+            result_data = self.db.doc_to_dict(result)
+            result_age = current_utc - result_data["timestamp"]
+            if result_age.days > 180 and not self.check_url_expires(result_data["url"]):
+                self.db.delete_doc("results", self.db.doc_id(result))
+
+    def check_url_exists(self, url):
+        """
+        Check if the S3 object at the URL exists.
+        """
+        response = requests.head(url)
+        return response.status_code == 200
+
+
 class DBUploader(object):
     """
     Handles the uploading of data to the database.
@@ -710,3 +735,19 @@ def compile_db_recipe_data(db_recipe_data, obj_dict, grad_dict, comp_dict):
         if grad_dict:
             recipe_data["gradients"] = [{**v} for v in grad_dict.values()]
         return recipe_data
+
+
+class DBMaintenance(object):
+    """
+    Handles the maintenance of the database.
+    """
+
+    def __init__(self, db_handler):
+        self.db = db_handler
+        self.result_doc = ResultDoc(self.db)
+
+    def cleanup_results(self):
+        """
+        Check if the results in the database are expired and delete them if the linked object expired.
+        """
+        self.result_doc.check_expired_results()
diff --git a/cellpack/autopack/FirebaseHandler.py b/cellpack/autopack/FirebaseHandler.py
index 69e1f0fe..693332d2 100644
--- a/cellpack/autopack/FirebaseHandler.py
+++ b/cellpack/autopack/FirebaseHandler.py
@@ -166,6 +166,13 @@ def update_or_create(self, collection, id, data):
         except NotFound:
             self.set_doc(collection, id, data)
 
+    # Delete methods
+    def delete_doc(self, collection, id):
+        doc_ref = self.db.collection(collection).document(id)
+        doc_ref.delete()
+        print(f"successfully deleted path: {doc_ref.path}")
+        return doc_ref.id
+
     # other utils
     @staticmethod
     def write_creds_path():