SNOW-1563083: snowflakeConn creates/manages its own context (#1196)

connection.go

@@ -65,6 +65,9 @@ const (
 	executionTypeStatement string  = "statement"
 )
 
+// snowflakeConn manages its own context.
+// External cancellation should not be supported because the connection
+// may be reused after the original query/request has completed.
 type snowflakeConn struct {
 	ctx                 context.Context
 	cfg                 *Config
@@ -746,10 +749,12 @@ func (scd *snowflakeArrowStreamChunkDownloader) GetBatches() (out []ArrowStreamB
 	return
 }
 
+// buildSnowflakeConn creates a new snowflakeConn.
+// The provided context is used only for establishing the initial connection.
 func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, error) {
 	sc := &snowflakeConn{
 		SequenceCounter:     0,
-		ctx:                 ctx,
+		ctx:                 context.Background(),
 		cfg:                 &config,
 		queryContextCache:   (&queryContextCache{}).init(),
 		currentTimeProvider: defaultTimeProvider,
@@ -773,7 +778,7 @@ func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, err
 		// use the custom transport
 		st = sc.cfg.Transporter
 	}
-	if err = setupOCSPEnvVars(sc.ctx, sc.cfg.Host); err != nil {
+	if err = setupOCSPEnvVars(ctx, sc.cfg.Host); err != nil {
 		return nil, err
 	}
 	var tokenAccessor TokenAccessor

connector_test.go

@@ -3,10 +3,13 @@
 package gosnowflake
 
 import (
+	"bytes"
 	"context"
 	"database/sql/driver"
 	"reflect"
+	"strings"
 	"testing"
+	"time"
 )
 
 type noopTestDriver struct {
@@ -67,3 +70,35 @@ func TestConnectorWithMissingConfig(t *testing.T) {
 	assertEqualE(t, driverErr.Number, expectedErr.Number)
 	assertEqualE(t, driverErr.Message, expectedErr.Message)
 }
+
+func TestConnectorCancelContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// restore the logger output after the test is complete.
+	logger := GetLogger().(*defaultLogger)
+	initialOutput := logger.inner.Out
+	defer logger.SetOutput(initialOutput)
+
+	// write logs to temp buffer so we can assert log output.
+	var buf bytes.Buffer
+	logger.SetOutput(&buf)
+
+	// pass in our context which should only be used for establishing the initial connection; not persisted.
+	sfConn, err := buildSnowflakeConn(ctx, Config{Params: make(map[string]*string)})
+	assertNilF(t, err)
+
+	// patch close handler
+	sfConn.rest = &snowflakeRestful{
+		FuncCloseSession: func(ctx context.Context, sr *snowflakeRestful, d time.Duration) error {
+			return ctx.Err()
+		},
+	}
+
+	// cancel context BEFORE closing the connection.
+	// this may occur if the *snowflakeConn was spawned by a QueryContext(), and the query has completed.
+	cancel()
+	assertNilF(t, sfConn.Close())
+
+	// if the following log is emitted, the connection is holding onto context that it shouldn't be.
+	assertFalseF(t, strings.Contains(buf.String(), "context canceled"))
+}