Repository: snowflakedb/gosnowflake
Pull request #541 (2022-01-30 04:33:47)

SNOW-440362 Implement UUIDv4
### Description
Implement UUID https://github.com/snowflakedb/gosnowflake/issues/173

### Checklist
- [ ] Code compiles correctly
- [ ] Run ``make fmt`` to fix inconsistent formats
- [ ] Run ``make lint`` to get lint errors and fix all of them
- [ ] Created tests which fail without the change (if possible)
- [ ] All tests passing
- [ ] Extended the README / documentation, if necessary



Patch:
diff --git a/auth.go b/auth.go
index f2534ec..1ca6614 100644
--- a/auth.go
+++ b/auth.go
@@ -18,7 +18,6 @@ import (
 	"time"
 
 	"github.com/form3tech-oss/jwt-go"
-	"github.com/google/uuid"
 )
 
 const (
@@ -199,7 +198,7 @@ func postAuth(
 	timeout time.Duration) (
 	data *authResponse, err error) {
 	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	params.Add(requestGUIDKey, uuid.New().String())
+	params.Add(requestGUIDKey, newUUID().String())
 
 	fullURL := sr.getFullURL(loginRequestPath, params)
 	logger.Infof("full URL: %v", fullURL)
diff --git a/bind_uploader.go b/bind_uploader.go
index b56ed97..e348907 100644
--- a/bind_uploader.go
+++ b/bind_uploader.go
@@ -10,8 +10,6 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
-
-	"github.com/google/uuid"
 )
 
 const (
@@ -173,7 +171,7 @@ func (sc *snowflakeConn) processBindings(
 	ctx context.Context,
 	bindings []driver.NamedValue,
 	describeOnly bool,
-	requestID uuid.UUID,
+	requestID uuid,
 	req *execRequest) error {
 	arrayBindThreshold := sc.getArrayBindStageThreshold()
 	numBinds := arrayBindValueCount(bindings)
diff --git a/go.mod b/go.mod
index b0eafe0..c49c028 100644
--- a/go.mod
+++ b/go.mod
@@ -12,7 +12,7 @@ require (
 	github.com/aws/smithy-go v1.9.0
 	github.com/form3tech-oss/jwt-go v3.2.5+incompatible
 	github.com/gabriel-vasile/mimetype v1.4.0
-	github.com/google/uuid v1.3.0
+	github.com/google/uuid v1.3.0 // indirect
 	github.com/klauspost/compress v1.13.6 // indirect
 	github.com/pierrec/lz4/v4 v4.1.11 // indirect
 	github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8
diff --git a/heartbeat.go b/heartbeat.go
index 71b8a7b..e475a4f 100644
--- a/heartbeat.go
+++ b/heartbeat.go
@@ -10,8 +10,6 @@ import (
 	"net/http"
 	"net/url"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 const (
@@ -56,8 +54,8 @@ func (hc *heartbeat) stop() {
 func (hc *heartbeat) heartbeatMain() error {
 	logger.Info("Heartbeating!")
 	params := &url.Values{}
-	params.Add(requestIDKey, uuid.New().String())
-	params.Add(requestGUIDKey, uuid.New().String())
+	params.Add(requestIDKey, newUUID().String())
+	params.Add(requestGUIDKey, newUUID().String())
 	headers := getHeaders()
 	token, _, _ := hc.restful.TokenAccessor.GetTokens()
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
diff --git a/monitoring.go b/monitoring.go
index 487f96b..f821596 100644
--- a/monitoring.go
+++ b/monitoring.go
@@ -10,8 +10,6 @@ import (
 	"net/url"
 	"strconv"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 const urlQueriesResultFmt = "/queries/%s/result"
@@ -135,7 +133,7 @@ func (sc *snowflakeConn) checkQueryStatus(
 	*retStatus, error) {
 	headers := make(map[string]string)
 	param := make(url.Values)
-	param.Add(requestGUIDKey, uuid.New().String())
+	param.Add(requestGUIDKey, newUUID().String())
 	if tok, _, _ := sc.rest.TokenAccessor.GetTokens(); tok != "" {
 		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, tok)
 	}
@@ -208,7 +206,7 @@ func (sc *snowflakeConn) getQueryResultResp(
 	param := make(url.Values)
 	param.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
 	param.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
-	param.Add(requestGUIDKey, uuid.New().String())
+	param.Add(requestGUIDKey, newUUID().String())
 	token, _, _ := sc.rest.TokenAccessor.GetTokens()
 	if token != "" {
 		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
diff --git a/restful.go b/restful.go
index f43ce9f..478d232 100644
--- a/restful.go
+++ b/restful.go
@@ -11,8 +11,6 @@ import (
 	"net/url"
 	"strconv"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 // HTTP headers
@@ -61,14 +59,14 @@ type snowflakeRestful struct {
 
 	Connection *snowflakeConn
 
-	FuncPostQuery       func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID, *Config) (*execResponse, error)
-	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid.UUID, *Config) (*execResponse, error)
+	FuncPostQuery       func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid, *Config) (*execResponse, error)
+	FuncPostQueryHelper func(context.Context, *snowflakeRestful, *url.Values, map[string]string, []byte, time.Duration, uuid, *Config) (*execResponse, error)
 	FuncPost            funcPostType
 	FuncGet             funcGetType
 	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, 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)
@@ -170,7 +168,7 @@ func postRestfulQuery(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	requestID uuid.UUID,
+	requestID uuid,
 	cfg *Config) (
 	data *execResponse, err error) {
 
@@ -194,13 +192,13 @@ func postRestfulQueryHelper(
 	headers map[string]string,
 	body []byte,
 	timeout time.Duration,
-	requestID uuid.UUID,
+	requestID uuid,
 	cfg *Config) (
 	data *execResponse, err error) {
 	logger.Infof("params: %v", params)
 	params.Add(requestIDKey, requestID.String())
 	params.Add("clientStartTime", strconv.FormatInt(time.Now().Unix(), 10))
-	params.Add(requestGUIDKey, uuid.New().String())
+	params.Add(requestGUIDKey, newUUID().String())
 	token, _, _ := sr.TokenAccessor.GetTokens()
 	if token != "" {
 		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
@@ -293,7 +291,7 @@ func closeSession(ctx context.Context, sr *snowflakeRestful, timeout time.Durati
 	params := &url.Values{}
 	params.Add("delete", "true")
 	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	params.Add(requestGUIDKey, uuid.New().String())
+	params.Add(requestGUIDKey, newUUID().String())
 	fullURL := sr.getFullURL(sessionRequestPath, params)
 
 	headers := getHeaders()
@@ -342,7 +340,7 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time
 	logger.WithContext(ctx).Info("start renew session")
 	params := &url.Values{}
 	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	params.Add(requestGUIDKey, uuid.New().String())
+	params.Add(requestGUIDKey, newUUID().String())
 	fullURL := sr.getFullURL(tokenRequestPath, params)
 
 	token, masterToken, _ := sr.TokenAccessor.GetTokens()
@@ -411,11 +409,11 @@ func getCancelRetry(ctx context.Context) int {
 	return cnt
 }
 
-func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID uuid.UUID, timeout time.Duration) error {
+func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID uuid, timeout time.Duration) error {
 	logger.WithContext(ctx).Info("cancel query")
 	params := &url.Values{}
 	params.Add(requestIDKey, getOrGenerateRequestIDFromContext(ctx).String())
-	params.Add(requestGUIDKey, uuid.New().String())
+	params.Add(requestGUIDKey, newUUID().String())
 
 	fullURL := sr.getFullURL(abortRequestPath, params)
 
diff --git a/retry.go b/retry.go
index e44d6f5..9d2b7a4 100644
--- a/retry.go
+++ b/retry.go
@@ -15,8 +15,6 @@ import (
 	"strings"
 	"sync"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 var random *rand.Rand
@@ -81,7 +79,7 @@ generated uuid
 */
 func (replacer *requestGUIDReplace) replace() *url.URL {
 	replacer.urlValues.Del(requestGUIDKey)
-	replacer.urlValues.Add(requestGUIDKey, uuid.New().String())
+	replacer.urlValues.Add(requestGUIDKey, newUUID().String())
 	replacer.urlPtr.RawQuery = replacer.urlValues.Encode()
 	return replacer.urlPtr
 }
diff --git a/util.go b/util.go
index eb3cdf0..33733bf 100644
--- a/util.go
+++ b/util.go
@@ -10,8 +10,6 @@ import (
 	"strings"
 	"sync"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 type contextKey string
@@ -49,7 +47,7 @@ func WithQueryIDChan(ctx context.Context, c chan<- string) context.Context {
 }
 
 // WithRequestID returns a new context with the specified snowflake request id
-func WithRequestID(ctx context.Context, requestID uuid.UUID) context.Context {
+func WithRequestID(ctx context.Context, requestID uuid) context.Context {
 	return context.WithValue(ctx, snowflakeRequestIDKey, requestID)
 }
 
@@ -86,12 +84,12 @@ func WithHigherPrecision(ctx context.Context) context.Context {
 }
 
 // Get the request ID from the context if specified, otherwise generate one
-func getOrGenerateRequestIDFromContext(ctx context.Context) uuid.UUID {
-	requestID, ok := ctx.Value(snowflakeRequestIDKey).(uuid.UUID)
-	if ok && requestID != uuid.Nil {
+func getOrGenerateRequestIDFromContext(ctx context.Context) uuid {
+	requestID, ok := ctx.Value(snowflakeRequestIDKey).(uuid)
+	if ok && requestID != nilUUID {
 		return requestID
 	}
-	return uuid.New()
+	return newUUID()
 }
 
 // integer min
diff --git a/uuid.go b/uuid.go
new file mode 100644
index 0000000..601175e
--- /dev/null
+++ b/uuid.go
@@ -0,0 +1,28 @@
+// Copyright (c) 2022 Snowflake Computing Inc. All rights reserved.
+
+package gosnowflake
+
+import (
+	"crypto/rand"
+	"fmt"
+)
+
+const rfc4122 = 0x40
+
+type uuid [16]byte
+
+var nilUUID uuid
+
+func newUUID() uuid {
+	var u uuid
+	rand.Read(u[:])
+	u[8] = (u[8] | rfc4122) & 0x7F
+
+	var version byte = 4
+	u[6] = (u[6] & 0xF) | (version << 4)
+	return u
+}
+
+func (u uuid) String() string {
+	return fmt.Sprintf("%x-%x-%x-%x-%x", u[0:4], u[4:6], u[6:8], u[8:10], u[10:])
+}


Test patch:
diff --git a/connection_test.go b/connection_test.go
index 010a004..b895917 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -16,8 +16,6 @@ import (
 	"sync"
 	"testing"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 const (
@@ -48,7 +46,7 @@ func TestInvalidConnection(t *testing.T) {
 // 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,
+	headers map[string]string, _ []byte, _ time.Duration, _ uuid,
 	_ *Config) (*execResponse, error) {
 	var serviceName string
 	if serviceHeader, ok := headers[httpHeaderServiceName]; ok {
@@ -69,10 +67,10 @@ func postQueryMock(_ context.Context, _ *snowflakeRestful, _ *url.Values,
 }
 
 func TestExecWithEmptyRequestID(t *testing.T) {
-	ctx := WithRequestID(context.Background(), uuid.Nil)
+	ctx := WithRequestID(context.Background(), nilUUID)
 	postQueryMock := func(_ context.Context, _ *snowflakeRestful,
 		_ *url.Values, _ map[string]string, _ []byte, _ time.Duration,
-		requestID uuid.UUID, _ *Config) (*execResponse, error) {
+		requestID uuid, _ *Config) (*execResponse, error) {
 		// ensure the same requestID from context is used
 		if len(requestID) == 0 {
 			t.Fatal("requestID is empty")
@@ -139,11 +137,11 @@ func TestGetQueryResultUsesTokenFromTokenAccessor(t *testing.T) {
 }
 
 func TestExecWithSpecificRequestID(t *testing.T) {
-	origRequestID := uuid.New()
+	origRequestID := newUUID()
 	ctx := WithRequestID(context.Background(), origRequestID)
 	postQueryMock := func(_ context.Context, _ *snowflakeRestful,
 		_ *url.Values, _ map[string]string, _ []byte, _ time.Duration,
-		requestID uuid.UUID, _ *Config) (*execResponse, error) {
+		requestID uuid, _ *Config) (*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 28414a0..cade0e4 100644
--- a/driver_test.go
+++ b/driver_test.go
@@ -16,8 +16,6 @@ import (
 	"syscall"
 	"testing"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 var (
@@ -445,7 +443,7 @@ 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())
+		ctx := WithRequestID(context.Background(), newUUID())
 		rows := dbt.db.QueryRowContext(ctx, query)
 		var v1 interface{}
 		if err := rows.Scan(&v1); err != nil {
diff --git a/restful_test.go b/restful_test.go
index db95d0b..2eedfc0 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -13,8 +13,6 @@ import (
 	"sync/atomic"
 	"testing"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 func postTestError(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
@@ -111,7 +109,7 @@ func postTestAfterRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ ma
 	}, nil
 }
 
-func cancelTestRetry(ctx context.Context, sr *snowflakeRestful, requestID uuid.UUID, timeout time.Duration) error {
+func cancelTestRetry(ctx context.Context, sr *snowflakeRestful, requestID uuid, timeout time.Duration) error {
 	ctxRetry := getCancelRetry(ctx)
 	u := url.URL{}
 	reqByte, err := json.Marshal(make(map[string]string))
@@ -144,20 +142,19 @@ func TestUnitPostQueryHelperError(t *testing.T) {
 		TokenAccessor: getSimpleTokenAccessor(),
 	}
 	var err error
-	var requestID uuid.UUID
-	requestID = uuid.New()
+	requestID := newUUID()
 	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID, &Config{})
 	if err == nil {
 		t.Fatalf("should have failed to post")
 	}
 	sr.FuncPost = postTestAppBadGatewayError
-	requestID = uuid.New()
+	requestID = newUUID()
 	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID, &Config{})
 	if err == nil {
 		t.Fatalf("should have failed to post")
 	}
 	sr.FuncPost = postTestSuccessButInvalidJSON
-	requestID = uuid.New()
+	requestID = newUUID()
 	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, requestID, &Config{})
 	if err == nil {
 		t.Fatalf("should have failed to post")
@@ -362,7 +359,7 @@ func TestUnitPostQueryHelperUsesToken(t *testing.T) {
 	accessor.SetTokens(token, "", 0)
 
 	var err error
-	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ uuid.UUID, _ *Config) (*execResponse, error) {
+	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, headers map[string]string, _ []byte, _ time.Duration, _ uuid, _ *Config) (*execResponse, error) {
 		if headers[headerAuthorizationKey] != fmt.Sprintf(headerSnowflakeToken, token) {
 			t.Fatalf("authorization key doesn't match, %v vs %v", headers[headerAuthorizationKey], fmt.Sprintf(headerSnowflakeToken, token))
 		}
@@ -380,7 +377,7 @@ func TestUnitPostQueryHelperUsesToken(t *testing.T) {
 		FuncRenewSession: renewSessionTest,
 		TokenAccessor:    accessor,
 	}
-	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, uuid.New(), &Config{})
+	_, err = postRestfulQueryHelper(context.Background(), sr, &url.Values{}, make(map[string]string), []byte{0x12, 0x34}, 0, newUUID(), &Config{})
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -388,8 +385,8 @@ func TestUnitPostQueryHelperUsesToken(t *testing.T) {
 
 func TestUnitPostQueryHelperRenewSession(t *testing.T) {
 	var err error
-	origRequestID := uuid.New()
-	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid.UUID, _ *Config) (*execResponse, error) {
+	origRequestID := newUUID()
+	postQueryTest := func(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, _ []byte, _ time.Duration, requestID uuid, _ *Config) (*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 99e5661..15153ca 100644
--- a/util_test.go
+++ b/util_test.go
@@ -10,8 +10,6 @@ import (
 	"sync"
 	"testing"
 	"time"
-
-	"github.com/google/uuid"
 )
 
 type tcIntMinMax struct {
@@ -76,15 +74,15 @@ func TestSimpleTokenAccessorGetTokensSynchronization(t *testing.T) {
 }
 
 func TestGetRequestIDFromContext(t *testing.T) {
-	expectedRequestID := uuid.New()
+	expectedRequestID := newUUID()
 	ctx := WithRequestID(context.Background(), expectedRequestID)
 	requestID := getOrGenerateRequestIDFromContext(ctx)
 	if requestID != expectedRequestID {
 		t.Errorf("unexpected request id: %v, expected: %v", requestID, expectedRequestID)
 	}
-	ctx = WithRequestID(context.Background(), uuid.Nil)
+	ctx = WithRequestID(context.Background(), nilUUID)
 	requestID = getOrGenerateRequestIDFromContext(ctx)
-	if requestID == uuid.Nil {
+	if requestID == nilUUID {
 		t.Errorf("unexpected request id, should not be nil")
 	}
 }