Repository: snowflakedb/gosnowflake
Pull request #342 (2020-10-19 21:19:05)

SNOW-206590 allow specifying the snowflake request id via context
### Description
Allow specific request ID to be set via the context key `SNOWFLAKE_REQUEST_ID`.
If the key is not specified or value is empty, fallback to existing behavior of generating UUID.

### 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 68c7a49..cbe4b97 100644
--- a/auth.go
+++ b/auth.go
@@ -195,7 +195,7 @@ func postAuth(
 	body []byte,
 	timeout time.Duration) (
 	data *authResponse, err error) {
-	params.Add(requestIDKey, uuid.New().String())
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
 	params.Add(requestGUIDKey, uuid.New().String())
 
 	fullURL := sr.getFullURL(loginRequestPath, params)
diff --git a/authokta.go b/authokta.go
index c937507..2f70570 100644
--- a/authokta.go
+++ b/authokta.go
@@ -13,8 +13,6 @@ import (
 	"net/url"
 	"strconv"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 type authOKTARequest struct {
@@ -208,7 +206,7 @@ func postAuthSAML(
 	data *authResponse, err error) {
 
 	params := &url.Values{}
-	params.Add(requestIDKey, uuid.New().String())
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
 	fullURL := sr.getFullURL(authenticatorRequestPath, params)
 
 	glog.V(2).Infof("fullURL: %v", fullURL)
diff --git a/connection.go b/connection.go
index dc7989a..7b50adc 100644
--- a/connection.go
+++ b/connection.go
@@ -127,8 +127,8 @@ func (sc *snowflakeConn) exec(
 
 	var data *execResponse
 
-	requestID := uuid.New()
-	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers, jsonBody, sc.rest.RequestTimeout, &requestID)
+	requestID := getOrGenerateRequestIDFromContext(ctx)
+	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers, jsonBody, sc.rest.RequestTimeout, requestID)
 	if err != nil {
 		return data, err
 	}
@@ -525,7 +525,7 @@ func (sc *snowflakeConn) getQueryResult(ctx context.Context, resultPath string)
 		headers["X-Snowflake-Service"] = *serviceName
 	}
 	param := make(url.Values)
-	param.Add(requestIDKey, uuid.New().String())
+	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
 	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 a94abd3..327e745 100644
--- a/doc.go
+++ b/doc.go
@@ -195,6 +195,14 @@ Note: The logger may be changed in the future for better logging. Currently if
 the applications use the same parameters as glog, you cannot collect both
 application and driver logs at the same time.
 
+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-...")
+	rows, err := db.QueryContext(ctxWithID, query)
+
 Canceling Query by CtrlC
 
 From 0.5.0, a signal handling responsibility has moved to the applications. If you want to cancel a
diff --git a/restful.go b/restful.go
index 1f9d4e4..784e6c3 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, *uuid.UUID) (*execResponse, error)
-	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, *uuid.UUID) (*execResponse, error)
+	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)
 	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, *uuid.UUID, time.Duration) error
+	FuncCancelQuery     func(context.Context, *snowflakeRestful, string, 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 *uuid.UUID) (
+	requestID string) (
 	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 *uuid.UUID) (
+	requestID string) (
 	data *execResponse, err error) {
 	glog.V(2).Infof("params: %v", params)
-	params.Add(requestIDKey, requestID.String())
+	params.Add(requestIDKey, requestID)
 	params.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
 	params.Add(requestGUIDKey, uuid.New().String())
 	if sr.Token != "" {
@@ -258,7 +258,7 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 	glog.V(2).Info("close session")
 	params := &url.Values{}
 	params.Add("delete", "true")
-	params.Add(requestIDKey, uuid.New().String())
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
 	params.Add(requestGUIDKey, uuid.New().String())
 	fullURL := sr.getFullURL(sessionRequestPath, params)
 
@@ -313,7 +313,7 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time.Duration) error {
 	glog.V(2).Info("start renew session")
 	params := &url.Values{}
-	params.Add(requestIDKey, uuid.New().String())
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
 	params.Add(requestGUIDKey, uuid.New().String())
 	fullURL := sr.getFullURL(tokenRequestPath, params)
 
@@ -377,10 +377,10 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time
 	}
 }
 
