[Storage] Allow setting http headers for default http client

cli/pflags/api/generator_test.go

@@ -78,7 +78,7 @@ func TestNewGenerator(t *testing.T) {
 
 	goldenOutput, err := ioutil.ReadFile(filepath.Clean(goldenFilePath))
 	assert.NoError(t, err)
-	assert.Equal(t, goldenOutput, codeBytes)
+	assert.Equal(t, string(goldenOutput), string(codeBytes))
 
 	goldenTestOutput, err := ioutil.ReadFile(filepath.Clean(goldenTestFilePath))
 	assert.NoError(t, err)

go.sum

@@ -52,8 +52,6 @@ github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRF
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
 github.com/aws/aws-sdk-go v1.23.4/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
-github.com/aws/aws-sdk-go v1.28.8 h1:kPGnElMdW0GDc54Giy1lcE/3gAr2Gzl6cMjYKoBNFhw=
-github.com/aws/aws-sdk-go v1.28.8/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
 github.com/aws/aws-sdk-go v1.28.9 h1:grIuBQc+p3dTRXerh5+2OxSuWFi0iXuxbFdTSg0jaW0=
 github.com/aws/aws-sdk-go v1.28.9/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
 github.com/benlaurie/objecthash v0.0.0-20180202135721-d1e3d6079fc1 h1:VRtJdDi2lqc3MFwmouppm2jlm6icF+7H3WYKpLENMTo=
@@ -167,7 +165,6 @@ github.com/gorilla/websocket v1.4.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoA
 github.com/graymeta/stow v0.2.4 h1:qDGstknYXqcnmBQ5TRJtxD9Qv1MuRbYRhLoSMeUDs7U=
 github.com/graymeta/stow v0.2.4/go.mod h1:+0vRL9oMECKjPMP7OeVWl8EIqRCpFwDlth3mrAeV2Kw=
 github.com/gregjones/httpcache v0.0.0-20170728041850-787624de3eb7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
-github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
@@ -350,7 +347,6 @@ golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8U
 golang.org/x/crypto v0.0.0-20190605123033-f99c8df09eb5/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20190611184440-5c40567a22f8/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20190701094942-4def268fd1a4/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/crypto v0.0.0-20190820162420-60c769a6c586/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20191206172530-e9b2fee46413/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
@@ -526,15 +522,11 @@ honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
 k8s.io/api v0.0.0-20191016110408-35e52d86657a/go.mod h1:/L5qH+AD540e7Cetbui1tuJeXdmNhO8jM6VkXeDdDhQ=
-k8s.io/api v0.17.2/go.mod h1:BS9fjjLc4CMuqfSO8vgbHPKMt5+SF0ET6u/RVDihTo4=
 k8s.io/apimachinery v0.0.0-20191004115801-a2eda9f80ab8/go.mod h1:llRdnznGEAqC3DcNm6yEj472xaFVfLM7hnYofMb12tQ=
 k8s.io/apimachinery v0.17.2 h1:hwDQQFbdRlpnnsR64Asdi55GyCaIP/3WQpMmbNBeWr4=
 k8s.io/apimachinery v0.17.2/go.mod h1:b9qmWdKlLuU9EBh+06BtLcSf/Mu89rWL33naRxs1uZg=
 k8s.io/client-go v0.0.0-20191016111102-bec269661e48 h1:C2XVy2z0dV94q9hSSoCuTPp1KOG7IegvbdXuz9VGxoU=
 k8s.io/client-go v0.0.0-20191016111102-bec269661e48/go.mod h1:hrwktSwYGI4JK+TJA3dMaFyyvHVi/aLarVHpbs8bgCU=
-k8s.io/client-go v0.17.2/go.mod h1:QAzRgsa0C2xl4/eVpeVAZMvikCn8Nm81yqVx3Kk9XYI=
-k8s.io/client-go v11.0.0+incompatible h1:LBbX2+lOwY9flffWlJM7f1Ct8V2SRNiMRDFeiwnJo9o=
-k8s.io/client-go v11.0.0+incompatible/go.mod h1:7vJpHMYJwNQCWgzmNV+VYUl1zCObLyodBc8nIyt8L5s=
 k8s.io/gengo v0.0.0-20190128074634-0689ccc1d7d6/go.mod h1:ezvh/TsK7cY6rbqRK0oQQ8IAqLxYwwyPxAX1Pzy0ii0=
 k8s.io/klog v0.0.0-20181102134211-b9b56d5dfc92/go.mod h1:Gq+BEi5rUBO/HRz0bTSXDUcqjScdoY3a9IHpCEIOOfk=
 k8s.io/klog v0.3.0/go.mod h1:Gq+BEi5rUBO/HRz0bTSXDUcqjScdoY3a9IHpCEIOOfk=
@@ -544,7 +536,6 @@ k8s.io/klog v1.0.0/go.mod h1:4Bi6QPql/J/LkTDqv7R/cd3hPo4k2DG6Ptcz060Ez5I=
 k8s.io/kube-openapi v0.0.0-20190816220812-743ec37842bf/go.mod h1:1TqjTSzOxsLGIKfj0lK8EeCP7K1iUG65v09OM0/WG5E=
 k8s.io/kube-openapi v0.0.0-20191107075043-30be4d16710a/go.mod h1:1TqjTSzOxsLGIKfj0lK8EeCP7K1iUG65v09OM0/WG5E=
 k8s.io/utils v0.0.0-20190801114015-581e00157fb1/go.mod h1:sZAwmy6armz5eXlNoLmJcl4F1QuKu7sr+mFQ0byX7Ew=
-k8s.io/utils v0.0.0-20191114184206-e782cd3c129f/go.mod h1:sZAwmy6armz5eXlNoLmJcl4F1QuKu7sr+mFQ0byX7Ew=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 sigs.k8s.io/structured-merge-diff v0.0.0-20190525122527-15d366b2352e/go.mod h1:wWxsB5ozmmv/SG7nM11ayaAW51xMvak/t1r0CSlcokI=
 sigs.k8s.io/yaml v1.1.0/go.mod h1:UJmg0vDUVViEyp3mgSv9WPwZCDxu4rQW1olrI1uml+o=

storage/config.go

@@ -52,8 +52,15 @@ type Config struct {
 	// inputs is accelerated. The size of the cache is large so understand how to configure the cache.
 	// TODO provide some default config choices
 	// If this section is skipped, Caching is disabled
-	Cache  CachingConfig `json:"cache"`
-	Limits LimitsConfig  `json:"limits" pflag:",Sets limits for stores."`
+	Cache             CachingConfig    `json:"cache"`
+	Limits            LimitsConfig     `json:"limits" pflag:",Sets limits for stores."`
+	DefaultHTTPClient HTTPClientConfig `json:"defaultHttpClient" pflag:",Sets the default http client config."`
+}
+
+// HTTPClientConfig encapsulates common settings that can be applied to an HTTP Client.
+type HTTPClientConfig struct {
+	Headers map[string][]string `json:"headers" pflag:"-,Sets http headers to set on the http client."`
+	Timeout config.Duration     `json:"timeout" pflag:"timeout,Sets time out on the http client."`
 }
 
 // Defines connection configurations.

storage/rawstores.go

@@ -2,6 +2,7 @@ package storage
 
 import (
 	"fmt"
+	"net/http"
 
 	"github.com/lyft/flytestdlib/promutils"
 )
@@ -16,8 +17,52 @@ var stores = map[string]dataStoreCreateFn{
 	TypeStow:   newStowRawStore,
 }
 
+type proxyTransport struct {
+	http.RoundTripper
+	defaultHeaders map[string][]string
+}
+
+func (p proxyTransport) RoundTrip(r *http.Request) (resp *http.Response, err error) {
+	applyDefaultHeaders(r, p.defaultHeaders)
+	return p.RoundTripper.RoundTrip(r)
+}
+
+func applyDefaultHeaders(r *http.Request, headers map[string][]string) {
+	if r.Header == nil {
+		r.Header = http.Header{}
+	}
+
+	for key, values := range headers {
+		for _, val := range values {
+			r.Header.Add(key, val)
+		}
+	}
+}
+
+func createHTTPClient(cfg HTTPClientConfig) *http.Client {
+	c := &http.Client{
+		Timeout: cfg.Timeout.Duration,
+	}
+
+	if len(cfg.Headers) > 0 {
+		c.Transport = &proxyTransport{
+			RoundTripper:   http.DefaultTransport,
+			defaultHeaders: cfg.Headers,
+		}
+	}
+
+	return c
+}
+
 // Creates a new Data Store with the supplied config.
 func NewDataStore(cfg *Config, metricsScope promutils.Scope) (s *DataStore, err error) {
+	defaultClient := http.DefaultClient
+	defer func() {
+		http.DefaultClient = defaultClient
+	}()
+
+	http.DefaultClient = createHTTPClient(cfg.DefaultHTTPClient)
+
 	var rawStore RawStore
 	if fn, found := stores[cfg.Type]; found {
 		rawStore, err = fn(cfg, metricsScope)

storage/rawstores_test.go

@@ -0,0 +1,60 @@
+package storage
+
+import (
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/lyft/flytestdlib/config"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_createHTTPClient(t *testing.T) {
+	t.Run("empty", func(t *testing.T) {
+		client := createHTTPClient(HTTPClientConfig{})
+		assert.Nil(t, client.Transport)
+	})
+
+	t.Run("Some headers", func(t *testing.T) {
+		m := map[string][]string{
+			"Header1": {"val1", "val2"},
+		}
+
+		client := createHTTPClient(HTTPClientConfig{
+			Headers: m,
+		})
+
+		assert.NotNil(t, client.Transport)
+		proxyTransport, casted := client.Transport.(*proxyTransport)
+		assert.True(t, casted)
+		assert.Equal(t, m, proxyTransport.defaultHeaders)
+	})
+
+	t.Run("Set empty timeout", func(t *testing.T) {
+		client := createHTTPClient(HTTPClientConfig{
+			Timeout: config.Duration{},
+		})
+
+		assert.Zero(t, client.Timeout)
+	})
+
+	t.Run("Set timeout", func(t *testing.T) {
+		client := createHTTPClient(HTTPClientConfig{
+			Timeout: config.Duration{Duration: 2 * time.Second},
+		})
+
+		assert.Equal(t, 2*time.Second, client.Timeout)
+	})
+}
+
+func Test_applyDefaultHeaders(t *testing.T) {
+	input := map[string][]string{
+		"Header1": {"val1", "val2"},
+	}
+
+	r := &http.Request{}
+	applyDefaultHeaders(r, input)
+
+	assert.Equal(t, http.Header(input), r.Header)
+}

storage/stowstore.go

@@ -83,10 +83,12 @@ func newStowRawStore(cfg *Config, metricsScope promutils.Scope) (RawStore, error
 	if !ok {
 		return nil, errors.Errorf("unsupported stow.kind [%s], add support in flytestdlib?", kind)
 	}
+
 	loc, err := stow.Dial(kind, cfgMap)
 	if err != nil {
 		return emptyStore, fmt.Errorf("unable to configure the storage for %s. Error: %v", kind, err)
 	}
+
 	c, err := loc.Container(cfg.InitContainer)
 	if err != nil {
 		if IsNotFound(err) || awsBucketIsNotFound(err) {
@@ -97,8 +99,10 @@ func newStowRawStore(cfg *Config, metricsScope promutils.Scope) (RawStore, error
 			}
 			return NewStowRawStore(fn(c.Name()), c, metricsScope)
 		}
+
 		return emptyStore, err
 	}
+
 	return NewStowRawStore(fn(c.Name()), c, metricsScope)
 }
 

storage/testdata/config.yaml

@@ -9,6 +9,9 @@ connection:
   region: us-east-1
   secret-key: miniostorage
 container: ""
+defaultHttpClient:
+  headers: null
+  timeout: 0s
 limits:
   maxDownloadMBs: 0
 type: s3

tests/testdata/combined.yaml

@@ -16,6 +16,9 @@ storage:
     region: us-east-1
     secret-key: miniostorage
   container: ""
+  defaultHttpClient:
+    headers: null
+    timeout: 0s
   limits:
     maxDownloadMBs: 0
   type: s3