Skip to content

Commit 6ece113

Browse files
sfc-gh-ext-simba-jysfc-gh-ext-simba-jy
authored
Merge branch 'master' into SNOW-1029631-optimize-put-mem
2 parents cd542f1 + 7d34091 commit 6ece113

18 files changed

+540
-124
lines changed

‎.github/workflows/jira_issue.yml

+1-1
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ jobs:
3737
summary: '${{ github.event.issue.title }}'
3838
description: |
3939
${{ github.event.issue.body }} \\ \\ _Created from GitHub Action_ for ${{ github.event.issue.html_url }}
40-
fields: '{ "customfield_11401": {"id": "14723"}, "assignee": {"id": "712020:3c0352b5-63f7-4e26-9afe-38f6f9f0f4c5"}, "components":[{"id":"19286"}] }'
40+
fields: '{ "customfield_11401": {"id": "14723"}, "assignee": {"id": "712020:e1f41916-da57-4fe8-b317-116d5229aa51"}, "components":[{"id":"19286"}], "labels": ["oss"], "priority": {"id": "10001"} }'
4141

4242
- name: Update GitHub Issue
4343
uses: ./jira/gajira-issue-update

‎async.go

+1
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ func (sr *snowflakeRestful) getAsync(
119119
rows.errChannel <- err
120120
return err
121121
}
122+
rows.format = respd.Data.QueryResultFormat
122123
rows.errChannel <- nil // mark query status complete
123124
}
124125
} else {

‎azure_storage_client.go

+22-11
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
)
2424

2525
type snowflakeAzureClient struct {
26+
cfg *Config
2627
}
2728

2829
type azureLocation struct {
@@ -85,9 +86,11 @@ func (util *snowflakeAzureClient) getFileHeader(meta *fileMetadata, filename str
8586
if meta.mockAzureClient != nil {
8687
blobClient = meta.mockAzureClient
8788
}
88-
resp, err := blobClient.GetProperties(context.Background(), &blob.GetPropertiesOptions{
89-
AccessConditions: &blob.AccessConditions{},
90-
CPKInfo: &blob.CPKInfo{},
89+
resp, err := withCloudStorageTimeout(util.cfg, func(ctx context.Context) (blob.GetPropertiesResponse, error) {
90+
return blobClient.GetProperties(ctx, &blob.GetPropertiesOptions{
91+
AccessConditions: &blob.AccessConditions{},
92+
CPKInfo: &blob.CPKInfo{},
93+
})
9194
})
9295
if err != nil {
9396
var se *azcore.ResponseError
@@ -203,9 +206,11 @@ func (util *snowflakeAzureClient) uploadFile(
203206
if meta.realSrcStream != nil {
204207
uploadSrc = meta.realSrcStream
205208
}
206-
_, err = blobClient.UploadStream(context.Background(), uploadSrc, &azblob.UploadStreamOptions{
207-
BlockSize: int64(uploadSrc.Len()),
208-
Metadata: azureMeta,
209+
_, err = withCloudStorageTimeout(util.cfg, func(ctx context.Context) (azblob.UploadStreamResponse, error) {
210+
return blobClient.UploadStream(ctx, uploadSrc, &azblob.UploadStreamOptions{
211+
BlockSize: int64(uploadSrc.Len()),
212+
Metadata: azureMeta,
213+
})
209214
})
210215
} else {
211216
var f *os.File
@@ -228,7 +233,9 @@ func (util *snowflakeAzureClient) uploadFile(
228233
if meta.options.putAzureCallback != nil {
229234
blobOptions.Progress = meta.options.putAzureCallback.call
230235
}
231-
_, err = blobClient.UploadFile(context.Background(), f, blobOptions)
236+
_, err = withCloudStorageTimeout(util.cfg, func(ctx context.Context) (azblob.UploadFileResponse, error) {
237+
return blobClient.UploadFile(ctx, f, blobOptions)
238+
})
232239
}
233240
if err != nil {
234241
var se *azcore.ResponseError
@@ -279,7 +286,9 @@ func (util *snowflakeAzureClient) nativeDownloadFile(
279286
blobClient = meta.mockAzureClient
280287
}
281288
if meta.options.GetFileToStream {
282-
blobDownloadResponse, err := blobClient.DownloadStream(context.Background(), &azblob.DownloadStreamOptions{})
289+
blobDownloadResponse, err := withCloudStorageTimeout(util.cfg, func(ctx context.Context) (azblob.DownloadStreamResponse, error) {
290+
return blobClient.DownloadStream(ctx, &azblob.DownloadStreamOptions{})
291+
})
283292
if err != nil {
284293
return err
285294
}
@@ -295,9 +304,11 @@ func (util *snowflakeAzureClient) nativeDownloadFile(
295304
return err
296305
}
297306
defer f.Close()
298-
_, err = blobClient.DownloadFile(
299-
context.Background(), f, &azblob.DownloadFileOptions{
300-
Concurrency: uint16(maxConcurrency)})
307+
_, err = withCloudStorageTimeout(util.cfg, func(ctx context.Context) (any, error) {
308+
return blobClient.DownloadFile(
309+
ctx, f, &azblob.DownloadFileOptions{
310+
Concurrency: uint16(maxConcurrency)})
311+
})
301312
if err != nil {
302313
return err
303314
}

‎azure_storage_client_test.go

+48-3
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,11 @@ func TestUploadFileWithAzureUploadFailedError(t *testing.T) {
177177
return azblob.UploadFileResponse{}, errors.New("unexpected error uploading file")
178178
},
179179
},
180+
sfa: &snowflakeFileTransferAgent{
181+
sc: &snowflakeConn{
182+
cfg: &Config{},
183+
},
184+
},
180185
}
181186

182187
uploadMeta.realSrcFileName = uploadMeta.srcFileName
@@ -230,6 +235,11 @@ func TestUploadStreamWithAzureUploadFailedError(t *testing.T) {
230235
return azblob.UploadStreamResponse{}, errors.New("unexpected error uploading file")
231236
},
232237
},
238+
sfa: &snowflakeFileTransferAgent{
239+
sc: &snowflakeConn{
240+
cfg: &Config{},
241+
},
242+
},
233243
}
234244

235245
uploadMeta.realSrcStream = uploadMeta.srcStream
@@ -291,6 +301,11 @@ func TestUploadFileWithAzureUploadTokenExpired(t *testing.T) {
291301
}
292302
},
293303
},
304+
sfa: &snowflakeFileTransferAgent{
305+
sc: &snowflakeConn{
306+
cfg: &Config{},
307+
},
308+
},
294309
}
295310

296311
uploadMeta.realSrcFileName = uploadMeta.srcFileName
@@ -362,6 +377,11 @@ func TestUploadFileWithAzureUploadNeedsRetry(t *testing.T) {
362377
}
363378
},
364379
},
380+
sfa: &snowflakeFileTransferAgent{
381+
sc: &snowflakeConn{
382+
cfg: &Config{},
383+
},
384+
},
365385
}
366386