-func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID *uuid.UUID, timeout time.Duration) error {
+func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID string, timeout time.Duration) error {
 	glog.V(2).Info("cancel query")
 	params := &url.Values{}
-	params.Add(requestIDKey, uuid.New().String())
+	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx))
 	params.Add(requestGUIDKey, uuid.New().String())
 
 	fullURL := sr.getFullURL(abortRequestPath, params)
@@ -392,7 +392,7 @@ func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID *uuid.UUID
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, sr.Token)
 
 	req := make(map[string]string)
-	req[requestIDKey] = requestID.String()
+	req[requestIDKey] = requestID
 
 	reqByte, err := json.Marshal(req)
 	if err != nil {
diff --git a/util.go b/util.go
index e4948f6..c45d8ae 100644
--- a/util.go
+++ b/util.go
@@ -3,10 +3,31 @@
 package gosnowflake
 
 import (
+	"context"
 	"database/sql/driver"
+	"github.com/google/uuid"
 	"time"
 )
 
+type contextKey string
+
+// SnowflakeRequestIDKey is optional context key to specify request id
+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 {
+	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 != "" {
+		return requestID
+	}
+	return uuid.New().String()
+}
+
 // integer min
 func intMin(a, b int) int {
 	if a < b {


Test patch:
diff --git a/connection_test.go b/connection_test.go
index 1fb3014..6713f19 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -2,7 +2,6 @@ package gosnowflake
 
 import (
 	"context"
-	"github.com/google/uuid"
 	"net/url"
 	"testing"
 	"time"
@@ -14,7 +13,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, _ *uuid.UUID) (*execResponse, error) {
+func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ string) (*execResponse, error) {
 	var serviceName string
 	if serviceHeader, ok := headers["X-Snowflake-Service"]; ok {
 		serviceName = serviceHeader + serviceNameAppend
@@ -33,6 +32,67 @@ func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values, header
 	}, nil
 }
 
+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) {
+		// ensure the same requestID from context is used
+		if len(requestID) == 0 {
+			t.Fatal("requestID is empty")
+		}
+		dd := &execResponseData{}
+		return &execResponse{
+			Data:    *dd,
+			Message: "",
+			Code:    "0",
+			Success: true,
+		}, nil
+	}
+
+	sr := &snowflakeRestful{
+		FuncPostQuery: postQueryMock,
+	}
+
+	sc := &snowflakeConn{
+		cfg:  &Config{Params: map[string]*string{}},
+		rest: sr,
+	}
+	_, err := sc.exec(ctx, "", false, false, nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+}
+
+func TestExecWithSpecificRequestID(t *testing.T) {
+	origRequestID := "specific-snowflake-request-id"
+	ctx := WithRequestID(context.Background(), origRequestID)
+	postQueryMock := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID string) (*execResponse, error) {
+		// ensure the same requestID from context is used
+		if requestID != origRequestID {
+			t.Fatal("requestID doesn't match")
+		}
+		dd := &execResponseData{}
+		return &execResponse{
+			Data:    *dd,
+			Message: "",
+			Code:    "0",
+			Success: true,
+		}, nil
+	}
+
+	sr := &snowflakeRestful{
+		FuncPostQuery: postQueryMock,
+	}
+
+	sc := &snowflakeConn{
+		cfg:  &Config{Params: map[string]*string{}},
+		rest: sr,
+	}
+	_, err := sc.exec(ctx, "", false, false, nil)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+}
+
 // TestServiceName tests two things:
 // 1. request header would contain X-Snowflake-Service if the cfg parameters contains SERVICE_NAME
 // 2. SERVICE_NAME would be update by response payload
diff --git a/restful_test.go b/restful_test.go
index 0cb8b81..b4acbeb 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -93,21 +93,21 @@ func TestUnitPostQueryHelperError(t *testing.T) {
 		FuncPost: postTestError,
 	}
 	var err error
-	var requestID uuid.UUID
-	requestID = uuid.New()
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, &requestID)
+	var requestID string
+	requestID = uuid.New().String()
+	_, 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()
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, &requestID)
+	requestID = uuid.New().String()
+	_, 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()
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, &requestID)
+	requestID = uuid.New().String()
+	_, 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,10 +123,10 @@ func renewSessionTestError(_ context.Context, _ *snowflakeRestful, _ time.Durati
 
 func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 	var err error
-	orgRequestID := uuid.New()
-	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID *uuid.UUID) (*execResponse, error) {
+	origRequestID := uuid.New().String()
+	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID string) (*execResponse, error) {
 		// ensure the same requestID is used after the session token is renewed.
-		if requestID.String() != orgRequestID.String() {
+		if requestID != origRequestID {
 			t.Fatal("requestID doesn't match")
 		}
 		dd := &execResponseData{}
@@ -144,12 +144,12 @@ func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 		FuncRenewSession: renewSessionTest,
 	}
 
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, &orgRequestID)
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, origRequestID)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	sr.FuncRenewSession = renewSessionTestError
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, &orgRequestID)
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, origRequestID)
 	if err == nil {
 		t.Fatal("should have failed to renew session")
 	}
