Adding a cleanup entrypoint to allow for deletion of all (or a subset of) reports older than a given number of days.

Author: jonbannister

notebooker/_entrypoints.py

@@ -1,5 +1,6 @@
 import os
 import uuid
+from typing import Optional
 
 import click
 
@@ -10,6 +11,7 @@
 from notebooker.serialization import SERIALIZER_TO_CLI_OPTIONS
 from notebooker.settings import BaseConfig, WebappConfig
 from notebooker.snapshot import snap_latest_successful_notebooks
+from notebooker.utils.cleanup import delete_old_reports
 from notebooker.web.app import main
 
 
@@ -267,6 +269,17 @@ def execute_notebook(
     )
 
 
+@base_notebooker.command()
+@click.option("--days", "--days-cutoff", "-d", type=int, required=True, help="Delete reports older than this many days")
+@click.option(
+    "--report-name", required=False, help="The name of the template to retrieve, relative to the template directory."
+)
+@click.option("--dry-run", is_flag=True, default=False, help="Show what would be deleted without actually deleting")
+@pass_config
+def cleanup_old_reports(config: BaseConfig, days: int, report_name: Optional[str], dry_run: bool):
+    delete_old_reports(config, days_cutoff=days, report_name=report_name, dry_run=dry_run)
+
+
 @base_notebooker.command()
 @click.option(
     "--report-name", required=True, help="The name of the template to retrieve, relative to the template directory."

notebooker/serialization/mongo.py

@@ -521,18 +521,33 @@ def get_latest_successful_job_ids_for_name_all_params(self, report_name: str) ->
     def n_all_results_for_report_name(self, report_name: str) -> int:
         return self._get_result_count({"report_name": report_name})
 
-    def delete_result(self, job_id: AnyStr) -> Dict[str, Any]:
+    def delete_result(self, job_id: AnyStr, dry_run: bool = False) -> Dict[str, Any]:
         result = self._get_raw_check_result(job_id)
         status = JobStatus.from_string(result["status"])
         gridfs_filenames = load_files_from_gridfs(self.result_data_store, result, do_read=False)
         if status in (JobStatus.ERROR, JobStatus.TIMEOUT, JobStatus.CANCELLED):
             gridfs_filenames.append(_error_info_filename(job_id))
-        self.update_check_status(job_id, JobStatus.DELETED)
+        if not dry_run:
+            self.update_check_status(job_id, JobStatus.DELETED)
+        deleted_gridfs_files = []
         for filename in gridfs_filenames:
-            logger.info(f"Deleting {filename}")
+            logger.debug(f"Deleting {filename}")
+            existed = False
             for grid_out in self.result_data_store.find({"filename": filename}):
-                self.result_data_store.delete(grid_out._id)
-        return {"deleted_result_document": result, "gridfs_filenames": gridfs_filenames}
+                existed = True
+                if not dry_run:
+                    self.result_data_store.delete(grid_out._id)
+            if existed:
+                deleted_gridfs_files.append(filename)
+        return {"deleted_result_document": result, "gridfs_filenames": deleted_gridfs_files}
+
+    def get_job_ids_older_than(self, cutoff: datetime.datetime, report_name: Optional[str] = None) -> List[str]:
+        query = {"job_start_time": {"$lte": cutoff}}
+        query = _add_deleted_status_to_filter(query)
+        if report_name:
+            query["report_name"] = report_name
+        to_delete = [d["job_id"] for d in self.library.find(query, {"_id": 0, "job_id": 1})]
+        return to_delete
 
 
 def _pdf_filename(job_id: str) -> str:

notebooker/utils/cleanup.py

@@ -0,0 +1,50 @@
+import datetime
+from typing import Optional
+
+from tqdm import tqdm
+import logging
+
+from notebooker.serialization.serialization import get_serializer_from_cls
+from notebooker.settings import BaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+def delete_old_reports(config: BaseConfig, days_cutoff: int, report_name: Optional[str], dry_run: bool = True) -> None:
+    """
+    Delete notebooker reports older than specified days.
+
+    Args:
+        config: The configuration which will point to the serializer class and config.
+        days_cutoff: Delete reports older than this many days
+        report_name: Optionally specify which report_name we should be removing old reports for.
+        dry_run: If True, only show what would be deleted without actually deleting
+    """
+    serializer = get_serializer_from_cls(config.SERIALIZER_CLS, **config.SERIALIZER_CONFIG)
+    cutoff_date = datetime.datetime.now() - datetime.timedelta(days=days_cutoff)
+
+    # Find reports to delete
+    to_delete = serializer.get_job_ids_older_than(cutoff_date, report_name=report_name)
+
+    num_reports = len(to_delete)
+
+    if num_reports == 0:
+        logger.info(f"No reports found older than {days_cutoff} days")
+        return
+
+    logger.info(f"Found {num_reports} reports older than {days_cutoff} days")
+
+    # Delete reports
+    logger.info("Starting deletion process...")
+    for report in tqdm(to_delete, desc="Deleting reports"):
+        try:
+            removed = serializer.delete_result(report, dry_run=dry_run)
+            logger.info(
+                f"{'Would have deleted' if dry_run else 'Deleted'}: "
+                f"Title={removed['deleted_result_document']['report_title']}, "
+                f"GridFS files={removed['gridfs_filenames']}"
+            )
+        except Exception as e:
+            logger.error(f"Failed to delete report {report}: {str(e)}")
+
+    logger.info(f"{'Would have' if dry_run else 'Successfully'} removed {num_reports} reports")

tests/integration/test_mongo.py

@@ -77,6 +77,7 @@ def test_delete(bson_library, webapp_config):
             raw_html_resources={"inlining": {"big_thing": "a" * 32 * (2**20)}, "other_stuff": "Yep"},
         )
     )
