From 1c412ebaceaba256e1416dfb855f8164c03a6663 Mon Sep 17 00:00:00 2001
From: Elliot Graebert <elliot.graebert@skydio.com>
Date: Thu, 30 May 2024 08:19:33 +0000
Subject: [PATCH] Added retry logic to user groups

---
 slack/resource_usergroup.go | 36 +++++++++++++++++++++++++++++++++---
 slack/retry.go              | 22 ++++++++++++++++++++++
 2 files changed, 55 insertions(+), 3 deletions(-)
 create mode 100644 slack/retry.go

diff --git a/slack/resource_usergroup.go b/slack/resource_usergroup.go
index c36d4b5..7df0cb7 100644
--- a/slack/resource_usergroup.go
+++ b/slack/resource_usergroup.go
@@ -109,11 +109,21 @@ func resourceSlackUserGroupRead(ctx context.Context, d *schema.ResourceData, m i
 	client := m.(*slack.Client)
 	id := d.Id()
 	var diags diag.Diagnostics
-	userGroups, err := client.GetUserGroupsContext(ctx, slack.GetUserGroupsOptionIncludeUsers(true))
+
+	operation := func() (interface{}, error) {
+		return client.GetUserGroupsContext(ctx, slack.GetUserGroupsOptionIncludeUsers(true))
+	}
+	result, err := retryOnRateLimit(operation)
+
 	if err != nil {
 		return diag.FromErr(fmt.Errorf("couldn't get usergroups: %w", err))
 	}
 
+	userGroups, ok := result.([]slack.UserGroup)
+	if !ok {
+		return diag.FromErr(fmt.Errorf("unexpected result type: %T", result))
+	}
+
 	for _, userGroup := range userGroups {
 		if userGroup.ID == id {
 			return updateUserGroupData(d, userGroup)
@@ -129,11 +139,21 @@ func resourceSlackUserGroupRead(ctx context.Context, d *schema.ResourceData, m i
 
 func findUserGroupByName(ctx context.Context, name string, includeDisabled bool, m interface{}) (slack.UserGroup, error) {
 	client := m.(*slack.Client)
-	userGroups, err := client.GetUserGroupsContext(ctx, slack.GetUserGroupsOptionIncludeDisabled(includeDisabled), slack.GetUserGroupsOptionIncludeUsers(true))
+
+	operation := func() (interface{}, error) {
+		return client.GetUserGroupsContext(ctx, slack.GetUserGroupsOptionIncludeDisabled(includeDisabled), slack.GetUserGroupsOptionIncludeUsers(true))
+	}
+	result, err := retryOnRateLimit(operation)
+
 	if err != nil {
 		return slack.UserGroup{}, err
 	}
 
+	userGroups, ok := result.([]slack.UserGroup)
+	if !ok {
+		return slack.UserGroup{}, fmt.Errorf("unexpected result type: %T", result)
+	}
+
 	for _, userGroup := range userGroups {
 		if userGroup.Name == name {
 			return userGroup, nil
@@ -145,11 +165,21 @@ func findUserGroupByName(ctx context.Context, name string, includeDisabled bool,
 
 func findUserGroupByID(ctx context.Context, id string, includeDisabled bool, m interface{}) (slack.UserGroup, error) {
 	client := m.(*slack.Client)
-	userGroups, err := client.GetUserGroupsContext(ctx, slack.GetUserGroupsOptionIncludeDisabled(includeDisabled), slack.GetUserGroupsOptionIncludeUsers(true))
+
+	operation := func() (interface{}, error) {
+		return client.GetUserGroupsContext(ctx, slack.GetUserGroupsOptionIncludeDisabled(includeDisabled), slack.GetUserGroupsOptionIncludeUsers(true))
+	}
+	result, err := retryOnRateLimit(operation)
+
 	if err != nil {
 		return slack.UserGroup{}, err
 	}
 
+	userGroups, ok := result.([]slack.UserGroup)
+	if !ok {
+		return slack.UserGroup{}, fmt.Errorf("unexpected result type: %T", result)
+	}
+
 	for _, userGroup := range userGroups {
 		if userGroup.ID == id {
 			return userGroup, nil
diff --git a/slack/retry.go b/slack/retry.go
new file mode 100644
index 0000000..29c8bf5
--- /dev/null
+++ b/slack/retry.go
@@ -0,0 +1,22 @@
+package slack
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/slack-go/slack"
+)
+
+func retryOnRateLimit(operation func() (interface{}, error)) (interface{}, error) {
+	for {
+		result, err := operation()
+		if rateLimitedError, ok := err.(*slack.RateLimitedError); ok {
+			fmt.Printf("Rate limited. Retrying after %v seconds...\n", rateLimitedError.RetryAfter)
+			time.Sleep(rateLimitedError.RetryAfter)
+			continue
+		} else if err != nil {
+			return nil, err
+		}
+		return result, nil
+	}
+}