From 388d32ab7ae7cb1dc56470938093844f0d5efd4b Mon Sep 17 00:00:00 2001
From: Prashansa Kulshrestha <prashkulshrestha@gmail.com>
Date: Tue, 3 Sep 2024 10:54:08 +0530
Subject: [PATCH] fix: added retry and backoff in konnect authenticate for 429
 errors

---
 cmd/common_konnect.go           | 31 +++++++++++++++++++++++++++++--
 tests/integration/test_utils.go |  5 +++--
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/cmd/common_konnect.go b/cmd/common_konnect.go
index 16dcde56e..e93fcfda0 100644
--- a/cmd/common_konnect.go
+++ b/cmd/common_konnect.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/kong/go-database-reconciler/pkg/diff"
 	"github.com/kong/go-database-reconciler/pkg/dump"
@@ -16,12 +17,38 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-const defaultControlPlaneName = "default"
+const (
+	defaultControlPlaneName         = "default"
+	maxRetriesForAuth               = 5
+	apiRateLimitExceededErrorString = "API rate limit exceeded"
+)
 
 func authenticate(
 	ctx context.Context, client *konnect.Client, konnectConfig utils.KonnectConfig,
 ) (konnect.AuthResponse, error) {
-	return client.Auth.LoginV2(ctx, konnectConfig.Email, konnectConfig.Password, konnectConfig.Token)
+	attempts := 0
+	backoff := 200 * time.Millisecond
+
+	for {
+		// fmt.Printf("Attempt #%d\n", attempts+1)
+		authResponse, err := client.Auth.LoginV2(ctx, konnectConfig.Email, konnectConfig.Password, konnectConfig.Token)
+		if err == nil {
+			return authResponse, nil
+		}
+
+		if !strings.Contains(err.Error(), apiRateLimitExceededErrorString) {
+			return authResponse, err
+		}
+
+		attempts++
+		if attempts > maxRetriesForAuth {
+			return authResponse, fmt.Errorf("maximum retries (%d) exceeded for authentication", maxRetriesForAuth)
+		}
+
+		time.Sleep(backoff)
+		backoff *= 2
+	}
+	// return client.Auth.LoginV2(ctx, konnectConfig.Email, konnectConfig.Password, konnectConfig.Token)
 }
 
 // GetKongClientForKonnectMode abstracts the different cloud environments users
diff --git a/tests/integration/test_utils.go b/tests/integration/test_utils.go
index d20600359..d1d37f676 100644
--- a/tests/integration/test_utils.go
+++ b/tests/integration/test_utils.go
@@ -43,7 +43,8 @@ func getTestClient() (*kong.Client, error) {
 		return cmd.GetKongClientForKonnectMode(ctx, &konnectConfig)
 	}
 	return utils.GetKongClient(utils.KongClientConfig{
-		Address: getKongAddress(),
+		Address:   getKongAddress(),
+		Retryable: true,
 	})
 }
 
@@ -236,7 +237,7 @@ func reset(t *testing.T, opts ...string) {
 	t.Helper()
 
 	deckCmd := cmd.NewRootCmd()
-	args := []string{"reset", "--force"}
+	args := []string{"gateway", "reset", "--force"}
 	if len(opts) > 0 {
 		args = append(args, opts...)
 	}