Context not propagated to snowflakeFileTransferAgent on PUT command; cancellation is not supported

[joellubi]

Context cancellation works for other queries executed using conn.ExecContext(), but is ignored for PUT/GET commands because it is not propagated any further than the snowflakeConn.processFileTransfer() method.

1. What version of GO driver are you using?
   github.com/snowflakedb/gosnowflake v1.7.1

2. What operating system and processor architecture are you using?
   macOS 14.1.1 x86_64

3. What version of GO are you using?
   go version go1.21.5 darwin/amd64

4. Server version:
   8.1.0

5. What did you do?

package main

import (
	"context"
	"database/sql"
	"log"
	"time"

	"github.com/snowflakedb/gosnowflake"
)

func main() {
	cfg := &gosnowflake.Config{
		// Redacted
	}

	dsn, err := gosnowflake.DSN(cfg)
	if err != nil {
		log.Fatal(err)
	}

	db, err := sql.Open("snowflake", dsn)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	conn, err := db.Conn(ctx)
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close()

	stageName := "my_temp_stage"
	createStageQuery := "CREATE OR REPLACE STAGE my_temp_stage" + stageName
	putQuery := "PUT file:///path/to/test.parquet @" + stageName + " OVERWRITE=true"

	_, err = conn.ExecContext(ctx, createStageQuery)
	if err != nil {
		log.Fatal(err)
	}

	go func() {
		time.Sleep(10 * time.Second)
		log.Println("Canceled")
		cancel()
	}()

	log.Println("Starting PUT upload")
	_, err = conn.ExecContext(ctx, putQuery)
	if err != nil {
		log.Fatal(err)
	}
	log.Println("Completed PUT upload")
}

6. What did you expect to see?

Expected Logs

% go run .
2024/01/10 21:11:22 Starting PUT upload
2024/01/10 21:11:32 Canceled
2024/01/10 21:11:32 context canceled
exit status 1

Actual Logs

% go run .
2024/01/10 21:11:22 Starting PUT upload
2024/01/10 21:11:32 Canceled
2024/01/10 21:12:36 Completed PUT upload

7. Can you set logging to DEBUG and collect the logs?

   https://community.snowflake.com/s/article/How-to-generate-log-file-on-Snowflake-connectors

Yes, if it will help. But it's clear in the code that ctx is not passed all the way down to the upload methods.

8. What is your Snowflake account identifier, if any? (Optional)

[sfc-gh-dszmolka]

Hi and thank you for raising this issue; especially with such a nice self contained reproduction! It indeed reproduces and when replacing putQuery with something else like a SELECT..., it works as expected. But not with PUT/GET.
We'll take a look.

[sfc-gh-dszmolka]

Short update; indeed the processFileTransfer does not seem to contain WithCancel. We're working on a fix and I'll keep this thread posted.

[sfc-gh-dszmolka]

we have a draft PR at #1108 for this, will keep this thread updated