From 01a5fc8a4d68019fdeccf2aee4e85f735b8dc962 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Oswaldo=20Alcal=C3=A1?=
 <40160237+odarotto@users.noreply.github.com>
Date: Wed, 11 Sep 2024 12:05:27 -0400
Subject: [PATCH] feat(buckets.py): added limit option to remove command (#121)

---
 workflow/cli/buckets.py  | 8 +++++++-
 workflow/http/buckets.py | 4 +++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/workflow/cli/buckets.py b/workflow/cli/buckets.py
index 34b418f..f041b02 100644
--- a/workflow/cli/buckets.py
+++ b/workflow/cli/buckets.py
@@ -38,6 +38,9 @@ def buckets():
 )
 @click.option("-p", "--parent", type=str, required=False, help="Filter by parent.")
 @click.option("-f", "--force", is_flag=True, help="Do not prompt for confirmation")
+@click.option(
+    "-l", "--limit", type=int, required=False, help="Limit of Work objects to remove."
+)
 def remove(
     bucket: str,
     status: Optional[str] = None,
@@ -45,6 +48,7 @@ def remove(
     tag: Optional[Tuple[str]] = None,
     parent: Optional[str] = None,
     force: bool = False,
+    limit: Optional[int] = 100,
 ):
     """Remove work[s] from the buckets.
 
@@ -55,8 +59,9 @@ def remove(
         tag (Optional[Tuple[str]], optional): Filter by tag. Defaults to None.
         parent (Optional[str], optional): Filter by parent. Defaults to None.
         force (bool, optional): Do not prompt for confirmation. Defaults to False.
+        limit (int, optional): Limit of Work objects to remove. Defaults to False.
     """
-    http = HTTPContext()
+    http = HTTPContext(backends=["buckets"])
     events: Optional[List[int]] = None
     tags: Optional[List[str]] = None
     if event:
@@ -70,6 +75,7 @@ def remove(
         tags=tags,
         parent=parent,
         force=force,
+        limit=limit,
     )
 
 
diff --git a/workflow/http/buckets.py b/workflow/http/buckets.py
index 757f5ce..6b36a75 100644
--- a/workflow/http/buckets.py
+++ b/workflow/http/buckets.py
@@ -148,6 +148,7 @@ def delete_many(
         tags: Optional[List[str]] = None,
         parent: Optional[str] = None,
         force: bool = False,
+        limit: Optional[int] = 100
     ) -> bool:
         """Delete works belonging to a pipeline from the buckets backend.
 
@@ -161,6 +162,7 @@ def delete_many(
             event (Optional[List[int]]): The event to delete works with.
             force (bool, optional): Whether to force the deletion without requiring
                 user confirmation. Defaults to False.
+            limit (int, optional): Limit of Work objects to remove. Defaults to False.
 
         Returns:
             bool: Whether any works were deleted.
@@ -171,7 +173,7 @@ def delete_many(
         query.update({"tags": {"$in": tags}} if tags else {})
         query.update({"config.parent": parent} if parent else {})
         projection = {"id": True}
-        result = self.view(query, projection)
+        result = self.view(query, projection, limit=limit)
         ids: List[str] = []
         if result:
             ids = [work["id"] for work in result]