From f72b947c24e73f0a6f72173ab87db1dcd8a27ad0 Mon Sep 17 00:00:00 2001
From: Thierry Lacour <tlacour@verifa.io>
Date: Fri, 10 Nov 2023 14:21:19 +0100
Subject: [PATCH] Add optional threads arg to expand/run commands

---
 README.md               |  6 +++++-
 commands/artifactory.go |  4 ++--
 commands/expand.go      | 17 +++++++++++++++--
 commands/run.go         | 14 +++++++++++++-
 4 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 54f4cb9..f4ef06e 100644
--- a/README.md
+++ b/README.md
@@ -35,7 +35,8 @@ To enforce them, set up a humble cron job running the plugin.
 
 - Options:
   - --dry-run    _do not delete artifacts [Default: **true**]_
-  - --recursive    _recursively find FileSpecs files in the given dir [Default: false]_
+  - --recursive  _recursively find FileSpecs files in the given dir [Default: false]_
+  - --threads    _amount of worker threads [Default: 3]_
 
 #### expand
 
@@ -45,6 +46,9 @@ To enforce them, set up a humble cron job running the plugin.
   - config-path    _(Path to the JSON config file)_
   - output-path    _(Path to output the generated FileSpecs)_
 
+- Options:
+  - --threads    _amount of worker threads [Default: 3]_
+
 ### running with verbose output
 
 For verbose output to aid in debugging, set `JFROG_CLI_LOG_LEVEL=DEBUG`.
diff --git a/commands/artifactory.go b/commands/artifactory.go
index 58302a1..a1fae53 100644
--- a/commands/artifactory.go
+++ b/commands/artifactory.go
@@ -41,14 +41,14 @@ func GetArtifactoryDetails(context *components.Context) (*config.ServerDetails,
 	return details, nil
 }
 
-func GetArtifactoryManager(context *components.Context, dryRun bool) (artifactory.ArtifactoryServicesManager, error) {
+func GetArtifactoryManager(context *components.Context, dryRun bool, threads int) (artifactory.ArtifactoryServicesManager, error) {
 	artifactoryDetails, cfgErr := GetArtifactoryDetails(context)
 	if cfgErr != nil {
 		return nil, cfgErr
 	}
 
 	log.Info("Configuring Artifactory manager")
-	artifactoryManager, rtfErr := core_utils.CreateServiceManager(artifactoryDetails, 3, 5000, dryRun)
+	artifactoryManager, rtfErr := core_utils.CreateServiceManagerWithThreads(artifactoryDetails, dryRun, threads, 3, 5000)
 	if rtfErr != nil {
 		return nil, rtfErr
 	}
diff --git a/commands/expand.go b/commands/expand.go
index 4fc889b..9531494 100644
--- a/commands/expand.go
+++ b/commands/expand.go
@@ -18,6 +18,7 @@ import (
 type ExpandConfiguration struct {
 	configFile string
 	outputPath string
+	threads    int
 }
 
 type Policy struct {
@@ -90,7 +91,14 @@ func GetExpandArguments() []components.Argument {
 }
 
 func GetExpandFlags() []components.Flag {
-	return []components.Flag{}
+	return []components.Flag{
+		components.StringFlag{
+			Name:         "threads",
+			Description:  "Number of worker threads",
+			DefaultValue: "3",
+			Mandatory:    false,
+		},
+	}
 }
 
 func GetExpandEnvVar() []components.EnvVar {
@@ -199,7 +207,7 @@ func ExpandCmd(context *components.Context) error {
 					return parseErr
 				}
 
-				artifactoryManager, rtfErr := GetArtifactoryManager(context, true)
+				artifactoryManager, rtfErr := GetArtifactoryManager(context, true, expandConfig.threads)
 				if rtfErr != nil {
 					return rtfErr
 				}
@@ -295,6 +303,11 @@ func ParseExpandConfig(context *components.Context) (*ExpandConfiguration, error
 	var expandConfig = new(ExpandConfiguration)
 	expandConfig.configFile = context.Arguments[0]
 	expandConfig.outputPath = context.Arguments[1]
+	threads, err := strconv.Atoi(context.GetStringFlagValue("threads"))
+	if err != nil {
+		return nil, err
+	}
+	expandConfig.threads = threads
 
 	return expandConfig, nil
 }
diff --git a/commands/run.go b/commands/run.go
index 32aac40..f87c101 100644
--- a/commands/run.go
+++ b/commands/run.go
@@ -13,6 +13,7 @@ type RunConfiguration struct {
 	fileSpecsPath string
 	dryRun        bool
 	recursive     bool
+	threads       int
 }
 
 func GetRunCommand() components.Command {
@@ -50,6 +51,12 @@ func GetRunFlags() []components.Flag {
 			Description:  "recursively find filespecs files in the given dir",
 			DefaultValue: false,
 		},
+		components.StringFlag{
+			Name:         "threads",
+			Description:  "Number of worker threads",
+			DefaultValue: "3",
+			Mandatory:    false,
+		},
 	}
 }
 
@@ -69,7 +76,7 @@ func RunCmd(context *components.Context) error {
 	log.Debug("    recursive:", runConfig.recursive)
 
 	log.Info("Configuring Artifactory manager")
-	artifactoryManager, rtfErr := GetArtifactoryManager(context, runConfig.dryRun)
+	artifactoryManager, rtfErr := GetArtifactoryManager(context, runConfig.dryRun, runConfig.threads)
 	if rtfErr != nil {
 		return rtfErr
 	}
@@ -107,6 +114,11 @@ func ParseRunConfig(context *components.Context) (*RunConfiguration, error) {
 	runConfig.fileSpecsPath = context.Arguments[0]
 	runConfig.dryRun = context.GetBoolFlagValue("dry-run")
 	runConfig.recursive = context.GetBoolFlagValue("recursive")
+	threads, err := strconv.Atoi(context.GetStringFlagValue("threads"))
+	if err != nil {
+		return nil, err
+	}
+	runConfig.threads = threads
 
 	return runConfig, nil
 }