Skip to content

Commit 0c28a5a

Browse files
Enable Transport configuration for http client (#129)
### Context: This change enables passing a RoundTripper as a configuration parameter (via the WithTransport() func) to use, for instance, a Socks5 proxy. If the transport is not configured, it remains with the original behavior. ### Changes made: - Changed Transport.Base to be of interface `http.RoundTripper` (is returned as that at `PooledTransport()` ) - Added 'WithTransport()' func to set the transport configuration. - Added validation to existing test case.
2 parents f0448ae + 54aef94 commit 0c28a5a

File tree

4 files changed

+36
-4
lines changed

4 files changed

+36
-4
lines changed

connector.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,3 +214,10 @@ func WithAuthenticator(authr auth.Authenticator) connOption {
214214
c.Authenticator = authr
215215
}
216216
}
217+
218+
// WithTransport sets up the transport configuration to be used by the httpclient.
219+
func WithTransport(t http.RoundTripper) connOption {
220+
return func(c *config.Config) {
221+
c.Transport = t
222+
}
223+
}

connector_test.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package dbsql
22

33
import (
4+
"net/http"
45
"testing"
56
"time"
67

@@ -22,6 +23,7 @@ func TestNewConnector(t *testing.T) {
2223
schema := "schema-string"
2324
userAgentEntry := "user-agent"
2425
sessionParams := map[string]string{"key": "value"}
26+
roundTripper := mockRoundTripper{}
2527
con, err := NewConnector(
2628
WithServerHostname(host),
2729
WithPort(port),
@@ -33,6 +35,7 @@ func TestNewConnector(t *testing.T) {
3335
WithUserAgentEntry(userAgentEntry),
3436
WithSessionParams(sessionParams),
3537
WithRetries(10, 3*time.Second, 60*time.Second),
38+
WithTransport(roundTripper),
3639
)
3740
expectedUserConfig := config.UserConfig{
3841
Host: host,
@@ -50,6 +53,7 @@ func TestNewConnector(t *testing.T) {
5053
RetryMax: 10,
5154
RetryWaitMin: 3 * time.Second,
5255
RetryWaitMax: 60 * time.Second,
56+
Transport: roundTripper,
5357
}
5458
expectedCfg := config.WithDefaults()
5559
expectedCfg.DriverVersion = DriverVersion
@@ -127,3 +131,11 @@ func TestNewConnector(t *testing.T) {
127131
assert.Equal(t, expectedCfg, coni.cfg)
128132
})
129133
}
134+
135+
type mockRoundTripper struct{}
136+
137+
var _ http.RoundTripper = mockRoundTripper{}
138+
139+
func (m mockRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {
140+
return &http.Response{StatusCode: 200}, nil
141+
}

internal/client/client.go

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -328,7 +328,7 @@ func isRetryableServerResponse(resp *http.Response) bool {
328328
}
329329

330330
type Transport struct {
331-
Base *http.Transport
331+
Base http.RoundTripper
332332
Authr auth.Authenticator
333333
trace bool
334334
}
@@ -405,10 +405,20 @@ func PooledClient(cfg *config.Config) *http.Client {
405405
if cfg.Authenticator == nil {
406406
return nil
407407
}
408-
tr := &Transport{
409-
Base: PooledTransport(),
410-
Authr: cfg.Authenticator,
408+
409+
var tr *Transport
410+
if cfg.Transport != nil {
411+
tr = &Transport{
412+
Base: cfg.Transport,
413+
Authr: cfg.Authenticator,
414+
}
415+
} else {
416+
tr = &Transport{
417+
Base: PooledTransport(),
418+
Authr: cfg.Authenticator,
419+
}
411420
}
421+
412422
return &http.Client{
413423
Transport: tr,
414424
Timeout: cfg.ClientTimeout,

internal/config/config.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"crypto/tls"
66
"fmt"
7+
"net/http"
78
"net/url"
89
"strconv"
910
"strings"
@@ -97,6 +98,7 @@ type UserConfig struct {
9798
RetryWaitMin time.Duration
9899
RetryWaitMax time.Duration
99100
RetryMax int
101+
Transport http.RoundTripper
100102
}
101103

102104
// DeepCopy returns a true deep copy of UserConfig
@@ -135,6 +137,7 @@ func (ucfg UserConfig) DeepCopy() UserConfig {
135137
RetryWaitMin: ucfg.RetryWaitMin,
136138
RetryWaitMax: ucfg.RetryWaitMax,
137139
RetryMax: ucfg.RetryMax,
140+
Transport: ucfg.Transport,
138141
}
139142
}
140143

0 commit comments

Comments
 (0)