Delete objects versions and incomplete multipart uploads

gen3workflow/aws/bucket.py

@@ -1,5 +1,6 @@
 import json
 import random
+from cachelib import SimpleCache
 from typing import Tuple, Union
 
 import asyncio
@@ -453,78 +454,44 @@ async def create_user_bucket(user_id: str) -> Tuple[str, str, str]:
             await asyncio.sleep(delay)
 
 
-def get_all_bucket_objects(user_bucket_name: str) -> list:
-    """
-    Get all objects from the specified S3 bucket.
-    """
-    response = clients.s3_client.list_objects_v2(Bucket=user_bucket_name)
-    object_list = response.get("Contents", [])
-
-    # list_objects_v2 can utmost return 1000 objects in a single response
-    # if there are more objects, the response will have a key "IsTruncated" set to True
-    # and a key "NextContinuationToken" which can be used to get the next set of objects
-
-    # TODO:
-    # Currently, all objects are loaded into memory, which can be problematic for large buckets.
-    # To optimize, convert this function into a generator that accepts a `batch_size` parameter
-    # (capped at 1,000) and yields objects in batches.
-    # This is fine for now because this code is only called during integration tests, with a small
-    # number of files in the bucket.
-    while response.get("IsTruncated"):
-        response = clients.s3_client.list_objects_v2(
-            Bucket=user_bucket_name,
-            ContinuationToken=response.get("NextContinuationToken"),
-        )
-        object_list += response.get("Contents", [])
-
-    return object_list
-
-
-# TODO: Handle `delete_all_bucket_objects` when bucket versioning is enabled. Needed for CI, when S3Files is enabled.
-def delete_all_bucket_objects(user_id: str, user_bucket_name: str) -> None:
+def _delete_all_bucket_objects(user_id: str, user_bucket_name: str) -> None:
     """
     Deletes all objects from the specified S3 bucket.
 
     Args:
         user_id (str): The user's unique Gen3 ID.
         user_bucket_name (str): The name of the S3 bucket.
     """
-    object_list = get_all_bucket_objects(user_bucket_name)
-
-    if not object_list:
-        return
-
     logger.debug(
         f"Deleting all contents from '{user_bucket_name}' for user '{user_id}' before deleting the bucket"
     )
-    keys = [{"Key": obj.get("Key")} for obj in object_list]
-
-    # According to the docs, up to 1000 objects can be deleted in a single request:
-    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.delete_objects
-
-    # TODO: When `get_all_bucket_objects` is converted to a generator,
-    # we can remove this batching logic and retrieve objects in batches of 1,000 for deletion.
-    limit = 1000
-    for offset in range(0, len(keys), limit):
-        response = clients.s3_client.delete_objects(
-            Bucket=user_bucket_name,
-            Delete={"Objects": keys[offset : offset + limit]},
-        )
+    bucket = clients.s3_resource.Bucket(user_bucket_name)
+
+    # Cancel incomplete multipart uploads to avoid storage charges for orphaned parts
+    for upload in bucket.multipart_uploads.all():
+        upload.abort()
+
+    for response in bucket.object_versions.delete():
+        # boto returns one response for each underlying batch
         if response.get("Errors"):
