Simplify implementation and reduce code duplication

Author: nss10

gen3workflow/aws_utils.py

@@ -409,26 +409,6 @@ def create_user_bucket(user_id: str) -> Tuple[str, str, str]:
     return user_bucket_name, "ga4gh-tes", config["USER_BUCKETS_REGION"], kms_key_arn
 
 
-def _get_user_bucket_if_present(user_id: str) -> Optional[str]:
-    """
-    Return the user's bucket name if it exists in S3, else None.
-
-    Derives the bucket name from the user_id and checks its existence
-    using a HeadBucket call. Raises for any AWS errors other than
-    bucket-not-found.
-    """
-    bucket = get_bucket_name_from_user_id(user_id)
-    try:
-        s3_client.head_bucket(Bucket=bucket)
-        return bucket
-    except ClientError as e:
-        code = e.response.get("Error", {}).get("Code")
-        if code == "404":
-            logger.warning(f"Bucket '{bucket}' not found for user '{user_id}'.")
-            return None
-        raise
-
-
 def get_all_bucket_objects(user_bucket_name: str) -> list:
     """
     Get all objects from the specified S3 bucket.
@@ -492,41 +472,44 @@ def delete_all_bucket_objects(user_id: str, user_bucket_name: str) -> None:
             raise Exception(response)
 
 
-def empty_user_bucket(user_id: str) -> Optional[str]:
-    bucket = _get_user_bucket_if_present(user_id)
-    if not bucket:
-        return None
-    try:
-        delete_all_bucket_objects(user_id, bucket)
-    except Exception as e:
-        logger.error(
-            f"Failed to empty the bucket: '{bucket}' for user '{user_id}': {e}"
-        )
-        raise
-    return bucket
-
-
-def delete_user_bucket(user_id: str) -> Union[str, None]:
+def cleanup_user_bucket(user_id: str, delete_bucket: bool = False) -> Union[str, None]:
     """
-    Deletes all objects from a user's S3 bucket before deleting the bucket itself.
+    Empty a user's S3 bucket and optionally delete the bucket.
 
     Args:
-        user_id (str): The user's unique Gen3 ID
+        user_id: User identifier used to derive the bucket name.
+        delete_bucket:  If True, delete the bucket after removing all objects.
+                        Defaults to False (only objects are removed).
+
+    Returns:
+        Bucket name if it exists and cleanup was performed, otherwise None
+        if the bucket does not exist.
 
     Raises:
-        Exception: If there is an error during the deletion process.
+        Exception: Propagates unexpected errors during cleanup or deletion.
     """
-    user_bucket_name = _get_user_bucket_if_present(user_id)
-    logger.info(f"Deleting bucket '{user_bucket_name}' for user '{user_id}'")
-    if not user_bucket_name:
-        return None
+    user_bucket_name = get_bucket_name_from_user_id(user_id)
+
+    try:
+        s3_client.head_bucket(Bucket=user_bucket_name)
+    except ClientError as e:
+        error_code = e.response["Error"]["Code"]
+        if error_code == "404":
+            logger.warning(
+                f"Bucket '{user_bucket_name}' not found for user '{user_id}'."
+            )
+            return None
     try:
         delete_all_bucket_objects(user_id, user_bucket_name)
-        s3_client.delete_bucket(Bucket=user_bucket_name)
+        if delete_bucket:
+            logger.info(
+                f"Initializing delete for bucket '{user_bucket_name}' for user '{user_id}'"
+            )
+            s3_client.delete_bucket(Bucket=user_bucket_name)
         return user_bucket_name
 
     except Exception as e:
         logger.error(
-            f"Failed to delete bucket '{user_bucket_name}' for user '{user_id}': {e}"
+            f"Failed to cleanup bucket '{user_bucket_name}' for user '{user_id}': {e}"
         )
         raise

gen3workflow/routes/storage.py

@@ -50,7 +50,7 @@ async def delete_user_bucket(request: Request, auth=Depends(Auth)) -> None:
     token_claims = await auth.get_token_claims()
     user_id = token_claims.get("sub")
     logger.info(f"User '{user_id}' deleting their storage bucket")
-    deleted_bucket_name = aws_utils.delete_user_bucket(user_id)
+    deleted_bucket_name = aws_utils.cleanup_user_bucket(user_id, delete_bucket=True)
 
     if not deleted_bucket_name:
         raise HTTPException(
@@ -85,7 +85,7 @@ async def empty_user_bucket(request: Request, auth=Depends(Auth)) -> None:
     token_claims = await auth.get_token_claims()
     user_id = token_claims.get("sub")
     logger.info(f"User '{user_id}' emptying their storage bucket")
-    deleted_bucket_name = aws_utils.empty_user_bucket(user_id)
+    deleted_bucket_name = aws_utils.cleanup_user_bucket(user_id)
 
     if not deleted_bucket_name:
         raise HTTPException(

tests/test_misc.py

@@ -324,7 +324,7 @@ async def test_delete_user_bucket_no_token(client, mock_aws_services):
     """
     mock_delete_bucket = MagicMock()
     # Delete the bucket
-    with patch("gen3workflow.aws_utils.delete_user_bucket", mock_delete_bucket):
+    with patch("gen3workflow.aws_utils.cleanup_user_bucket", mock_delete_bucket):
         res = await client.delete("/storage/user-bucket")
         assert res.status_code == 401, res.text
         assert res.json() == {"detail": "Must provide an access token"}
@@ -346,7 +346,7 @@ async def test_delete_user_bucket_unauthorized(
     """
     mock_delete_bucket = MagicMock()
     # Delete the bucket
-    with patch("gen3workflow.aws_utils.delete_user_bucket", mock_delete_bucket):
+    with patch("gen3workflow.aws_utils.cleanup_user_bucket", mock_delete_bucket):
         res = await client.delete(
             "/storage/user-bucket",
             headers={"Authorization": f"bearer {TEST_USER_TOKEN}"},