Test GCS HEAD requests with retry

Author: sfc-gh-pfus

gcs_storage_client.go

@@ -3,6 +3,7 @@
 package gosnowflake
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -11,6 +12,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 )
 
 const (
@@ -74,7 +76,12 @@ func (util *snowflakeGcsClient) getFileHeader(meta *fileMetadata, filename strin
 		if meta.mockGcsClient != nil {
 			client = meta.mockGcsClient
 		}
-		resp, err := client.Do(req)
+		req.Close = true
+		r := newRetryHTTP(context.Background(), client, http.NewRequest, URL, gcsHeaders, time.Second, 3, defaultTimeProvider, nil) // TODO replace with timeout context
+		r.doHead()
+		fmt.Printf("Before calling HEAD to GCS\n")
+		resp, err := r.execute()
+		fmt.Printf("After calling HEAD to GCS, err: %v\n", err)
 		if err != nil {
 			return nil, err
 		}
@@ -394,6 +401,6 @@ func (util *snowflakeGcsClient) isTokenExpired(resp *http.Response) bool {
 
 func newGcsClient() gcsAPI {
 	return &http.Client{
-		Transport: SnowflakeTransport,
+		//Transport: SnowflakeTransport,
 	}
 }

put_get_test.go

@@ -253,6 +253,11 @@ func TestPutLocalFile(t *testing.T) {
 }
 
 func TestPutGetWithAutoCompressFalse(t *testing.T) {
+	level := logger.GetLogLevel()
+	_ = logger.SetLogLevel("debug")
+	defer func() {
+		_ = logger.SetLogLevel(level)
+	}()
 	tmpDir := t.TempDir()
 	testData := filepath.Join(tmpDir, "data.txt")
 	f, err := os.Create(testData)
@@ -267,14 +272,17 @@ func TestPutGetWithAutoCompressFalse(t *testing.T) {
 		assertNilF(t, f.Close())
 	}()
 
+	fmt.Printf("Starting TestPutGetWithAutoCompressFalse at %v\n", time.Now())
 	runDBTest(t, func(dbt *DBTest) {
 		stageDir := "test_put_uncompress_file_" + randomString(10)
 		dbt.mustExec("rm @~/" + stageDir)
 
 		// PUT test
+		fmt.Printf("Running PUT at %v\n", time.Now())
 		sqlText := fmt.Sprintf("put 'file://%v' @~/%v auto_compress=FALSE", testData, stageDir)
 		sqlText = strings.ReplaceAll(sqlText, "\\", "\\\\")
 		dbt.mustExec(sqlText)
+		fmt.Printf("Finished PUT at %v\n", time.Now())
 		defer dbt.mustExec("rm @~/" + stageDir)
 		rows := dbt.mustQuery("ls @~/" + stageDir)
 		defer func() {
@@ -327,6 +335,7 @@ func TestPutGetWithAutoCompressFalse(t *testing.T) {
 }
 
 func TestPutOverwrite(t *testing.T) {
+	fmt.Printf("Starting TestPutOverwrite at %v\n", time.Now())
 	tmpDir := t.TempDir()
 	testData := filepath.Join(tmpDir, "data.txt")
 	f, err := os.Create(testData)

retry.go

@@ -282,6 +282,11 @@ func (r *retryHTTP) doPost() *retryHTTP {
 	return r
 }
 
+func (r *retryHTTP) doHead() *retryHTTP {
+	r.method = "HEAD"
+	return r
+}
+
 func (r *retryHTTP) setBody(body []byte) *retryHTTP {
 	r.bodyCreator = func() ([]byte, error) {
 		return body, nil