Repository: snowflakedb/gosnowflake
Pull request #1196 (2024-08-12 11:09:14)

SNOW-1563083: snowflakeConn creates/manages its own context
### Description

SNOW-1563083 Avoid persisting query context in `snowflakeConn`, using it only for dialing purposes as described in [docs](https://pkg.go.dev/database/sql/driver#Connector).

Fixes: #1186

### Checklist
- [x] Created tests which fail without the change
- [x] Extended the README / documentation (yes, inline comments added)


github-actions[bot]: <br/>Thank you for your submission, we really appreciate it. Like many open-source projects, we ask that you sign our [Contributor License Agreement](https://github.com/snowflakedb/CLA/blob/main/README.md) before we can accept your contribution. You can sign the CLA by just posting a Pull Request Comment same as the below format.<br/>
   - - -
   I have read the CLA Document and I hereby sign the CLA
   - - -
   <sub>You can retrigger this bot by commenting **recheck** in this Pull Request. Posted by the **CLA Assistant Lite bot**.</sub>
joellubi: I have read the CLA Document and I hereby sign the CLA
sfc-gh-dheyman-1: @joellubi Thank you for the contribution! You're actually right, we shouldn't propagate the context to build the connection. I'm running this PR against our CI tests and we'll get back to you (probably on Friday).

Patch:
diff --git a/connection.go b/connection.go
index 7851c30..fdd986e 100644
--- a/connection.go
+++ b/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


Test patch:
diff --git a/connector_test.go b/connector_test.go
index 7688619..d9d901b 100644
--- a/connector_test.go
+++ b/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"))
+}