Changed Logic to remove old execution directories instead of just ipyb files (#1134)

Author: jayasimha-raghavan-unskript

unskript-ctl/unskript_audit_cleanup.py

@@ -11,6 +11,7 @@
 
 import os
 import yaml
+import shutil
 import ZODB
 import datetime
 import ZODB.FileStorage
@@ -19,48 +20,57 @@
 from unskript_utils import *
 
 
-def remove_old_ipynb_files():
+# We also need to remove old directories in the execution directory
+# That are older than the threshold date
+def remove_old_directories():
     # Read the Audit period from config file
     audit_period = get_audit_period()
     current_date = datetime.datetime.now()
     threshold_date = current_date - datetime.timedelta(days=audit_period)
-    files_deleted = False
-    for ipynb_file in os.listdir(UNSKRIPT_EXECUTION_DIR):
-        if ipynb_file.endswith('.ipynb'):
-            file_path = os.path.join(UNSKRIPT_EXECUTION_DIR, ipynb_file)
-            file_ts = datetime.datetime.fromtimestamp(os.path.getmtime(file_path))
-            if file_ts < threshold_date:
+    directories_deleted = False
+    for directory in os.listdir(UNSKRIPT_EXECUTION_DIR):
+        dir_path = os.path.join(UNSKRIPT_EXECUTION_DIR, directory)
+        if os.path.isdir(dir_path):
+            dir_ts = datetime.datetime.fromtimestamp(os.path.getmtime(dir_path))
+            if dir_ts < threshold_date:
                 try:
-                    os.remove(file_path)
-                    print(f"Deleted {file_path}")
-                    files_deleted = True
+                    # Use shutil.rmtree instead of os.rmdir to remove non-empty directories
+                    shutil.rmtree(dir_path)
+                    print(f"Deleted {dir_path}")
+                    directories_deleted = True
                 except Exception as e:
-                    print(f"ERROR: {e}")
-                    return 
+                    print(f"ERROR: Failed to delete {dir_path}: {e}")
+                    # Continue with other directories rather than returning
+                    continue
 
-    if files_deleted:
-        print(f"Deleted ipynb files older than {audit_period} days!")
+    if directories_deleted:
+        print(f"Deleted directories older than {audit_period} days!")
     else:
-        print(f"No ipynb files are older than {audit_period}. Nothing to delete")
+        print(f"No directories are older than {audit_period}. Nothing to delete")
     return
 
+
 def get_audit_period():
     audit_period = 90
     try:
         if os.path.exists(GLOBAL_CTL_CONFIG) is True:
-            with open(GLOBAL_CTL_CONFIG, 'r', encoding='utf-8') as f:
+            with open(GLOBAL_CTL_CONFIG, "r", encoding="utf-8") as f:
                 data = yaml.safeload(f.read())
-                if data and data.get('global') and data.get('global').get('audit_period'):
-                    audit_period = data.get('global').get('audit_period')
+                if (
+                    data
+                    and data.get("global")
+                    and data.get("global").get("audit_period")
+                ):
+                    audit_period = data.get("global").get("audit_period")
     except:
         # We use 90 days as the default period to cleanup  then.
-        pass 
+        pass
     return audit_period
 
 
-def clean_db() -> DB:
+def clean_db() -> None:
     """clean_db This function calls the db.pack(...) function to cleanup the ZoDB of data that are audit_period old
-       default is 90 days. This function can be called as docker cron job to cleanup ZoDB data that are 90 days old
+    default is 90 days. This function can be called as docker cron job to cleanup ZoDB data that are 90 days old
     """
     audit_period = get_audit_period()
 
@@ -73,6 +83,6 @@ def clean_db() -> DB:
         print(f"ERROR: {e}")
 
 
-if __name__ == '__main__':
-    remove_old_ipynb_files()
-    clean_db()
+if __name__ == "__main__":
+    remove_old_directories()
+    clean_db()