367387
uploadMeta.realSrcFileName = uploadMeta.srcFileName
@@ -418,6 +438,11 @@ func TestDownloadOneFileToAzureFailed(t *testing.T) {
418438
return blob.GetPropertiesResponse{}, nil
419439
},
420440
},
441+
sfa: &snowflakeFileTransferAgent{
442+
sc: &snowflakeConn{
443+
cfg: &Config{},
444+
},
445+
},
421446
}
422447
err = new(remoteStorageUtil).downloadOneFile(&downloadMeta)
423448
if err == nil {
@@ -444,9 +469,14 @@ func TestGetFileHeaderErrorStatus(t *testing.T) {
444469
return blob.GetPropertiesResponse{}, errors.New("failed to retrieve headers")
445470
},
446471
},
472+
sfa: &snowflakeFileTransferAgent{
473+
sc: &snowflakeConn{
474+
cfg: &Config{},
475+
},
476+
},
447477
}
448478

449-
if header, err := new(snowflakeAzureClient).getFileHeader(&meta, "file.txt"); header != nil || err == nil {
479+
if header, err := (&snowflakeAzureClient{cfg: &Config{}}).getFileHeader(&meta, "file.txt"); header != nil || err == nil {
450480
t.Fatalf("expected null header, got: %v", header)
451481
}
452482
if meta.resStatus != errStatus {
@@ -477,9 +507,14 @@ func TestGetFileHeaderErrorStatus(t *testing.T) {
477507
}
478508
},
479509
},
510+
sfa: &snowflakeFileTransferAgent{
511+
sc: &snowflakeConn{
512+
cfg: &Config{},
513+
},
514+
},
480515
}
481516

