retry async request if still in progress

async.go

@@ -63,25 +63,36 @@ func (sr *snowflakeRestful) getAsync(
 	defer close(errChannel)
 	token, _, _ := sr.TokenAccessor.GetTokens()
 	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
-	resp, err := sr.FuncGet(ctx, sr, URL, headers, timeout)
-	if err != nil {
-		logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
-		sfError.Message = err.Error()
-		errChannel <- sfError
-		return err
-	}
-	if resp.Body != nil {
-		defer resp.Body.Close()
-	}
 
-	respd := execResponse{}
-	err = json.NewDecoder(resp.Body).Decode(&respd)
-	resp.Body.Close()
-	if err != nil {
-		logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
-		sfError.Message = err.Error()
-		errChannel <- sfError
-		return err
+	var err error
+	var respd execResponse
+
+	for {
+		resp, err := sr.FuncGet(ctx, sr, URL, headers, timeout)
+		if err != nil {
+			logger.WithContext(ctx).Errorf("failed to get response. err: %v", err)
+			sfError.Message = err.Error()
+			errChannel <- sfError
+			return err
+		}
+		if resp.Body != nil {
+			defer resp.Body.Close()
+		}
+		respd = execResponse{} // reset the response
+		err = json.NewDecoder(resp.Body).Decode(&respd)
+		resp.Body.Close()
+		if err != nil {
+			logger.WithContext(ctx).Errorf("failed to decode JSON. err: %v", err)
+			sfError.Message = err.Error()
+			errChannel <- sfError
+			return err
+		}
+		if respd.Code != queryInProgressAsyncCode {
+			// If the query takes longer than 45 seconds to complete the results are not returned.
+			// Retry the request when the query is still in progress.
+			// For all other scenarios continue processing response.
+			break
+		}
 	}
 
 	sc := &snowflakeConn{rest: sr, cfg: cfg}

async_test.go

@@ -136,3 +136,33 @@ func retrieveRows(rows *RowsExtended, ch chan string) {
 	ch <- s
 	close(ch)
 }
+
+func TestLongRunningAsyncQuery(t *testing.T) {
+	ctx, _ := WithMultiStatement(context.Background(), 0)
+	query := "CALL SYSTEM$WAIT(50, 'SECONDS');use snowflake_sample_data"
+	runTests(t, dsn, func(dbt *DBTest) {
+		rows := dbt.mustQueryContext(WithAsyncMode(ctx), query)
+		defer rows.Close()
+		var v interface{}
+		i := 0
+		for {
+			for rows.Next() {
+				err := rows.Scan(&v)
+				if err != nil {
+					t.Fatalf("failed to get result. err: %v", err)
+				}
+				if v == nil {
+					t.Fatal("should have returned a result")
+				}
+				results := []string{"waited 50 seconds", "Statement executed successfully."}
+				if v != results[i] {
+					t.Fatalf("unexpected result returned. expected: %v, but got: %v", results[i], v)
+				}
+				i++
+			}
+			if !rows.NextResultSet() {
+				break
+			}
+		}
+	})
+}