Use retries when contacting Stitch's Import API

planetscale · Aug 4, 2022 · 5336834 · 5336834
1 parent dee5b06
commit 5336834
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 11 deletions.
diff --git a/cmd/internal/batch_writer.go b/cmd/internal/batch_writer.go
@@ -4,30 +4,28 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"github.com/hashicorp/go-retryablehttp"
 	"io/ioutil"
-	"net/http"
 	"time"
 )
 
-// MaxImportBatchSize represents the maximum items we can send in each batch request to Stitch.
-// Each data object in the request body cannot exceed 10,000 individual data points.
-const MaxImportBatchSize = 10000
-
 type BatchWriter interface {
 	Flush(stream *Stream) error
 	Send(record *Record, stream *Stream) error
 }
 
 func NewBatchWriter(batchSize int, logger Logger, apiURL, apiToken string) BatchWriter {
+	client := retryablehttp.NewClient()
+	// Wait 3 seconds before retrying
+	client.RetryWaitMin = 3 * time.Second
+	client.Logger = nil
 	return &httpBatchWriter{
 		batchSize: batchSize,
 		apiURL:    apiURL,
 		apiToken:  apiToken,
 		logger:    logger,
-		client: &http.Client{
-			Timeout: time.Second * 10,
-		},
-		messages: make([]ImportMessage, 0, batchSize),
+		client:    client,
+		messages:  make([]ImportMessage, 0, batchSize),
 	}
 }
 
@@ -36,7 +34,7 @@ type httpBatchWriter struct {
 	apiURL    string
 	apiToken  string
 	logger    Logger
-	client    *http.Client
+	client    *retryablehttp.Client
 	messages  []ImportMessage
 }
 
@@ -64,7 +62,7 @@ func (h *httpBatchWriter) Flush(stream *Stream) error {
 		return err
 	}
 
-	stitch, err := http.NewRequest("POST", h.apiURL+"/v2/import/batch", bytes.NewBuffer(b))
+	stitch, err := retryablehttp.NewRequest("POST", h.apiURL+"/v2/import/batch", bytes.NewBuffer(b))
 	if err != nil {
 		return err
 	}

diff --git a/go.mod b/go.mod
@@ -29,6 +29,8 @@ require (
 	github.com/google/uuid v1.3.0 // indirect
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 // indirect
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 // indirect
+	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
+	github.com/hashicorp/go-retryablehttp v0.7.1 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 // indirect
 	github.com/opentracing-contrib/go-grpc v0.0.0-20210225150812-73cb765af46e // indirect
 	github.com/opentracing/opentracing-go v1.2.0 // indirect

diff --git a/go.sum b/go.sum
@@ -380,6 +380,8 @@ github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9
 github.com/hashicorp/go-plugin v1.0.1/go.mod h1:++UyYGoz3o5w9ZzAdZxtQKrWWP+iqPBn3cQptSMzBuY=
 github.com/hashicorp/go-retryablehttp v0.5.3/go.mod h1:9B5zBasrRhHXnJnui7y6sL7es7NDiJgTc6Er0maI1Xs=
 github.com/hashicorp/go-retryablehttp v0.6.6/go.mod h1:vAew36LZh98gCBJNLH42IQ1ER/9wtLZZ8meHqQvEYWY=
+github.com/hashicorp/go-retryablehttp v0.7.1 h1:sUiuQAnLlbvmExtFQs72iFW/HXeUn8Z1aJLQ4LJJbTQ=
+github.com/hashicorp/go-retryablehttp v0.7.1/go.mod h1:vAew36LZh98gCBJNLH42IQ1ER/9wtLZZ8meHqQvEYWY=
 github.com/hashicorp/go-rootcerts v1.0.0/go.mod h1:K6zTfqpRlCUIjkwsN4Z+hiSfzSTQa6eBIzfwKfwNnHU=
 github.com/hashicorp/go-rootcerts v1.0.2 h1:jzhAVGtqPKbwpyCPELlgNWhE1znq+qwJtW5Oi2viEzc=
 github.com/hashicorp/go-rootcerts v1.0.2/go.mod h1:pqUvnprVnM5bf7AOirdbb01K4ccR319Vf4pU3K5EGc8=