482-
if header, err := new(snowflakeAzureClient).getFileHeader(&meta, "file.txt"); header != nil || err == nil {
517+
if header, err := (&snowflakeAzureClient{cfg: &Config{}}).getFileHeader(&meta, "file.txt"); header != nil || err == nil {
483518
t.Fatalf("expected null header, got: %v", header)
484519
}
485520
if meta.resStatus != notFoundFile {
@@ -505,7 +540,7 @@ func TestGetFileHeaderErrorStatus(t *testing.T) {
505540
},
506541
}
507542

508-
if header, err := new(snowflakeAzureClient).getFileHeader(&meta, "file.txt"); header != nil || err == nil {
543+
if header, err := (&snowflakeAzureClient{cfg: &Config{}}).getFileHeader(&meta, "file.txt"); header != nil || err == nil {
509544
t.Fatalf("expected null header, got: %v", header)
510545
}
511546
if meta.resStatus != renewToken {
@@ -540,6 +575,11 @@ func TestUploadFileToAzureClientCastFail(t *testing.T) {
540575
options: &SnowflakeFileTransferOptions{
541576
MultiPartThreshold: dataSizeThreshold,
542577
},
578+
sfa: &snowflakeFileTransferAgent{
579+
sc: &snowflakeConn{
580+
cfg: &Config{},
581+
},
582+
},
543583
}
544584

545585
uploadMeta.realSrcFileName = uploadMeta.srcFileName
@@ -573,6 +613,11 @@ func TestAzureGetHeaderClientCastFail(t *testing.T) {
573613
return blob.GetPropertiesResponse{}, nil
574614
},
575615
},
616+
sfa: &snowflakeFileTransferAgent{
617+
sc: &snowflakeConn{
618+
cfg: &Config{},
619+
},
620+
},
576621
}
577622

578623
_, err = new(snowflakeAzureClient).getFileHeader(&meta, "file.txt")

‎chunk_test.go

+39
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"context"
88
"database/sql/driver"
99
"encoding/json"
10+
"errors"
1011
"fmt"
1112
"io"
1213
"math/rand"
@@ -533,6 +534,44 @@ func TestWithArrowBatchesAsync(t *testing.T) {
533534
})
534535
}
535536