-            logger.error(
-                f"Failed to delete objects from bucket '{user_bucket_name}' for user '{user_id}': {response}"
+            raise Exception(
+                f"Unable to delete bucket object versions: {response['Errors']}"
+            )
+
+    for response in bucket.objects.delete():
+        if response.get("Errors"):
+            raise Exception(
+                f"Unable to delete bucket object versions: {response['Errors']}"
             )
-            raise Exception(response)
 
 
-def cleanup_user_bucket(user_id: str, delete_bucket: bool = False) -> Union[str, None]:
+def cleanup_user_bucket(user_id: str, delete_bucket: bool) -> Union[str, None]:
     """
     Empty a user's S3 bucket and optionally delete the bucket.
 
     Args:
         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
@@ -545,7 +512,7 @@ def cleanup_user_bucket(user_id: str, delete_bucket: bool = False) -> Union[str,
             )
             return None
     try:
-        delete_all_bucket_objects(user_id, user_bucket_name)
+        _delete_all_bucket_objects(user_id, user_bucket_name)
         if delete_bucket:
             logger.info(
                 f"Initializing delete for bucket '{user_bucket_name}' for user '{user_id}'"

gen3workflow/aws/clients.py

@@ -2,7 +2,7 @@
 from gen3workflow.config import config
 
 
-def get_boto3_client(service_name: str, **kwargs):
+def get_boto3_client(service_name: str, type: str = "client", **kwargs):
     """
     Create a boto3 client for the specified AWS service,
     using credentials from the config if provided,
@@ -16,11 +16,17 @@ def get_boto3_client(service_name: str, **kwargs):
             kwargs["aws_secret_access_key"] = config[
                 "S3_ENDPOINTS_AWS_SECRET_ACCESS_KEY"
             ]
-    return boto3.client(service_name, **kwargs)
+    if type == "client":
+        return boto3.client(service_name, **kwargs)
+    else:
+        return boto3.resource(service_name, **kwargs)
 
 
 iam_client = get_boto3_client("iam")
 s3_client = get_boto3_client("s3", region_name=config["USER_BUCKETS_REGION"])
+s3_resource = get_boto3_client(
+    "s3", type="resource", region_name=config["USER_BUCKETS_REGION"]
+)
 kms_client = get_boto3_client("kms", region_name=config["USER_BUCKETS_REGION"])
 sts_client = get_boto3_client("sts")
 eks_client = get_boto3_client("eks", region_name=config["EKS_CLUSTER_REGION"])

gen3workflow/aws/s3_files.py

@@ -253,12 +253,8 @@ def _get_eks_security_groups() -> List[tuple[str, str]]:
     Return a list of tuples consisting of (sg_id, sg_name) for the EKS security group that Karpenter attaches
     to an EKS worker node in this cluster.
     """
-    eks_sg_names = config.get("EKS_SECURITY_GROUP_NAMES")
-    if not eks_sg_names:
-        eks_sg_names = [f"{config["EKS_CLUSTER_NAME"]}_EKS_workers_sg"]
-
     security_groups = clients.ec2_client.describe_security_groups(
-        Filters=[{"Name": "group-name", "Values": eks_sg_names}]
+        Filters=[{"Name": "group-name", "Values": config["EKS_SECURITY_GROUP_NAMES"]}]
     )["SecurityGroups"]
     return [(group["GroupId"], group["GroupName"]) for group in security_groups]
 

gen3workflow/config-default.yaml

@@ -101,6 +101,5 @@ WORKER_PODS_NAMESPACE: default-namespace
 
 # Names of the EKS security groups that Karpenter attaches to an EKS worker node in this cluster
 # Usually going to be in the lines of f"{config["EKS_CLUSTER_NAME"]}_EKS_workers_sg"
-# Steps to fetch the values are mentioned in the (docs/s3Files.md#determining-the-eks-security-group-names-for-s3-files-setup)
-
+# Steps to fetch the values: https://github.com/uc-cdis/gen3-workflow/blob/master/docs/s3Files.md#determining-the-eks-security-group-names-for-s3-files-setup
 EKS_SECURITY_GROUP_NAMES: []

gen3workflow/config.py

@@ -72,6 +72,11 @@ def validate_top_level_configs(self) -> None:
             self["S3_ENDPOINTS_AWS_SECRET_ACCESS_KEY"]
         ), "Both 'S3_ENDPOINTS_AWS_ACCESS_KEY_ID' and 'S3_ENDPOINTS_AWS_SECRET_ACCESS_KEY' must be configured, or both must be left empty"
 
+        if self["ENABLE_S3_FILES"]:
+            assert (
+                len(self["EKS_SECURITY_GROUP_NAMES"]) > 0
+            ), "EKS_SECURITY_GROUP_NAMES must be configured when ENABLE_S3_FILES is True"
+
 
 config = Gen3WorkflowConfig(DEFAULT_CFG_PATH)
 try:

gen3workflow/routes/storage.py

@@ -11,10 +11,7 @@
 from gen3workflow import logger
 from gen3workflow.auth import Auth
 from gen3workflow.aws import s3_files
-from gen3workflow.aws.bucket import (
-    cleanup_user_bucket,
-    create_user_bucket,
-)
+from gen3workflow.aws.bucket import cleanup_user_bucket, create_user_bucket
 from gen3workflow.config import config
 
 router = APIRouter(prefix="/storage")
@@ -133,7 +130,7 @@ async def empty_user_bucket(request: Request, auth=Depends(Auth)) -> None:
         "delete", [f"/services/workflow/gen3-workflow/storage/{user_id}"]
     )
     logger.info(f"User '{user_id}' emptying their storage bucket")
-    deleted_bucket_name = cleanup_user_bucket(user_id)
+    deleted_bucket_name = cleanup_user_bucket(user_id, delete_bucket=False)
 
     if not deleted_bucket_name:
         raise HTTPException(