Repository: snowflakedb/gosnowflake
Pull request #1108 (2024-04-19 06:14:34)

SNOW-1006312: Cancel context not propagated to snowflakeFileTransferAgent on PUT/GET command
### Description

issue 822: https://github.com/snowflakedb/snowflake-sdks-drivers-issues-teamwork/issues/822
Context not propagated to snowflakeFileTransferAgent on PUT/GET command; cancellation is not supported

This PR fixes the issue by adding `ctx` to `snowflakeFileTransferAgent` and adding a channel to detect if the file transfer process has been cancelled during `processFileTransfer`.

### Checklist
- [x] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1108?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
Attention: Patch coverage is `85.71429%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 82.25%. Comparing base [(`0e1648e`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/0e1648ec97d82c4c1706069fe8df67f29311f74c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) to head [(`df23661`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/df23661afc1824a7816257e696326bb976a8a21b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).

| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1108?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Patch % | Lines |
|---|---|---|
| [connection.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1108?src=pr&el=tree&filepath=connection.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Y29ubmVjdGlvbi5nbw==) | 84.61% | [2 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1108?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #1108      +/-   ##
==========================================
- Coverage   82.48%   82.25%   -0.24%     
==========================================
  Files          54       53       -1     
  Lines       12832    12840       +8     
==========================================
- Hits        10585    10561      -24     
- Misses       2247     2279      +32     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1108?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).


Patch:
diff --git a/connection.go b/connection.go
index 6508f09..e386a63 100644
--- a/connection.go
+++ b/connection.go
@@ -169,10 +169,21 @@ func (sc *snowflakeConn) exec(
 	}
 
 	// handle PUT/GET commands
+	fileTransferChan := make(chan error, 1)
 	if isFileTransfer(query) {
-		data, err = sc.processFileTransfer(ctx, data, query, isInternal)
-		if err != nil {
-			return nil, err
+		go func() {
+			data, err = sc.processFileTransfer(ctx, data, query, isInternal)
+			fileTransferChan <- err
+		}()
+
+		select {
+		case <-ctx.Done():
+			logger.WithContext(ctx).Info("File transfer has been cancelled")
+			return nil, ctx.Err()
+		case err := <-fileTransferChan:
+			if err != nil {
+				return nil, err
+			}
 		}
 	}
 
diff --git a/connection_util.go b/connection_util.go
index 31b97c9..78ea9d1 100644
--- a/connection_util.go
+++ b/connection_util.go
@@ -89,6 +89,7 @@ func (sc *snowflakeConn) processFileTransfer(
 	isInternal bool) (
 	*execResponse, error) {
 	sfa := snowflakeFileTransferAgent{
+		ctx:          ctx,
 		sc:           sc,
 		data:         &data.Data,
 		command:      query,
diff --git a/file_transfer_agent.go b/file_transfer_agent.go
index 03cd1dc..c22d748 100644
--- a/file_transfer_agent.go
+++ b/file_transfer_agent.go
@@ -106,6 +106,7 @@ type SnowflakeFileTransferOptions struct {
 }
 
 type snowflakeFileTransferAgent struct {
+	ctx                         context.Context
 	sc                          *snowflakeConn
 	data                        *execResponseData
 	command                     string


Test patch:
diff --git a/file_transfer_agent_test.go b/file_transfer_agent_test.go
index b94261c..3ffd6a8 100644
--- a/file_transfer_agent_test.go
+++ b/file_transfer_agent_test.go
@@ -30,6 +30,7 @@ func TestGetBucketAccelerateConfiguration(t *testing.T) {
 	}
 	runSnowflakeConnTest(t, func(sct *SCTest) {
 		sfa := &snowflakeFileTransferAgent{
+			ctx:         context.Background(),
 			sc:          sct.sc,
 			commandType: uploadCommand,
 			srcFiles:    make([]string, 0),
@@ -81,6 +82,7 @@ func TestUnitDownloadWithInvalidLocalPath(t *testing.T) {
 func TestUnitGetLocalFilePathFromCommand(t *testing.T) {
 	runSnowflakeConnTest(t, func(sct *SCTest) {
 		sfa := &snowflakeFileTransferAgent{
+			ctx:         context.Background(),
 			sc:          sct.sc,
 			commandType: uploadCommand,
 			srcFiles:    make([]string, 0),
@@ -110,6 +112,7 @@ func TestUnitGetLocalFilePathFromCommand(t *testing.T) {
 func TestUnitProcessFileCompressionType(t *testing.T) {
 	runSnowflakeConnTest(t, func(sct *SCTest) {
 		sfa := &snowflakeFileTransferAgent{
+			ctx:         context.Background(),
 			sc:          sct.sc,
 			commandType: uploadCommand,
 			srcFiles:    make([]string, 0),
@@ -156,6 +159,7 @@ func TestUnitProcessFileCompressionType(t *testing.T) {
 func TestParseCommandWithInvalidStageLocation(t *testing.T) {
 	runSnowflakeConnTest(t, func(sct *SCTest) {
 		sfa := &snowflakeFileTransferAgent{
+			ctx:         context.Background(),
 			sc:          sct.sc,
 			commandType: uploadCommand,
 			srcFiles:    make([]string, 0),
@@ -190,6 +194,7 @@ func TestParseCommandEncryptionMaterialMismatchError(t *testing.T) {
 		}
 
 		sfa := &snowflakeFileTransferAgent{
+			ctx:         context.Background(),
 			sc:          sct.sc,
 			commandType: uploadCommand,
 			srcFiles:    make([]string, 0),
@@ -226,6 +231,7 @@ func TestParseCommandInvalidStorageClientException(t *testing.T) {
 		}
 
 		sfa := &snowflakeFileTransferAgent{
+			ctx:         context.Background(),
 			sc:          sct.sc,
 			commandType: uploadCommand,
 			srcFiles:    make([]string, 0),
@@ -253,6 +259,7 @@ func TestParseCommandInvalidStorageClientException(t *testing.T) {
 func TestInitFileMetadataError(t *testing.T) {
 	runSnowflakeConnTest(t, func(sct *SCTest) {
 		sfa := &snowflakeFileTransferAgent{
+			ctx:         context.Background(),
 			sc:          sct.sc,
 			commandType: uploadCommand,
 			srcFiles:    []string{"fileDoesNotExist.txt"},
@@ -352,6 +359,7 @@ func TestUpdateMetadataWithPresignedUrl(t *testing.T) {
 
 		sct.sc.rest.FuncPostQuery = presignedURLMock
 		sfa := &snowflakeFileTransferAgent{
+			ctx:               context.Background(),
 			sc:                sct.sc,
 			commandType:       uploadCommand,
 			command:           "put file:///tmp/test_data/data1.txt @~",
@@ -400,6 +408,7 @@ func TestUpdateMetadataWithPresignedUrlForDownload(t *testing.T) {
 		}
 
 		sfa := &snowflakeFileTransferAgent{
+			ctx:               context.Background(),
 			sc:                sct.sc,
 			commandType:       downloadCommand,
 			command:           "get @~/data1.txt.gz file:///tmp/testData",
@@ -421,6 +430,7 @@ func TestUpdateMetadataWithPresignedUrlForDownload(t *testing.T) {
 func TestUpdateMetadataWithPresignedUrlError(t *testing.T) {
 	runSnowflakeConnTest(t, func(sct *SCTest) {
 		sfa := &snowflakeFileTransferAgent{
+			ctx:               context.Background(),
 			sc:                sct.sc,
 			command:           "get @~/data1.txt.gz file:///tmp/testData",
 			stageLocationType: gcsClient,
@@ -486,6 +496,7 @@ func TestUploadWhenFilesystemReadOnlyError(t *testing.T) {
 	}
 
 	sfa := &snowflakeFileTransferAgent{
+		ctx: context.Background(),
 		sc: &snowflakeConn{
 			cfg: &Config{},
 		},
@@ -585,6 +596,7 @@ func TestCustomTmpDirPath(t *testing.T) {
 	}
 
 	sfa := snowflakeFileTransferAgent{
+		ctx: context.Background(),
 		sc: &snowflakeConn{
 			cfg: &Config{
 				TmpDirPath: tmpDir,
@@ -646,6 +658,7 @@ func TestReadonlyTmpDirPathShouldFail(t *testing.T) {
 	}
 
 	sfa := snowflakeFileTransferAgent{
+		ctx: context.Background(),
 		sc: &snowflakeConn{
 			cfg: &Config{
 				TmpDirPath: tmpDir,
@@ -720,6 +733,7 @@ func testUploadDownloadOneFile(t *testing.T, isStream bool) {
 	}
 
 	sfa := snowflakeFileTransferAgent{
+		ctx: context.Background(),
 		sc: &snowflakeConn{
 			cfg: &Config{
 				TmpDirPath: tmpDir,
diff --git a/put_get_test.go b/put_get_test.go
index 2c03822..8a4322a 100644
--- a/put_get_test.go
+++ b/put_get_test.go
@@ -11,10 +11,12 @@ import (
 	"math/rand"
 	"os"
 	"os/user"
+	"path"
 	"path/filepath"
 	"strconv"
 	"strings"
 	"testing"
+	"time"
 )
 
 const createStageStmt = "CREATE OR REPLACE STAGE %v URL = '%v' CREDENTIALS = (%v)"
@@ -48,6 +50,7 @@ func TestPutError(t *testing.T) {
 	}
 
 	fta := &snowflakeFileTransferAgent{
+		ctx:  context.Background(),
 		data: data,
 		options: &SnowflakeFileTransferOptions{
 			RaisePutGetError: false,
@@ -64,6 +67,7 @@ func TestPutError(t *testing.T) {
 	}
 
 	fta = &snowflakeFileTransferAgent{
+		ctx:  context.Background(),
 		data: data,
 		options: &SnowflakeFileTransferOptions{
 			RaisePutGetError: true,
@@ -829,3 +833,34 @@ func TestPutGetMaxLOBSize(t *testing.T) {
 		}
 	})
 }
+
+func TestPutCancel(t *testing.T) {
+	sourceDir, err := os.Getwd()
+	assertNilF(t, err)
+	testData := path.Join(sourceDir, "/test_data/largefile.txt")
+
+	runDBTest(t, func(dbt *DBTest) {
+		c := make(chan error)
+		ctx, cancel := context.WithCancel(context.Background())
+		go func() {
+			// attempt to upload a large file, but it should be canceled in 3 seconds
+			_, err := dbt.conn.ExecContext(
+				ctx,
+				fmt.Sprintf("put 'file://%v' @~/test_put_cancel overwrite=true",
+					strings.ReplaceAll(testData, "\\", "/")))
+			if err != nil {
+				c <- err
+				return
+			}
+			c <- nil
+		}()
+		// cancel after 3 seconds
+		time.Sleep(3 * time.Second)
+		fmt.Println("Canceled")
+		cancel()
+		ret := <-c
+		assertNotNilF(t, ret)
+		assertStringContainsF(t, ret.Error(), "context canceled", "failed to cancel.")
+		close(c)
+	})
+}
diff --git a/telemetry_test.go b/telemetry_test.go
index 4ff9447..6c6835f 100644
--- a/telemetry_test.go
+++ b/telemetry_test.go
@@ -53,6 +53,7 @@ func TestTelemetrySQLException(t *testing.T) {
 			flushSize: defaultFlushSize,
 		}
 		sfa := &snowflakeFileTransferAgent{
+			ctx:         context.Background(),
 			sc:          sct.sc,
 			commandType: uploadCommand,
 			srcFiles:    make([]string, 0),