537+
func TestWithArrowBatchesButReturningJSON(t *testing.T) {
538+
testWithArrowBatchesButReturningJSON(t, false)
539+
}
540+
541+
func TestWithArrowBatchesButReturningJSONAsync(t *testing.T) {
542+
testWithArrowBatchesButReturningJSON(t, true)
543+
}
544+
545+
func testWithArrowBatchesButReturningJSON(t *testing.T, async bool) {
546+
runSnowflakeConnTest(t, func(sct *SCTest) {
547+
requestID := NewUUID()
548+
pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
549+
defer pool.AssertSize(t, 0)
550+
ctx := WithArrowAllocator(context.Background(), pool)
551+
ctx = WithArrowBatches(ctx)
552+
ctx = WithRequestID(ctx, requestID)
553+
if async {
554+
ctx = WithAsyncMode(ctx)
555+
}
556+
557+
sct.mustExec(forceJSON, nil)
558+
rows := sct.mustQueryContext(ctx, "SELECT 'hello'", nil)
559+
defer rows.Close()
560+
_, err := rows.(SnowflakeRows).GetArrowBatches()
561+
assertNotNilF(t, err)
562+
var se *SnowflakeError
563+
errors.As(err, &se)
564+
assertEqualE(t, se.Message, errJSONResponseInArrowBatchesMode)
565+
566+
ctx = WithRequestID(context.Background(), requestID)
567+
rows2 := sct.mustQueryContext(ctx, "SELECT 'hello'", nil)
568+
defer rows2.Close()
569+
scanValues := make([]driver.Value, 1)
570+
assertNilF(t, rows2.Next(scanValues))
571+
assertEqualE(t, scanValues[0], "hello")
572+
})
573+
}
574+
536575
func TestQueryArrowStream(t *testing.T) {
537576
runSnowflakeConnTest(t, func(sct *SCTest) {
538577
numrows := 50000 // approximately 10 ArrowBatch objects

‎connection.go

+1
Original file line numberDiff line numberDiff line change
@@ -435,6 +435,7 @@ func (sc *snowflakeConn) queryContextInternal(
435435
rows.sc = sc
436436
rows.queryID = data.Data.QueryID
437437
rows.ctx = ctx
438+
rows.format = data.Data.QueryResultFormat
438439

439440
if isMultiStmt(&data.Data) {
440441
// handleMultiQuery is responsible to fill rows with childResults

‎dsn.go

+24-1
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ const (
2525
defaultRequestTimeout = 0 * time.Second // Timeout for retry for request EXCLUDING clientTimeout
2626
defaultJWTTimeout = 60 * time.Second
2727
defaultExternalBrowserTimeout = 120 * time.Second // Timeout for external browser login
28+
defaultCloudStorageTimeout = -1 // Timeout for calling cloud storage.
2829
defaultMaxRetryCount = 7 // specifies maximum number of subsequent retries
2930
defaultDomain = ".snowflakecomputing.com"
3031
cnDomain = ".snowflakecomputing.cn"
@@ -77,6 +78,7 @@ type Config struct {
7778
ClientTimeout time.Duration // Timeout for network round trip + read out http response
7879
JWTClientTimeout time.Duration // Timeout for network round trip + read out http response used when JWT token auth is taking place
7980
ExternalBrowserTimeout time.Duration // Timeout for external browser login
81+
CloudStorageTimeout time.Duration // Timeout for a single call to a cloud storage provider
8082
MaxRetryCount int // Specifies how many times non-periodic HTTP request can be retried
8183

8284
Application string // application name.
@@ -139,11 +141,16 @@ func (c *Config) ocspMode() string {
139141

140142
// DSN constructs a DSN for Snowflake db.
141143
func DSN(cfg *Config) (dsn string, err error) {
142-
if cfg.Region == "us-west-2" {
144+
if strings.ToLower(cfg.Region) == "us-west-2" {
143145
cfg.Region = ""
144146
}
145147
// in case account includes region
146148
region, posDot := extractRegionFromAccount(cfg.Account)
149+
if strings.ToLower(region) == "us-west-2" {
150+
region = ""
151+
cfg.Account = cfg.Account[:posDot]
152+
logger.Info("Ignoring default region .us-west-2 in DSN from Account configuration.")
153+
}
147154
if region != "" {
148155
if cfg.Region != "" {
149156
return "", errRegionConflict()
@@ -215,6 +222,9 @@ func DSN(cfg *Config) (dsn string, err error) {
215222
if cfg.ExternalBrowserTimeout != defaultExternalBrowserTimeout {
216223
params.Add("externalBrowserTimeout", strconv.FormatInt(int64(cfg.ExternalBrowserTimeout/time.Second), 10))
217224
}
225+
if cfg.CloudStorageTimeout != defaultCloudStorageTimeout {
226+
params.Add("cloudStorageTimeout", strconv.FormatInt(int64(cfg.CloudStorageTimeout/time.Second), 10))
227+
}
218228
if cfg.MaxRetryCount != defaultMaxRetryCount {
219229
params.Add("maxRetryCount", strconv.Itoa(cfg.MaxRetryCount))
220230
}
@@ -498,6 +508,9 @@ func fillMissingConfigParameters(cfg *Config) error {
498508
if cfg.ExternalBrowserTimeout == 0 {
499509
cfg.ExternalBrowserTimeout = defaultExternalBrowserTimeout
500510
}
511+
if cfg.CloudStorageTimeout == 0 {
512+
cfg.CloudStorageTimeout = defaultCloudStorageTimeout
513+
}
501514
if cfg.MaxRetryCount == 0 {
502515
cfg.MaxRetryCount = defaultMaxRetryCount
503516
}
@@ -579,6 +592,11 @@ func transformAccountToHost(cfg *Config) (err error) {
579592
// account name is specified instead of host:port
580593
cfg.Account = cfg.Host
581594
region, posDot := extractRegionFromAccount(cfg.Account)
595+
if strings.ToLower(region) == "us-west-2" {
596+
region = ""
597+
cfg.Account = cfg.Account[:posDot]
598+
logger.Info("Ignoring default region .us-west-2 from Account configuration.")
599+
}
582600
if region != "" {
583601
cfg.Region = region
584602
cfg.Account = cfg.Account[:posDot]
@@ -714,6 +732,11 @@ func parseDSNParams(cfg *Config, params string) (err error) {
714732
if err != nil {
715733
return err
716734
}
735+
case "cloudStorageTimeout":
736+
cfg.CloudStorageTimeout, err = parseTimeout(value)
737+
if err != nil {
738+
return err
739+
}
717740
case "maxRetryCount":
718741
cfg.MaxRetryCount, err = strconv.Atoi(value)
719742
if err != nil {

0 commit comments

Comments
 (0)