+    assert serializer.get_job_ids_older_than(datetime.datetime(2020, 1, 1), report_name=report_name) == [job_id]
     assert bson_library.find_one({"job_id": job_id}) is not None
     result = serializer.get_check_result(job_id)
     assert result is not None

tests/unit/serialization/test_mongoose.py renamed to tests/unit/serialization/test_mongo.py

@@ -1,3 +1,5 @@
+from unittest.mock import Mock, MagicMock
+
 from mock import patch
 
 from notebooker.serialization.mongo import JobStatus, MongoResultSerializer
@@ -42,3 +44,33 @@ def test__get_all_job_ids(conn, db, gridfs):
             {"$project": {"report_name": 1, "job_id": 1}},
         ]
     )
+
+
+@patch("notebooker.serialization.mongo.gridfs")
+@patch("notebooker.serialization.mongo.MongoResultSerializer.get_mongo_database")
+@patch("notebooker.serialization.mongo.MongoResultSerializer.get_mongo_connection")
+def test_delete_result_dry_run(mock_conn, mock_db, mock_gridfs):
+    # Setup
+    serializer = MongoResultSerializer()
+    mock_result = {
+        "job_id": "test_job",
+        "status": JobStatus.DONE.value,
+        "raw_html_resources": {"outputs": ["file1.html"]},
+        "generate_pdf_output": True,
+    }
+
+    serializer._get_raw_check_result = Mock(return_value=mock_result)
+    mock_gridfs_instance = MagicMock()
+    serializer.result_data_store = mock_gridfs_instance
+    mock_gridfs_instance.find.return_value = [Mock(_id="id1")]
+
+    # Execute with dry_run=True
+    result = serializer.delete_result("test_job", dry_run=True)
+
+    # Verify no actual deletions occurred
+    assert not serializer.library.find_one_and_update.called
+    assert not mock_gridfs_instance.delete.called
+
+    # But verify the result contains what would be deleted
+    assert result["deleted_result_document"] == mock_result
+    assert len(result["gridfs_filenames"]) > 0