@@ -211,27 +211,23 @@ func TestUnitCancelQuery(t *testing.T) {
 	sr := &snowflakeRestful{
 		FuncPost: postTestAfterRenew,
 	}
-	var requestID uuid.UUID
-	requestID = uuid.New()
-	err := cancelQuery(context.Background(), sr, &requestID, time.Second)
+	ctx := context.Background()
+	err := cancelQuery(ctx, sr, getOrGenerateRequestIDFromContext(ctx), time.Second)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	sr.FuncPost = postTestError
-	requestID = uuid.New()
-	err = cancelQuery(context.Background(), sr, &requestID, time.Second)
+	err = cancelQuery(ctx, sr, getOrGenerateRequestIDFromContext(ctx), time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
-	requestID = uuid.New()
-	err = cancelQuery(context.Background(), sr, &requestID, time.Second)
+	err = cancelQuery(context.Background(), sr, getOrGenerateRequestIDFromContext(ctx), time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
-	requestID = uuid.New()
-	err = cancelQuery(context.Background(), sr, &requestID, time.Second)
+	err = cancelQuery(context.Background(), sr, getOrGenerateRequestIDFromContext(ctx), time.Second)
 	if err == nil {
 		t.Fatal("should have failed to close session")
 	}
diff --git a/util_test.go b/util_test.go
index be46c93..1186a53 100644
--- a/util_test.go
+++ b/util_test.go
@@ -3,6 +3,7 @@
 package gosnowflake
 
 import (
+	"context"
 	"database/sql/driver"
 	"testing"
 	"time"
@@ -14,6 +15,23 @@ type tcIntMinMax struct {
 	out int
 }
 
+func TestGetRequestIDFromContext(t *testing.T) {
+	expectedRequestID := "snowflake-request-id"
+	ctx := WithRequestID(context.Background(), expectedRequestID)
+	requestID := getOrGenerateRequestIDFromContext(ctx)
+	if requestID != expectedRequestID {
+		t.Errorf("unexpected request id")
+	}
+}
+
+func TestGenerateRequestID(t *testing.T) {
+	firstRequestID := getOrGenerateRequestIDFromContext(context.Background())
+	otherRequestID := getOrGenerateRequestIDFromContext(context.Background())
+	if firstRequestID == otherRequestID {
+		t.Errorf("request id should not be the same")
+	}
+}
+
 func TestIntMin(t *testing.T) {
 	testcases := []tcIntMinMax{
 		{1, 3, 1},