Repository: snowflakedb/gosnowflake
Pull request #377 (2021-02-03 01:01:56)

SNOW-274063 QueryRowContext hangs on bad request

### Description
GS rest endpoint requires UUID in the V1 endpoint.
It will return a bad request if the request ID is not UUID.

On HTTP400, the http retry logic still kicks in and end up
causing a hang due to retrying a bad query.

To fix this, I'm proposing the following changes:
1. Make request ID always UUID type, and check against UUID.nil
2. postQueryHelper should not retry on 4xx error, any 4xx error should be bubbled up

Testing
Added new unit test

### 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
- [x] Extended the README / documentation, if necessary



Patch:
diff --git a/auth.go b/auth.go
index 39d9bfd..ae65b83 100644
--- a/auth.go
+++ b/auth.go
@@ -4,6 +4,9 @@ package gosnowflake
 
 import (
 	"context"
+	"crypto/sha256"
+	"crypto/x509"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -16,10 +19,6 @@ import (
 
 	"github.com/form3tech-oss/jwt-go"
 	"github.com/google/uuid"
-
-	"crypto/sha256"
-	"crypto/x509"
-	"encoding/base64"
 )
 
 const (
@@ -195,7 +194,7 @@ func postAuth(
 	body []byte,
 	timeout time.Duration) (
 	data *authResponse, err error) {
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	params.Add(requestGUIDKey, uuid.New().String())
 
 	fullURL := sr.getFullURL(loginRequestPath, params)
diff --git a/authokta.go b/authokta.go
index bd59239..55a7735 100644
--- a/authokta.go
+++ b/authokta.go
@@ -205,7 +205,7 @@ func postAuthSAML(
 	data *authResponse, err error) {
 
 	params := &url.Values{}
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	fullURL := sr.getFullURL(authenticatorRequestPath, params)
 
 	logger.Infof("fullURL: %v", fullURL)
diff --git a/connection.go b/connection.go
index 302e528..84c70b6 100644
--- a/connection.go
+++ b/connection.go
@@ -8,13 +8,14 @@ import (
 	"database/sql/driver"
 	"encoding/json"
 	"fmt"
-	"github.com/google/uuid"
 	"net/url"
 	"reflect"
 	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
+
+	"github.com/google/uuid"
 )
 
 const (
@@ -560,7 +561,7 @@ func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string)
 		headers["X-Snowflake-Service"] = *serviceName
 	}
 	param := make(url.Values)
-	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
+	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	param.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
 	param.Add(requestGUIDKey, uuid.New().String())
 	if sc.rest.Token != "" {
diff --git a/doc.go b/doc.go
index 16292be..3eae79b 100644
--- a/doc.go
+++ b/doc.go
@@ -158,7 +158,8 @@ Query request ID
 Specific query request ID can be set in the context and will be passed through
 in place of the default randomized request ID. For example:
 
-	ctxWithID := WithRequestID(ctx, "app-request-id-...")
+	requestID := uuid.MustParse("6ba7b812-9dad-11d1-80b4-00c04fd430c8")
+	ctxWithID := WithRequestID(ctx, requestID)
 	rows, err := db.QueryContext(ctxWithID, query)
 
 Canceling Query by CtrlC
diff --git a/restful.go b/restful.go
index aac54a1..8c03914 100644
--- a/restful.go
+++ b/restful.go
@@ -56,14 +56,14 @@ type snowflakeRestful struct {
 	HeartBeat   *heartbeat
 
 	Connection          *snowflakeConn
-	FuncPostQuery       func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, string) (*execResponse, error)
-	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, string) (*execResponse, error)
+	FuncPostQuery       func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID) (*execResponse, error)
+	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID) (*execResponse, error)
 	FuncPost            func(context.Context, *snowflakeRestful, *url.URL, map[string]string, []byte, time.Duration, bool) (*http.Response, error)
 	FuncGet             func(context.Context, *snowflakeRestful, *url.URL, map[string]string, time.Duration) (*http.Response, error)
 	FuncRenewSession    func(context.Context, *snowflakeRestful, time.Duration) error
 	FuncPostAuth        func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration) (*authResponse, error)
 	FuncCloseSession    func(context.Context, *snowflakeRestful, time.Duration) error
-	FuncCancelQuery     func(context.Context, *snowflakeRestful, string, time.Duration) error
+	FuncCancelQuery     func(context.Context, *snowflakeRestful, uuid.UUID, time.Duration) error
 
 	FuncPostAuthSAML func(context.Context, *snowflakeRestful, map[string]string, []byte, time.Duration) (*authResponse, error)
 	FuncPostAuthOKTA func(context.Context, *snowflakeRestful, map[string]string, []byte, string, time.Duration) (*authOKTAResponse, error)
