Merge pull request #959 from synfinatic/merge-v2

Fix --threads cli arg
synfinatic · Jul 10, 2024 · e542e22 · e542e22
2 parents 672c655 + 51fb4ee
commit e542e22
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Bugs
+
+ * No longer ignore the `--threads` CLI option
+
 ### New Features
 
  * Add support for HTTP Auth/`$AWS_CONTAINER_AUTHORIZATION_TOKEN` env variable #516

diff --git a/cmd/aws-sso/auth.go b/cmd/aws-sso/auth.go
@@ -117,7 +117,7 @@ func doAuth(ctx *RunContext) *sso.AWSSSO {
 		if err != nil {
 			log.Fatalf(err.Error())
 		}
-		if err = ctx.Settings.Cache.Refresh(AwsSSO, s, ssoName); err != nil {
+		if err = ctx.Settings.Cache.Refresh(AwsSSO, s, ssoName, ctx.Cli.Threads); err != nil {
 			log.WithError(err).Fatalf("Unable to refresh cache")
 		}
 		if err = ctx.Settings.Cache.Save(true); err != nil {

diff --git a/cmd/aws-sso/cache_cmd.go b/cmd/aws-sso/cache_cmd.go
@@ -36,7 +36,7 @@ func (cc *CacheCmd) Run(ctx *RunContext) error {
 		log.Fatalf(err.Error())
 	}
 
-	err = ctx.Settings.Cache.Refresh(awssso, s, ssoName)
+	err = ctx.Settings.Cache.Refresh(awssso, s, ssoName, ctx.Cli.Threads)
 	if err != nil {
 		return fmt.Errorf("unable to refresh role cache: %s", err.Error())
 	}

diff --git a/cmd/aws-sso/tags_cmd.go b/cmd/aws-sso/tags_cmd.go
@@ -47,7 +47,7 @@ func (cc *TagsCmd) Run(ctx *RunContext) error {
 			log.Fatalf(err.Error())
 		}
 
-		err = set.Cache.Refresh(awssso, s, ssoName)
+		err = set.Cache.Refresh(awssso, s, ssoName, ctx.Cli.Threads)
 		if err != nil {
 			log.WithError(err).Fatalf("Unable to refresh role cache")
 		}

diff --git a/sso/cache.go b/sso/cache.go
@@ -278,7 +278,7 @@ func (c *Cache) deleteOldHistory() {
 
 // Refresh updates our cached Roles based on AWS SSO & our Config
 // but does not save this data!
-func (c *Cache) Refresh(sso *AWSSSO, config *SSOConfig, ssoName string) error {
+func (c *Cache) Refresh(sso *AWSSSO, config *SSOConfig, ssoName string, threads int) error {
 	// Only refresh once per execution
 	if c.refreshed {
 		return nil
@@ -314,7 +314,7 @@ func (c *Cache) Refresh(sso *AWSSSO, config *SSOConfig, ssoName string) error {
 	c.SSO[ssoName].ConfigHash = config.GetConfigHash(c.settings.ProfileFormat)
 
 	// load our AWSSSO & Config
-	r, err := c.NewRoles(sso, config)
+	r, err := c.NewRoles(sso, config, threads)
 	if err != nil {
 		return err
 	}
@@ -426,7 +426,7 @@ func (c *Cache) GetRole(arn string) (*AWSRoleFlat, error) {
 
 // Merges the AWS SSO and our Config file to create our Roles struct
 // which is defined in cache_roles.go
-func (c *Cache) NewRoles(as *AWSSSO, config *SSOConfig) (*Roles, error) {
+func (c *Cache) NewRoles(as *AWSSSO, config *SSOConfig, threads int) (*Roles, error) {
 	r := Roles{
 		SSORegion:     config.SSORegion,
 		StartUrl:      config.StartUrl,
@@ -435,7 +435,7 @@ func (c *Cache) NewRoles(as *AWSSSO, config *SSOConfig) (*Roles, error) {
 		ssoName:       config.settings.DefaultSSO,
 	}
 
-	if err := c.addSSORoles(&r, as); err != nil {
+	if err := c.addSSORoles(&r, as, threads); err != nil {
 		return &Roles{}, err
 	}
 
@@ -507,7 +507,7 @@ func processSSORoles(roles []RoleInfo, cache *SSOCache, r *Roles) {
 }
 
 // addSSORoles retrieves all the SSO Roles from AWS SSO and places them in r
-func (c *Cache) addSSORoles(r *Roles, as *AWSSSO) error {
+func (c *Cache) addSSORoles(r *Roles, as *AWSSSO, threads int) error {
 	cache := c.GetSSO()
 
 	accounts, err := as.GetAccounts()
@@ -531,9 +531,9 @@ func (c *Cache) addSSORoles(r *Roles, as *AWSSSO) error {
 	// Per #448, doing this serially is too slow for many accounts.  Hence,
 	// we'll use a worker pool.
 	if len(accounts) > 0 {
-		workers := 1
-		if c.settings.Threads > 0 {
-			workers = c.settings.Threads
+		workers := c.settings.Threads
+		if threads > 0 {
+			workers = threads
 		}
 		if workers > len(accounts) {
 			workers = len(accounts)