Repository: snowflakedb/gosnowflake
Pull request #440 (2021-06-17 03:06:23)

SNOW-342743 Cancellation Retry
### Description
Enable retry for cancellation queries. Default value is 5

### 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/restful.go b/restful.go
index 8e348bb..0c64a6d 100644
--- a/restful.go
+++ b/restful.go
@@ -29,6 +29,7 @@ const (
 	queryInProgressCode      = "333333"
 	queryInProgressAsyncCode = "333334"
 	sessionExpiredCode       = "390112"
+	queryNotExecuting        = "000605"
 )
 
 // Snowflake Server Endpoints
@@ -420,6 +421,15 @@ func renewRestfulSession(ctx context.Context, sr *snowflakeRestful, timeout time
 	}
 }
 
+func getCancelRetry(ctx context.Context) int {
+	val := ctx.Value(cancelRetry)
+	if val == nil {
+		return 5
+	}
+	cnt, _ := val.(int)
+	return cnt
+}
+
 func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID uuid.UUID, timeout time.Duration) error {
 	logger.WithContext(ctx).Info("cancel query")
 	params := &url.Values{}
@@ -452,12 +462,15 @@ func cancelQuery(ctx context.Context, sr *snowflakeRestful, requestID uuid.UUID,
 			logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
 			return err
 		}
+		ctxRetry := getCancelRetry(ctx)
 		if !respd.Success && respd.Code == sessionExpiredCode {
 			err := sr.FuncRenewSession(ctx, sr, timeout)
 			if err != nil {
 				return err
 			}
 			return sr.FuncCancelQuery(ctx, sr, requestID, timeout)
+		} else if !respd.Success && respd.Code == queryNotExecuting && ctxRetry != 0 {
+			return sr.FuncCancelQuery(context.WithValue(ctx, cancelRetry, ctxRetry-1), sr, requestID, timeout)
 		} else if respd.Success {
 			return nil
 		} else {
diff --git a/util.go b/util.go
index 3e0244b..3a4ca31 100644
--- a/util.go
+++ b/util.go
@@ -37,6 +37,10 @@ const (
 	describeOnly contextKey = "DESCRIBE_ONLY"
 )
 
+const (
+	cancelRetry contextKey = "CANCEL_RETRY"
+)
+
 // WithMultiStatement returns a context that allows the user to execute the desired number of sql queries in one query
 func WithMultiStatement(ctx context.Context, num int) (context.Context, error) {
 	return context.WithValue(ctx, multiStatementCount, num), nil


Test patch:
diff --git a/restful_test.go b/restful_test.go
index 4c3ae1d..1c3ab34 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -52,6 +52,25 @@ func postTestAppUnexpectedError(_ context.Context, _ *snowflakeRestful, _ *url.U
 	}, nil
 }
 
+func postTestQueryNotExecuting(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
+	dd := &execResponseData{}
+	er := &execResponse{
+		Data:    *dd,
+		Message: "",
+		Code:    queryNotExecuting,
+		Success: false,
+	}
+	ba, err := json.Marshal(er)
+	if err != nil {
+		panic(err)
+	}
+
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       &fakeResponseBody{body: ba},
+	}, nil
+}
+
 func postTestRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ map[string]string, _ []byte, _ time.Duration, _ bool) (*http.Response, error) {
 	dd := &execResponseData{}
 	er := &execResponse{
@@ -92,6 +111,30 @@ func postTestAfterRenew(_ context.Context, _ *snowflakeRestful, _ *url.URL, _ ma
 	}, nil
 }
 
+func cancelTestRetry(ctx context.Context, sr *snowflakeRestful, requestID uuid.UUID, timeout time.Duration) error {
+	ctxRetry := getCancelRetry(ctx)
+	u := url.URL{}
+	reqByte, _ := json.Marshal(make(map[string]string))
+	resp, err := sr.FuncPost(ctx, sr, &u, getHeaders(), reqByte, timeout, false)
+	if err != nil {
+		return err
+	}
+	if resp.StatusCode == http.StatusOK {
+		var respd cancelQueryResponse
+		err = json.NewDecoder(resp.Body).Decode(&respd)
+		if err != nil {
+			return err
+		}
+		if !respd.Success && respd.Code == queryNotExecuting && ctxRetry != 0 {
+			return sr.FuncCancelQuery(context.WithValue(ctx, cancelRetry, ctxRetry-1), sr, requestID, timeout)
+		}
+		if ctxRetry == 0 {
+			return nil
+		}
+	}
+	return fmt.Errorf("cancel retry failed")
+}
+
 func TestUnitPostQueryHelperError(t *testing.T) {
 	sr := &snowflakeRestful{
 		FuncPost:      postTestError,
@@ -494,3 +537,16 @@ func TestUnitCancelQuery(t *testing.T) {
 		t.Fatal("should have failed to close session")
 	}
 }
+
+func TestCancelRetry(t *testing.T) {
+	sr := &snowflakeRestful{
+		TokenAccessor:   getSimpleTokenAccessor(),
+		FuncPost:        postTestQueryNotExecuting,
+		FuncCancelQuery: cancelTestRetry,
+	}
+	ctx := context.Background()
+	err := cancelQuery(ctx, sr, getOrGenerateRequestIDFromContext(ctx), time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+}