From 3c185e49f909e76b29cbf3f24fd3825f5bce89a7 Mon Sep 17 00:00:00 2001
From: Joe Luttrell <joe@readytensor.com>
Date: Thu, 16 May 2024 13:13:44 +0200
Subject: [PATCH] use aws http client builder and add back aws cfg settings

---
 go.mod        |  2 +-
 go.sum        |  4 ++--
 timestream.go | 29 +++++++++++------------------
 3 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/go.mod b/go.mod
index 16ececa..6e37d19 100644
--- a/go.mod
+++ b/go.mod
@@ -32,7 +32,7 @@ require (
 )
 
 require (
-	github.com/aws/aws-sdk-go-v2 v1.26.1
+	github.com/aws/aws-sdk-go-v2 v1.26.2
 	github.com/aws/aws-sdk-go-v2/config v1.27.13
 	github.com/aws/aws-sdk-go-v2/service/timestreamquery v1.23.1
 	github.com/aws/aws-sdk-go-v2/service/timestreamwrite v1.25.6
diff --git a/go.sum b/go.sum
index 8a9921c..49ea81a 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-github.com/aws/aws-sdk-go-v2 v1.26.1 h1:5554eUqIYVWpU0YmeeYZ0wU64H2VLBs8TlhRB2L+EkA=
-github.com/aws/aws-sdk-go-v2 v1.26.1/go.mod h1:ffIFB97e2yNsv4aTSGkqtHnppsIJzw7G7BReUZ3jCXM=
+github.com/aws/aws-sdk-go-v2 v1.26.2 h1:OTRAL8EPdNoOdiq5SUhCaHhVPBU2wxAUe5uwasoJGRM=
+github.com/aws/aws-sdk-go-v2 v1.26.2/go.mod h1:ffIFB97e2yNsv4aTSGkqtHnppsIJzw7G7BReUZ3jCXM=
 github.com/aws/aws-sdk-go-v2/config v1.27.13 h1:WbKW8hOzrWoOA/+35S5okqO/2Ap8hkkFUzoW8Hzq24A=
 github.com/aws/aws-sdk-go-v2/config v1.27.13/go.mod h1:XLiyiTMnguytjRER7u5RIkhIqS8Nyz41SwAWb4xEjxs=
 github.com/aws/aws-sdk-go-v2/credentials v1.17.13 h1:XDCJDzk/u5cN7Aple7D/MiAhx1Rjo/0nueJ0La8mRuE=
diff --git a/timestream.go b/timestream.go
index 7c57e8f..9ee5634 100644
--- a/timestream.go
+++ b/timestream.go
@@ -22,13 +22,13 @@ import (
 	"context"
 	"fmt"
 	"math"
-	"net"
 	"net/http"
 	"strconv"
 	"strings"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
+	awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/timestreamquery"
 	"github.com/aws/aws-sdk-go-v2/service/timestreamwrite"
@@ -82,25 +82,18 @@ func newTimestreamQueryPaginator(client TimestreamQueryApi, params *timestreamqu
 }
 
 func newTimeStreamAdapter(ctx context.Context, logger *zap.SugaredLogger, cfg *adapterCfg, newQueryPaginator NewQueryPaginator, writeSvc *timestreamwrite.Client, querySvc *timestreamquery.Client) TimeStreamAdapter {
-	tr := &http.Transport{
-		ResponseHeaderTimeout: 20 * time.Second,
-		// Using DefaultTransport values for other parameters: https://golang.org/pkg/net/http/#RoundTripper
-		Proxy: http.ProxyFromEnvironment,
-		DialContext: (&net.Dialer{
-			KeepAlive: 30 * time.Second,
-			Timeout:   30 * time.Second,
-		}).DialContext,
-		MaxIdleConns:          100,
-		IdleConnTimeout:       90 * time.Second,
-		TLSHandshakeTimeout:   10 * time.Second,
-		ExpectContinueTimeout: 1 * time.Second,
-	}
+	if writeSvc == nil || querySvc == nil {
+		client := awshttp.NewBuildableClient().WithTransportOptions(func(tr *http.Transport) {
+			tr.ResponseHeaderTimeout = 20 * time.Second
 
-	// So client makes HTTP/2 requests
-	_ = http2.ConfigureTransport(tr)
+			// Enable HTTP/2
+			err := http2.ConfigureTransport(tr)
+			if err != nil {
+				logger.Fatalw("Unable to configure HTTP/2 transport", "error", err)
+			}
+		})
 
-	if writeSvc == nil || querySvc == nil {
-		awsCfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(cfg.awsRegion))
+		awsCfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(cfg.awsRegion), config.WithHTTPClient(client), config.WithRetryMaxAttempts(10))
 		if err != nil {
 			logger.Fatalw("Unable to load AWS SDK config", "error", err)
 		}