From e2c4a0de3ac0432ad308a61c855a5100594627ff Mon Sep 17 00:00:00 2001
From: Roman Roberman <roman@appsflyer.com>
Date: Tue, 2 Apr 2024 12:52:43 +0300
Subject: [PATCH] add --worker-count (#132)

* add threads flag to use threads and add missing readme for ignore-not-found

* add test context manager inside thread

* review

* review
---
 README.md                                 |  4 ++++
 artifactory_cleanup/artifactorycleanup.py |  9 ++++++++-
 artifactory_cleanup/cli.py                | 10 ++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 38d2319..acce4b6 100644
--- a/README.md
+++ b/README.md
@@ -113,6 +113,10 @@ artifactory-cleanup --help
 - Use CI servers or cron-like utilities to run `artifactory-cleanup` every day (or every hour). TeamCity and GitHub have
   built-in support and show additional logs format
 - Do not save credentials in the configuration file, use environment variables.
+- Use `--ignore-not-found` flag to ignore errors when the repository is not found. It's useful when you have a
+  configuration for multiple repositories and some of them are not found.
+- Use `--worker-count=<WORKER_NUM>` to increase the number of workers. By default, it's 1. It's useful when you have a lot of
+  artifacts and you want to speed up the process.
 
 ## Commands ##
 
diff --git a/artifactory_cleanup/artifactorycleanup.py b/artifactory_cleanup/artifactorycleanup.py
index 68615fd..af447b7 100644
--- a/artifactory_cleanup/artifactorycleanup.py
+++ b/artifactory_cleanup/artifactorycleanup.py
@@ -1,3 +1,4 @@
+from concurrent.futures import ThreadPoolExecutor
 from datetime import date
 from typing import List, Iterator
 
@@ -23,11 +24,13 @@ def __init__(
         destroy: bool,
         today: date,
         ignore_not_found: bool,
+        worker_count: int,
     ):
         self.session = session
         self.policies = policies
         self.destroy = destroy
         self.ignore_not_found = ignore_not_found
+        self.worker_count = worker_count
 
         self._init_policies(today)
 
@@ -56,10 +59,14 @@ def cleanup(self, block_ctx_mgr, test_ctx_mgr) -> Iterator[CleanupSummary]:
                 print(f"Found {len(artifacts_to_remove)} artifacts AFTER filtering")
 
                 # Delete artifacts
-                for artifact in artifacts_to_remove:
+                def _delete(artifact):
                     with test_ctx_mgr(get_name_for_ci(artifact)):
                         policy.delete(artifact, destroy=self.destroy, ignore_not_found=self.ignore_not_found)
 
+                with ThreadPoolExecutor(max_workers=int(self.worker_count)) as executor:
+                    for artifact in artifacts_to_remove:
+                        executor.submit(_delete, artifact=artifact)
+
             # Show summary
             print(f"Deleted artifacts count: {len(artifacts_to_remove)}")
             try:
diff --git a/artifactory_cleanup/cli.py b/artifactory_cleanup/cli.py
index ac77fd2..5ad6f3a 100644
--- a/artifactory_cleanup/cli.py
+++ b/artifactory_cleanup/cli.py
@@ -61,6 +61,14 @@ class ArtifactoryCleanupCLI(cli.Application):
         envname="ARTIFACTORY_CLEANUP_IGNORE_NOT_FOUND",
     )
 
+    _worker_count = cli.SwitchAttr(
+        "--worker-count",
+        help="Number of workers to use",
+        mandatory=False,
+        default=1,
+        envname="ARTIFACTORY_CLEANUP_WORKER_COUNT",
+    )
+
     _days_in_future = cli.SwitchAttr(
         "--days-in-future",
         help="Simulate future behaviour",
@@ -156,12 +164,14 @@ def main(self):
         session.auth = HTTPBasicAuth(user, password)
 
         self._destroy_or_verbose()
+        print(f"Using {self._worker_count} workers")
         cleanup = ArtifactoryCleanup(
             session=session,
             policies=policies,
             destroy=self._destroy,
             today=today,
             ignore_not_found=self._ignore_not_found,
+            worker_count=self._worker_count,
         )
 
         # Filter policies by name