Repository: snowflakedb/gosnowflake
Pull request #364 (2021-01-01 20:00:45)

SNOW-209089 Add HTTP transport override to support future token accessor
SNOW-209089

### Description
In order to support manual auth token accessor / management, we
need a clean way to intercept HTTP requests and responses and manually
set AUTH header, retry, and reauth.

Add an optional value to the config so that an override Http transport
can be specified. When the Transport is specified, the REST client will
use this transport instead of the default snowflake transport.

Testing
Added unit test to verify that config transport is indeed the one
used by the snowflake restful http client

This is built upon https://github.com/snowflakedb/gosnowflake/pull/343 which allows
opening connection with custom Config

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [ ] Extended the README / documentation, if necessary



Patch:
diff --git a/connector.go b/connector.go
new file mode 100644
index 0000000..a70241a
--- /dev/null
+++ b/connector.go
@@ -0,0 +1,40 @@
+// Copyright (c) 2020 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"context"
+	"database/sql/driver"
+)
+
+// InternalSnowflakeDriver is the interface for an internal Snowflake driver
+type InternalSnowflakeDriver interface {
+	Open(dsn string) (driver.Conn, error)
+	OpenWithConfig(ctx context.Context, config Config) (driver.Conn, error)
+}
+
+// Connector creates Driver with the specified Config
+type Connector struct {
+	driver InternalSnowflakeDriver
+	cfg    Config
+}
+
+// NewConnector creates a new connector with the given SnowflakeDriver and Config.
+func NewConnector(driver InternalSnowflakeDriver, config Config) Connector {
+	return Connector{driver, config}
+}
+
+// Connect creates a new connection.
+func (t Connector) Connect(ctx context.Context) (driver.Conn, error) {
+	cfg := t.cfg
+	err := fillMissingConfigParameters(&cfg)
+	if err != nil {
+		return nil, err
+	}
+	return t.driver.OpenWithConfig(ctx, cfg)
+}
+
+// Driver creates a new driver.
+func (t Connector) Driver() driver.Driver {
+	return t.driver
+}
diff --git a/doc.go b/doc.go
index 52a345e..16292be 100644
--- a/doc.go
+++ b/doc.go
@@ -125,6 +125,8 @@ A complete connection string looks similar to the following:
 Session-level parameters can also be set by using the SQL command "ALTER SESSION"
 (https://docs.snowflake.com/en/sql-reference/sql/alter-session.html).
 
+Alternatively, use OpenWithConfig() function to create a database handle with the specified Config.
+
 Proxy
 
 The Go Snowflake Driver honors the environment variables HTTP_PROXY, HTTPS_PROXY and NO_PROXY for the forward proxy setting.
diff --git a/driver.go b/driver.go
index e77f6cd..149ecc1 100644
--- a/driver.go
+++ b/driver.go
@@ -17,26 +17,37 @@ type SnowflakeDriver struct {
 // Open creates a new connection.
 func (d SnowflakeDriver) Open(dsn string) (driver.Conn, error) {
 	logger.Info("Open")
+	ctx := context.TODO()
+	cfg, err := ParseDSN(dsn)
+	if err != nil {
+		return nil, err
+	}
+	return d.OpenWithConfig(ctx, *cfg)
+}
+
+// OpenWithConfig creates a new connection with the given Config.
+func (d SnowflakeDriver) OpenWithConfig(ctx context.Context, config Config) (driver.Conn, error) {
+	logger.Info("OpenWithConfig")
 	var err error
 	sc := &snowflakeConn{
 		SequenceCounter: 0,
-		ctx:             context.TODO(),
+		ctx: ctx,
+		cfg: &config,
 	}
-
-	sc.cfg, err = ParseDSN(dsn)
-	if err != nil {
-		sc.cleanup()
-		return nil, err
-	}
-	st := SnowflakeTransport
-	if sc.cfg.InsecureMode {
-		// no revocation check with OCSP. Think twice when you want to enable this option.
-		st = snowflakeInsecureTransport
+	var st http.RoundTripper = SnowflakeTransport
+	if sc.cfg.Transporter == nil {
+		if sc.cfg.InsecureMode {
+			// no revocation check with OCSP. Think twice when you want to enable this option.
+			st = snowflakeInsecureTransport
+		} else {
+			// set OCSP fail open mode
+			ocspResponseCacheLock.Lock()
+			atomic.StoreUint32((*uint32)(&ocspFailOpen), uint32(sc.cfg.OCSPFailOpen))
+			ocspResponseCacheLock.Unlock()
+		}
 	} else {
-		// set OCSP fail open mode
-		ocspResponseCacheLock.Lock()
-		atomic.StoreUint32((*uint32)(&ocspFailOpen), uint32(sc.cfg.OCSPFailOpen))
-		ocspResponseCacheLock.Unlock()
+		// use the custom transport
+		st = sc.cfg.Transporter
 	}
 	// authenticate
 	sc.rest = &snowflakeRestful{
diff --git a/dsn.go b/dsn.go
index d7b2b37..56c71c3 100644
--- a/dsn.go
+++ b/dsn.go
@@ -6,6 +6,7 @@ import (
 	"crypto/rsa"
 	"encoding/base64"
 	"fmt"
+	"net/http"
 	"net/url"
 	"strconv"
 	"strings"
@@ -70,6 +71,8 @@ type Config struct {
 	Token string // Token to use for OAuth other forms of token based auth
 
 	PrivateKey *rsa.PrivateKey // Private key used to sign JWT
+
+	Transporter http.RoundTripper // RoundTripper to intercept HTTP requests and responses
 }
 
 // ocspMode returns the OCSP mode in string INSECURE, FAIL_OPEN, FAIL_CLOSED
@@ -413,7 +416,7 @@ func fillMissingConfigParameters(cfg *Config) error {
 	return nil
 }
 
-// transformAccountToHost transforms host to accout name
+// transformAccountToHost transforms host to account name
 func transformAccountToHost(cfg *Config) (err error) {
 	if cfg.Port == 0 && !strings.HasSuffix(cfg.Host, defaultDomain) && cfg.Host != "" {
 		// account name is specified instead of host:port


Test patch:
diff --git a/connector_test.go b/connector_test.go
new file mode 100644
index 0000000..4bf46ba
--- /dev/null
+++ b/connector_test.go
@@ -0,0 +1,49 @@
+// Copyright (c) 2020 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"context"
+	"database/sql/driver"
+	"reflect"
+	"testing"
+)
+
+type noopTestDriver struct {
+	config Config
+	conn   *snowflakeConn
+}
+
+func (d noopTestDriver) Open(_ string) (driver.Conn, error) {
+	return nil, nil
+}
+
+func (d noopTestDriver) OpenWithConfig(_ context.Context, config Config) (driver.Conn, error) {
+	d.config = config
+	return d.conn, nil
+}
+
+func TestConnector(t *testing.T) {
+	conn := snowflakeConn{}
+	mock := noopTestDriver{conn: &conn}
+	createDSN("UTC")
+	config, err := ParseDSN(dsn)
+	if err != nil {
+		t.Fatalf("Failed to parse dsn %s", dsn)
+	}
+	connector := NewConnector(mock, *config)
+	connection, err := connector.Connect(context.Background())
+	if err != nil {
+		t.Fatalf("Connect error %s", err)
+	}
+	if connection != &conn {
+		t.Fatalf("Connection mismatch %s", connection)
+	}
+	fillMissingConfigParameters(config)
+	if reflect.DeepEqual(config, mock.config) {
+		t.Fatalf("Config should be equal, expected %v, actual %v", config, mock.config)
+	}
+	if connector.Driver() == nil {
+		t.Fatalf("Missing driver")
+	}
+}
diff --git a/driver_test.go b/driver_test.go
index d97012e..a3093ce 100644
--- a/driver_test.go
+++ b/driver_test.go
@@ -10,6 +10,7 @@ import (
 	"flag"
 	"fmt"
 	"math/big"
+	"net/http"
 	"net/url"
 	"os"
 	"os/signal"
@@ -2388,6 +2389,67 @@ func TestPingInvalidHost(t *testing.T) {
 	}
 }
 
+func TestOpenWithConfig(t *testing.T) {
+	config, err := ParseDSN(dsn)
+	if err != nil {
+		t.Fatalf("failed to parse dsn. dsn: %v, err: %v", dsn, err)
+	}
+	driver := SnowflakeDriver{}
+	db, err := driver.OpenWithConfig(context.Background(), *config)
+	if err != nil {
+		t.Fatalf("failed to open with config. config: %v, err: %v", config, err)
+	}
+	db.Close()
+}
+
+type CountingTransport struct {
+	requests int
+}
+
+func (t *CountingTransport) RoundTrip(r *http.Request) (*http.Response, error) {
+	t.requests++
+	return snowflakeInsecureTransport.RoundTrip(r)
+}
+
+func TestOpenWithTransport(t *testing.T) {
+	config, err := ParseDSN(dsn)
+	if err != nil {
+		t.Fatalf("failed to parse dsn. dsn: %v, err: %v", dsn, err)
+	}
+	countingTransport := CountingTransport{}
+	var transport http.RoundTripper = &countingTransport
+	config.Transporter = transport
+	driver := SnowflakeDriver{}
+	db, err := driver.OpenWithConfig(context.Background(), *config)
+	if err != nil {
+		t.Fatalf("failed to open with config. config: %v, err: %v", config, err)
+	}
+	conn := db.(*snowflakeConn)
+	if conn.rest.Client.Transport != transport {
+		t.Fatal("transport doesn't match")
+	}
+	db.Close()
+	if countingTransport.requests == 0 {
+		t.Fatal("transport did not receive any requests")
+	}
+
+	// Test that transport override also works in insecure mode
+	countingTransport.requests = 0
+	config.InsecureMode = true
+	db, err = driver.OpenWithConfig(context.Background(), *config)
+	if err != nil {
+		t.Fatalf("failed to open with config. config: %v, err: %v", config, err)
+	}
+	conn = db.(*snowflakeConn)
+	if conn.rest.Client.Transport != transport {
+		t.Fatal("transport doesn't match")
+	}
+	db.Close()
+	if countingTransport.requests == 0 {
+		t.Fatal("transport did not receive any requests")
+	}
+}
+
 func createDSNWithClientSessionKeepAlive() {
 	dsn = fmt.Sprintf("%s:%s@%s/%s/%s", user, pass, host, dbname, schemaname)