From a19bf2ef076916c45f38ee0e0ad093b652db1268 Mon Sep 17 00:00:00 2001
From: Ronen Lubin <ronen7822@gmail.com>
Date: Tue, 31 Dec 2024 13:31:20 +0200
Subject: [PATCH] atlasaction/cloud: use http retry client

---
 atlasaction/cloud/client.go      | 34 ++++++++++++++++++++------------
 atlasaction/cloud/client_test.go | 33 +++++++++++++++++++++++++++++++
 go.mod                           |  2 ++
 go.sum                           |  4 ++++
 4 files changed, 60 insertions(+), 13 deletions(-)
 create mode 100644 atlasaction/cloud/client_test.go

diff --git a/atlasaction/cloud/client.go b/atlasaction/cloud/client.go
index 5b64546f..165a9205 100644
--- a/atlasaction/cloud/client.go
+++ b/atlasaction/cloud/client.go
@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/hashicorp/go-retryablehttp"
 	"github.com/vektah/gqlparser/gqlerror"
 )
 
@@ -17,7 +18,7 @@ const cloudURL = "https://api.atlasgo.cloud/query"
 type (
 	// Client is a client for the Atlas Cloud API.
 	Client struct {
-		client   *http.Client
+		client   *retryablehttp.Client
 		endpoint string
 	}
 	// roundTripper is a http.RoundTripper that adds the Authorization header.
@@ -34,21 +35,28 @@ func (r *roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	return http.DefaultTransport.RoundTrip(req)
 }
 
-// New creates a new Client for the Atlas Cloud API.
-func New(token, version, cliVersion string) *Client {
+func newClient(endpoint, token, version, cliVersion string) *Client {
+	if endpoint == "" {
+		endpoint = cloudURL
+	}
+	client := retryablehttp.NewClient()
+	client.HTTPClient.Timeout = time.Second * 30
+	client.HTTPClient.Transport = &roundTripper{
+		token:      token,
+		version:    version,
+		cliVersion: cliVersion,
+	}
 	return &Client{
-		endpoint: cloudURL,
-		client: &http.Client{
-			Transport: &roundTripper{
-				token:      token,
-				version:    version,
-				cliVersion: cliVersion,
-			},
-			Timeout: time.Second * 30,
-		},
+		endpoint: endpoint,
+		client:   client,
 	}
 }
 
+// New creates a new Client for the Atlas Cloud API.
+func New (token, version, cliVersion string) *Client {
+	return newClient("", token, version, cliVersion)
+}
+
 type (
 	ScopeIdent struct {
 		URL     string   `json:"url"`               // The (Atlas CLI) URL of the database snapshot.
@@ -133,7 +141,7 @@ func (c *Client) post(ctx context.Context, query string, vars, data any) error {
 	if err != nil {
 		return err
 	}
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, bytes.NewReader(body))
+	req, err := retryablehttp.NewRequestWithContext(ctx, http.MethodPost, c.endpoint, bytes.NewReader(body))
 	if err != nil {
 		return err
 	}
diff --git a/atlasaction/cloud/client_test.go b/atlasaction/cloud/client_test.go
new file mode 100644
index 00000000..d982f068
--- /dev/null
+++ b/atlasaction/cloud/client_test.go
@@ -0,0 +1,33 @@
+package cloud
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestClient_Retry(t *testing.T) {
+	var (
+		ctx   = context.Background()
+		calls = []int{http.StatusInternalServerError, http.StatusInternalServerError, http.StatusOK}
+		srv   = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			code := calls[0]
+			w.WriteHeader(code)
+			calls = calls[1:]
+			if code != http.StatusOK {
+				return
+			}
+			_, err := fmt.Fprint(w, `{"data":{"snapshotHash":{"hash":"hash"}}}`)
+			require.NoError(t, err)
+		}))
+		client = newClient(srv.URL, "token", "version", "cliVersion")
+	)
+	defer srv.Close()
+	_, err := client.SnapshotHash(ctx, &SnapshotHashInput{})
+	require.NoError(t, err)
+	require.Empty(t, calls)
+}
diff --git a/go.mod b/go.mod
index 8b3eab52..6057a86c 100644
--- a/go.mod
+++ b/go.mod
@@ -23,6 +23,8 @@ require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/go-openapi/inflect v0.19.0 // indirect
 	github.com/google/go-cmp v0.6.0 // indirect
+	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
+	github.com/hashicorp/go-retryablehttp v0.7.7 // indirect
 	github.com/hashicorp/hcl/v2 v2.18.1 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
diff --git a/go.sum b/go.sum
index 5c8b95a9..399cef3a 100644
--- a/go.sum
+++ b/go.sum
@@ -30,6 +30,10 @@ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY=
 github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ=
+github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
+github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
+github.com/hashicorp/go-retryablehttp v0.7.7 h1:C8hUCYzor8PIfXHa4UrZkU4VvK8o9ISHxT2Q8+VepXU=
+github.com/hashicorp/go-retryablehttp v0.7.7/go.mod h1:pkQpWZeYWskR+D1tR2O5OcBFOxfA7DoAO6xtkuQnHTk=
 github.com/hashicorp/hcl/v2 v2.18.1 h1:6nxnOJFku1EuSawSD81fuviYUV8DxFr3fp2dUi3ZYSo=
 github.com/hashicorp/hcl/v2 v2.18.1/go.mod h1:ThLC89FV4p9MPW804KVbe/cEXoQ8NZEh+JtMeeGErHE=
 github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=