Merge pull request #30 from dbt-labs/harry/cherry_pick_context_deadline_exceeded_fix

Cherry-pick a change to mitigate the context deadline exceeded error

Author: xuliangs

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 ## Upcoming Release
 
+Bug fixes:
+
+- Handle HTTP307 & 308 in drivers to achieve better resiliency to backend errors (snowflakedb/gosnowflake#1616).
+
+- Fix unsafe reflection of nil pointer on DECFLOAT func in bind uploader (snowflakedb/gosnowflake#1604).
 - Added temporary download files cleanup (snowflakedb/gosnowflake#1577)
 - Marked fields as deprecated (snowflakedb/gosnowflake#1556)
 - Exposed `QueryStatus` from `SnowflakeResult` and `SnowflakeRows` in `GetStatus()` function (snowflakedb/gosnowflake#1556)

retry.go

@@ -3,7 +3,6 @@ package gosnowflake
 import (
 	"bytes"
 	"context"
-	"errors"
 	"fmt"
 	"io"
 	"math"
@@ -323,7 +322,7 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 		}
 		res, err = r.client.Do(req)
 		// check if it can retry.
-		retryable, err := isRetryableError(req, res, err)
+		retryable, err := isRetryableError(r.ctx, req, res, err)
 		if !retryable {
 			return res, err
 		}
@@ -387,11 +386,11 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 	}
 }
 
-func isRetryableError(req *http.Request, res *http.Response, err error) (bool, error) {
+func isRetryableError(ctx context.Context, req *http.Request, res *http.Response, err error) (bool, error) {
+	if ctx.Err() != nil {
+		return false, ctx.Err()
+	}
 	if err != nil && res == nil { // Failed http connection. Most probably client timeout.
-		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-			return false, err
-		}
 		return true, err
 	}
 	if res == nil || req == nil {

retry_test.go

@@ -3,6 +3,7 @@ package gosnowflake
 import (
 	"bytes"
 	"context"
+	"database/sql"
 	"fmt"
 	"io"
 	"net/http"
@@ -486,70 +487,84 @@ func TestLoginRetry429(t *testing.T) {
 }
 
 func TestIsRetryable(t *testing.T) {
+	deadLineCtx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
+	defer cancel()
+	time.Sleep(2 * time.Nanosecond)
+
 	tcs := []struct {
+		ctx      context.Context
 		req      *http.Request
 		res      *http.Response
 		err      error
 		expected bool
 	}{
 		{
+			ctx:      context.Background(),
 			req:      nil,
 			res:      nil,
 			err:      nil,
 			expected: false,
 		},
 		{
+			ctx:      context.Background(),
 			req:      nil,
 			res:      &http.Response{StatusCode: http.StatusBadRequest},
 			err:      nil,
 			expected: false,
 		},
 		{
+			ctx:      context.Background(),
 			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
 			res:      nil,
 			err:      nil,
 			expected: false,
 		},
 		{
+			ctx:      context.Background(),
 			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
 			res:      &http.Response{StatusCode: http.StatusNotFound},
 			expected: false,
 		},
 		{
-			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
-			res:      nil,
-			err:      &url.Error{Err: context.Canceled},
-			expected: false,
-		},
-		{
+			ctx:      context.Background(),
 			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
 			res:      nil,
 			err:      &url.Error{Err: context.DeadlineExceeded},
-			expected: false,
+			expected: true,
 		},
 		{
+			ctx:      context.Background(),
 			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
 			res:      nil,
 			err:      errUnknownError(),
 			expected: true,
 		},
 		{
+			ctx:      context.Background(),
 			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
 			res:      &http.Response{StatusCode: http.StatusTooManyRequests},
 			err:      nil,
 			expected: true,
 		},
 		{
+			ctx:      deadLineCtx,
+			req:      &http.Request{URL: &url.URL{Path: loginRequestPath}},
+			res:      nil,
+			err:      &url.Error{Err: context.DeadlineExceeded},
+			expected: false,
+		},
+		{
+			ctx:      deadLineCtx,
 			req:      &http.Request{URL: &url.URL{Path: queryRequestPath}},
-			res:      &http.Response{StatusCode: http.StatusServiceUnavailable},
-			err:      nil,
-			expected: true,
+			res:      nil,
+			err:      &url.Error{Err: context.DeadlineExceeded},
+			expected: false,
 		},
 	}
 
 	for _, tc := range tcs {
 		t.Run(fmt.Sprintf("req %v, resp %v", tc.req, tc.res), func(t *testing.T) {
-			result, _ := isRetryableError(tc.req, tc.res, tc.err)
+			result, _ := isRetryableError(tc.ctx, tc.req, tc.res, tc.err)
 			if result != tc.expected {
 				t.Fatalf("expected %v, got %v; request: %v, response: %v", tc.expected, result, tc.req, tc.res)
 			}
@@ -673,3 +688,12 @@ func TestCalculateRetryWaitForNonAuthRequests(t *testing.T) {
 		})
 	}
 }
+
+func TestRedirectRetry(t *testing.T) {
+	wiremock.registerMappings(t, newWiremockMapping("retry/redirection_retry_workflow.json"))
+	cfg := wiremock.connectionConfig()
+	cfg.ClientTimeout = 3 * time.Second
+	connector := NewConnector(SnowflakeDriver{}, *cfg)
+	db := sql.OpenDB(connector)
+	runSmokeQuery(t, db)
+}