New API and a few enhancements (#77)

- Added `delete_files_and_snapshots()`
- Modify SubmitAndMonitorMultipleJobs to collect job failures and report instead of exiting early
- Added  the `dry_run` property to `TDR`

Author: pshapiro4broad

VERSION.txt

@@ -1,2 +1,4 @@
-11.6.0
-- Adding methods to get more specific workflow details in workspace
+11.7.0
+- Added delete_files_and_snapshots()
+- Modify SubmitAndMonitorMultipleJobs to collect job failures and report instead of exiting early
+- Added dry_run property to TDR

ops_utils/tdr_utils/tdr_api_utils.py

@@ -20,14 +20,16 @@ class TDR:
     DEV_LINK = "https://jade.datarepo-dev.broadinstitute.org/api/repository/v1"
     """(str): The base URL for the TDR API."""
 
-    def __init__(self, request_util: RunRequest, env: str = 'prod'):
+    def __init__(self, request_util: RunRequest, env: str = 'prod', dry_run: bool = False):
         """
         Initialize the TDR class (A class to interact with the Terra Data Repository (TDR) API).
 
         **Args:**
         - request_util (`ops_utils.request_util.RunRequest`): Utility for making HTTP requests.
         """
         self.request_util = request_util
+        # NOTE: dry_run is not fully implemented in this class, only in delete_files_and_snapshots
+        self.dry_run = dry_run
         if env.lower() == 'prod':
             self.tdr_link = self.PROD_LINK
         elif env.lower() == 'dev':
@@ -180,6 +182,47 @@ def delete_files(
             check_interval=check_interval
         ).run()
 
+    def _delete_snapshots_for_files(self, dataset_id: str, file_ids: set[str]) -> None:
+        """Delete snapshots that reference any of the provided file IDs."""
+        snapshots_resp = self.get_dataset_snapshots(dataset_id=dataset_id)
+        snapshot_items = snapshots_resp.json().get('items', [])
+        snapshots_to_delete = []
+        logging.info(
+            "Checking %d snapshots for references",
+            len(snapshot_items),
+        )
+        for snap in snapshot_items:
+            snap_id = snap.get('id')
+            if not snap_id:
+                continue
+            snap_files = self.get_files_from_snapshot(snapshot_id=snap_id)
+            snap_file_ids = {
+                fd.get('fileId') for fd in snap_files if fd.get('fileId')
+            }
+            # Use set intersection to check for any matching file IDs
+            if snap_file_ids & file_ids:
+                snapshots_to_delete.append(snap_id)
+        if snapshots_to_delete:
+            self.delete_snapshots(snapshot_ids=snapshots_to_delete)
+        else:
+            logging.info("No snapshots reference the provided file ids")
+
+    def _dry_run_msg(self) -> str:
+        return '[Dry run] ' if self.dry_run else ''
+
+    def delete_files_and_snapshots(self, dataset_id: str, file_ids: set[str]) -> None:
+        """Delete files from a dataset by their IDs, handling snapshots."""
+        self._delete_snapshots_for_files(dataset_id=dataset_id, file_ids=file_ids)
+
+        logging.info(
+            f"{self._dry_run_msg()}Submitting delete request for {len(file_ids)} files in "
+            f"dataset {dataset_id}")
+        if not self.dry_run:
+            self.delete_files(
+                file_ids=list(file_ids),
+                dataset_id=dataset_id
+            )
+
     def add_user_to_dataset(self, dataset_id: str, user: str, policy: str) -> requests.Response:
         """
         Add a user to a dataset with a specified policy.
@@ -322,14 +365,16 @@ def delete_snapshots(
         - check_interval (int, optional): The interval in seconds to wait between status checks. Defaults to `10`.
         - verbose (bool, optional): Whether to log detailed information about each job. Defaults to `False`.
         """
-        SubmitAndMonitorMultipleJobs(
-            tdr=self,
-            job_function=self.delete_snapshot,
-            job_args_list=[(snapshot_id,) for snapshot_id in snapshot_ids],
-            batch_size=batch_size,
-            check_interval=check_interval,
-            verbose=verbose
-        ).run()
+        logging.info(f"{self._dry_run_msg()}Deleting {len(snapshot_ids)} snapshots")
+        if not self.dry_run:
+            SubmitAndMonitorMultipleJobs(
+                tdr=self,
+                job_function=self.delete_snapshot,
+                job_args_list=[(snapshot_id,) for snapshot_id in snapshot_ids],
+                batch_size=batch_size,
+                check_interval=check_interval,
+                verbose=verbose
+            ).run()
 
     def delete_snapshot(self, snapshot_id: str) -> requests.Response:
         """
@@ -937,6 +982,10 @@ def _get_response_from_batched_endpoint(self, uri: str, limit: int = 1000) -> li
                 break
 
             metadata.extend(response_json)
+            if len(response_json) < limit:
+                logging.info(f"Retrieved final batch of results, found {len(metadata)} total records")
+                break
+
             # Increment the offset by limit for the next page
             offset += limit
             batch += 1

ops_utils/tdr_utils/tdr_job_utils.py

@@ -110,7 +110,10 @@ def run(self) -> None:
         Run the process to submit and monitor multiple jobs in batches.
 
         Logs the progress and status of each batch and job.
+
+        Failed jobs are collected and printed out at the end of processing.
         """
+        failed_jobs = []
         total_jobs = len(self.job_args_list)
         logging.info(f"Processing {total_jobs} {self.job_function.__name__} jobs in batches of {self.batch_size}")
 
@@ -133,13 +136,21 @@ def run(self) -> None:
             # Monitor jobs for the current batch
             logging.info(f"Monitoring {len(current_batch)} jobs in batch {i // self.batch_size + 1}")
             for job_id in job_ids:
-                MonitorTDRJob(
-                    tdr=self.tdr,
-                    job_id=job_id,
-                    check_interval=self.check_interval,
-                    return_json=False
-                ).run()
+                try:
+                    MonitorTDRJob(
+                        tdr=self.tdr,
+                        job_id=job_id,
+                        check_interval=self.check_interval,
+                        return_json=False
+                    ).run()
+                except Exception as e:
+                    logging.error(f"Job {job_id} failed: {e}")
+                    failed_jobs.append(job_id)
 
             logging.info(f"Completed batch {i // self.batch_size + 1} with {len(current_batch)} jobs.")
 
-        logging.info(f"Successfully processed {total_jobs} jobs.")
+        logging.info(f"Successfully processed {total_jobs - len(failed_jobs)} jobs.")
+
+        if len(failed_jobs) > 0:
+            raise Exception(
+                f"The following job IDs failed: {', '.join(failed_jobs)}")