@@ -142,7 +142,7 @@ func postRestfulQuery(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	requestID string) (
+	requestID uuid.UUID) (
 	data *execResponse, err error) {
 
 	data, err = sr.FuncPostQueryHelper(ctx, sr, params, headers, body, timeout, requestID)
@@ -166,10 +166,10 @@ func postRestfulQueryHelper(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	requestID string) (
+	requestID uuid.UUID) (
 	data *execResponse, err error) {
 	logger.Infof("params: %v", params)
-	params.Add(requestIDKey, requestID)
+	params.Add(requestIDKey, requestID.String())
 	params.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
 	params.Add(requestGUIDKey, uuid.New().String())
 	if sr.Token != "" {
@@ -178,7 +178,7 @@ func postRestfulQueryHelper(
 
 	var resp *http.Response
 	fullURL := sr.getFullURL(queryRequestPath, params)
-	resp, err = sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, false)
+	resp, err = sr.FuncPost(ctx, sr, fullURL, headers, body, timeout, true)
 
 	if err != nil {
 		return nil, err
@@ -282,7 +282,7 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 	logger.WithContext(ctx).Info("close session")
 	params := &url.Values{}
 	params.Add("delete", "true")
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	params.Add(requestGUIDKey, uuid.New().String())
 	fullURL := sr.getFullURL(sessionRequestPath, params)
 
@@ -334,7 +334,7 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time.Duration) error {
 	logger.WithContext(ctx).Info("start renew session")
 	params := &url.Values{}
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	params.Add(requestGUIDKey, uuid.New().String())
 	fullURL := sr.getFullURL(tokenRequestPath, params)
 
@@ -395,10 +395,10 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time
 	}
 }
 
-func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID string, timeout time.Duration) error {
+func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID uuid.UUID, timeout time.Duration) error {
 	logger.WithContext(ctx).Info("cancel query")
 	params := &url.Values{}
-	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	params.Add(requestGUIDKey, uuid.New().String())
 
 	fullURL := sr.getFullURL(abortRequestPath, params)
@@ -410,7 +410,7 @@ func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID string, ti
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
 
 	req := make(map[string]string)
-	req[requestIDKey] = requestID
+	req[requestIDKey] = requestID.String()
 
 	reqByte, err := json.Marshal(req)
 	if err != nil {
diff --git a/util.go b/util.go
index c45d8ae..ee8593d 100644
--- a/util.go
+++ b/util.go
@@ -5,8 +5,9 @@ package gosnowflake
 import (
 	"context"
 	"database/sql/driver"
-	"github.com/google/uuid"
 	"time"
+
+	"github.com/google/uuid"
 )
 
 type contextKey string
@@ -15,17 +16,17 @@ type contextKey string
 const SnowflakeRequestIDKey contextKey = "SNOWFLAKE_REQUEST_ID"
 
 // WithRequestID returns a new context with the specified snowflake request id
-func WithRequestID(ctx context.Context, requestID string) context.Context {
+func WithRequestID(ctx context.Context, requestID uuid.UUID) context.Context {
 	return context.WithValue(ctx, SnowflakeRequestIDKey, requestID)
 }
 
 // Get the request ID from the context if specified, otherwise generate one
-func getOrGenerateRequestIDFromContext(ctx context.Context) string {
-	requestID, ok := ctx.Value(SnowflakeRequestIDKey).(string)
-	if ok && requestID != "" {
+func getOrGenerateRequestIDFromContext(ctx context.Context) uuid.UUID {
+	requestID, ok := ctx.Value(SnowflakeRequestIDKey).(uuid.UUID)
+	if ok && requestID != uuid.Nil {
 		return requestID
 	}
-	return uuid.New().String()
+	return uuid.New()
 }
 
 // integer min


Test patch:
diff --git a/auth_test.go b/auth_test.go
index da309cf..b2729b5 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -9,10 +9,11 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"github.com/form3tech-oss/jwt-go"
 	"net/url"
 	"testing"
 	"time"
+
+	"github.com/form3tech-oss/jwt-go"
 )
 
 func TestUnitPostAuth(t *testing.T) {
diff --git a/connection_test.go b/connection_test.go
index 6713f19..a9fc7d5 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -5,6 +5,8 @@ import (
 	"net/url"
 	"testing"
 	"time"
+
+	"github.com/google/uuid"
 )
 
 const serviceNameStub = "SV"
@@ -13,7 +15,7 @@ const serviceNameAppend = "a"
 // postQueryMock would generate a response based on the X-Snowflake-Service header, to generate a response
 // with the SERVICE_NAME field appending a character at the end of the header
 // This way it could test both the send and receive logic
-func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ string) (*execResponse, error) {
+func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ uuid.UUID) (*execResponse, error) {
 	var serviceName string
 	if serviceHeader, ok := headers["X-Snowflake-Service"]; ok {
 		serviceName = serviceHeader + serviceNameAppend
@@ -33,8 +35,8 @@ func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, header
 }
 
 func TestExecWithEmptyRequestID(t *testing.T) {
-	ctx := WithRequestID(context.Background(), "")
-	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID string) (*execResponse, error) {
+	ctx := WithRequestID(context.Background(), uuid.Nil)
+	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID) (*execResponse, error) {
 		// ensure the same requestID from context is used
 		if len(requestID) == 0 {
 			t.Fatal("requestID is empty")
@@ -63,9 +65,9 @@ func TestExecWithEmptyRequestID(t *testing.T) {
 }
 
 func TestExecWithSpecificRequestID(t *testing.T) {
-	origRequestID := "specific-snowflake-request-id"
+	origRequestID := uuid.New()
 	ctx := WithRequestID(context.Background(), origRequestID)
-	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID string) (*execResponse, error) {
+	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID) (*execResponse, error) {
 		// ensure the same requestID from context is used
 		if requestID != origRequestID {
 			t.Fatal("requestID doesn't match")
diff --git a/driver_test.go b/driver_test.go
index a3093ce..52eb65b 100644
--- a/driver_test.go
+++ b/driver_test.go
@@ -19,6 +19,8 @@ import (
 	"syscall"
 	"testing"
 	"time"
+
+	"github.com/google/uuid"
 )
 
 var (
@@ -414,6 +416,19 @@ func TestEmptyQuery(t *testing.T) {
 	})
 }
 
+func TestEmptyQueryWithRequestID(t *testing.T) {
+	runTests(t, dsn, func(dbt *DBTest) {
+		query := "select 1"
+		ctx := WithRequestID(context.Background(), uuid.New())
+		rows := dbt.db.QueryRowContext(ctx, query)
+		var v1 interface{}
+		err := rows.Scan(&v1)
+		if err != nil {
+			dbt.Errorf("should not have failed with valid request id. err: %v", err)
+		}
+	})
+}
+
 func TestCRUD(t *testing.T) {
 	runTests(t, dsn, func(dbt *DBTest) {
 		// Create Table
diff --git a/restful_test.go b/restful_test.go
index 5010af9..5b7da90 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -5,11 +5,12 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
-	"github.com/google/uuid"
 	"net/http"
 	"net/url"
 	"testing"
 	"time"
+
+	"github.com/google/uuid"
 )
 
 func postTestError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
@@ -93,20 +94,20 @@ func TestUnitPostQueryHelperError(t *testing.T) {
 		FuncPost: postTestError,
 	}
 	var err error
-	var requestID string
-	requestID = uuid.New().String()
+	var requestID uuid.UUID
+	requestID = uuid.New()
 	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID)
 	if err == nil {
 		t.Fatalf("should have failed to post")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
-	requestID = uuid.New().String()
+	requestID = uuid.New()
 	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID)
 	if err == nil {
 		t.Fatalf("should have failed to post")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
-	requestID = uuid.New().String()
+	requestID = uuid.New()
 	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID)
 	if err == nil {
 		t.Fatalf("should have failed to post")
@@ -123,8 +124,8 @@ func renewSessionTestError(_ context.Context, _ *snowflakeRestful, _ time.Durati
 
 func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 	var err error
-	origRequestID := uuid.New().String()
-	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID string) (*execResponse, error) {
+	origRequestID := uuid.New()
+	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID) (*execResponse, error) {
 		// ensure the same requestID is used after the session token is renewed.
 		if requestID != origRequestID {
 			t.Fatal("requestID doesn't match")
diff --git a/util_test.go b/util_test.go
index 1186a53..59aa07b 100644
--- a/util_test.go
+++ b/util_test.go
@@ -7,6 +7,8 @@ import (
 	"database/sql/driver"
 	"testing"
 	"time"
+
+	"github.com/google/uuid"
 )
 
 type tcIntMinMax struct {
@@ -16,11 +18,16 @@ type tcIntMinMax struct {
 }
 
 func TestGetRequestIDFromContext(t *testing.T) {
-	expectedRequestID := "snowflake-request-id"
+	expectedRequestID := uuid.New()
 	ctx := WithRequestID(context.Background(), expectedRequestID)
 	requestID := getOrGenerateRequestIDFromContext(ctx)
 	if requestID != expectedRequestID {
-		t.Errorf("unexpected request id")
+		t.Errorf("unexpected request id: %v, expected: %v", requestID, expectedRequestID)
+	}
+	ctx = WithRequestID(context.Background(), uuid.Nil)
+	requestID = getOrGenerateRequestIDFromContext(ctx)
+	if requestID == uuid.Nil {
+		t.Errorf("unexpected request id, should not be nil")
 	}
 }