Repository: snowflakedb/gosnowflake
Pull request #370 (2021-01-26 19:48:59)

SNOW-212952 Emit Query ID
### Description
Interface to capture query ID as soon as it's ready using a channel passed through a context.

### 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/connection.go b/connection.go
index 84c70b6..9d8373d 100644
--- a/connection.go
+++ b/connection.go
@@ -644,7 +644,6 @@ func getAsync(
 		if resType == execResultType {
 			res.affectedRows, _ = updateRows(respd.Data)
 			res.insertID = -1
-			res.queryID = respd.Data.QueryID
 			res.errChannel <- nil // mark exec status complete
 		} else {
 			sc := &snowflakeConn{rest: sr}
@@ -669,7 +668,6 @@ func getAsync(
 					RowSetBase64: respd.Data.RowSetBase64,
 				},
 			}
-			rows.queryID = respd.Data.QueryID
 			rows.ChunkDownloader.start()
 			rows.errChannel <- nil // mark query status complete
 		}
@@ -692,6 +690,15 @@ func getAsync(
 	}
 }
 
+func getQueryIDChan(ctx context.Context) chan<- string {
+	v := ctx.Value(QueryIDChan)
+	if v == nil {
+		return nil
+	}
+	c, _ := v.(chan<- string)
+	return c
+}
+
 func populateChunkDownloader(ctx context.Context, sc *snowflakeConn, data execResponseData) *snowflakeChunkDownloader {
 	return &snowflakeChunkDownloader{
 		sc:                 sc,
diff --git a/restful.go b/restful.go
index 8c03914..4e32941 100644
--- a/restful.go
+++ b/restful.go
@@ -201,6 +201,12 @@ func postRestfulQueryHelper(
 			return sr.FuncPostQuery(ctx, sr, params, headers, body, timeout, requestID)
 		}
 
+		if queryIDChan := getQueryIDChan(ctx); queryIDChan != nil {
+			queryIDChan <- respd.Data.QueryID
+			close(queryIDChan)
+			ctx = WithQueryIDChan(ctx, nil)
+		}
+
 		var resultURL string
 		isSessionRenewed := false
 		noResult, _ := isAsyncMode(ctx)
diff --git a/statement.go b/statement.go
index cce8211..73a5af3 100644
--- a/statement.go
+++ b/statement.go
@@ -14,6 +14,8 @@ const (
 	MultiStatementCount paramKey = "MULTI_STATEMENT_COUNT"
 	// AsyncMode tells the server to not block the request on executing the entire query
 	AsyncMode paramKey = "ASYNC_MODE_QUERY"
+	// QueryIDChan is the channel to receive the query ID from
+	QueryIDChan paramKey = "QUERY_ID_CHANNEL"
 )
 
 type snowflakeStmt struct {
@@ -62,3 +64,8 @@ func WithMultiStatement(ctx context.Context, num int) (context.Context, error) {
 func WithAsyncMode(ctx context.Context) (context.Context, error) {
 	return context.WithValue(ctx, AsyncMode, true), nil
 }
+
+// WithQueryIDChan returns a context that contains the channel to receive the query ID
+func WithQueryIDChan(ctx context.Context, c chan<- string) context.Context {
+	return context.WithValue(ctx, QueryIDChan, c)
+}


Test patch:
diff --git a/statement_test.go b/statement_test.go
index 642fdc2..b11d382 100644
--- a/statement_test.go
+++ b/statement_test.go
@@ -7,6 +7,7 @@ import (
 	"database/sql"
 	"database/sql/driver"
 	"fmt"
+	"io"
 	"testing"
 )
 
@@ -630,3 +631,63 @@ func retrieveRows(rows *sql.Rows, ch chan string) {
 	ch <- s
 	close(ch)
 }
+
+func TestEmitQueryID(t *testing.T) {
+	var db *sql.DB
+	var err error
+
+	if db, err = sql.Open("snowflake", dsn); err != nil {
+		t.Fatalf("failed to open db. %v, err: %v", dsn, err)
+	}
+	defer db.Close()
+
+	queryIDChan := make(chan string, 1)
+	numrows := 100000
+	ctx, _ := WithAsyncMode(context.Background())
+	ctx = WithQueryIDChan(ctx, queryIDChan)
+	conn, _ := db.Conn(ctx)
+
+	dest := make([]driver.Value, 2)
+
+	goRoutineChan := make(chan string)
+	go func(grCh chan string, qIDch chan string) {
+		queryID := <- queryIDChan
+		grCh <- queryID
+	}(goRoutineChan, queryIDChan)
+
+	err = conn.Raw(func(x interface{}) error {
+		stmt, err := x.(driver.ConnPrepareContext).PrepareContext(ctx, fmt.Sprintf("SELECT SEQ8(), RANDSTR(1000, RANDOM()) FROM TABLE(GENERATOR(ROWCOUNT=>%v))", numrows))
+		if err != nil {
+			return err
+		}
+		rows, err := stmt.(driver.StmtQueryContext).QueryContext(ctx, nil)
+		if err != nil {
+			return err
+		}
+		defer rows.Close()
+
+		cnt := 0
+		for {
+			err = rows.Next(dest)
+			if err != nil {
+				if err == io.EOF {
+					break
+				} else {
+					t.Fatalf("unexpected error: %v", err)
+				}
+			}
+			cnt++
+		}
+		queryID := <- goRoutineChan
+		if queryID == "" {
+			t.Fatal("expected a nonempty query ID")
+		}
+		if cnt != numrows {
+			t.Errorf("number of rows didn't match. expected: %v, got: %v", numrows, cnt)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("failed to execute query. err: %v", err)
+	}
+}