Skip to content

Commit

Permalink
Test GCS HEAD requests with retry
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-pfus committed Dec 12, 2024
1 parent 3364567 commit c2d88a0
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 1 deletion.
9 changes: 8 additions & 1 deletion gcs_storage_client.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
package gosnowflake

import (
"context"
"encoding/json"
"fmt"
"io"
Expand All @@ -11,6 +12,7 @@ import (
"os"
"strconv"
"strings"
"time"
)

const (
Expand Down Expand Up @@ -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
}
Expand Down
9 changes: 9 additions & 0 deletions put_get_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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() {
Expand Down Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions retry.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit c2d88a0

Please sign in to comment.