@@ -23,6 +23,7 @@ import (
23
23
)
24
24
25
25
type snowflakeAzureClient struct {
26
+ cfg * Config
26
27
}
27
28
28
29
type azureLocation struct {
@@ -85,9 +86,11 @@ func (util *snowflakeAzureClient) getFileHeader(meta *fileMetadata, filename str
85
86
if meta .mockAzureClient != nil {
86
87
blobClient = meta .mockAzureClient
87
88
}
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
+ })
91
94
})
92
95
if err != nil {
93
96
var se * azcore.ResponseError
@@ -203,9 +206,11 @@ func (util *snowflakeAzureClient) uploadFile(
203
206
if meta .realSrcStream != nil {
204
207
uploadSrc = meta .realSrcStream
205
208
}
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
+ })
209
214
})
210
215
} else {
211
216
var f * os.File
@@ -228,7 +233,9 @@ func (util *snowflakeAzureClient) uploadFile(
228
233
if meta .options .putAzureCallback != nil {
229
234
blobOptions .Progress = meta .options .putAzureCallback .call
230
235
}
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
+ })
232
239
}
233
240
if err != nil {
234
241
var se * azcore.ResponseError
@@ -279,7 +286,9 @@ func (util *snowflakeAzureClient) nativeDownloadFile(
279
286
blobClient = meta .mockAzureClient
280
287
}
281
288
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
+ })
283
292
if err != nil {
284
293
return err
285
294
}
@@ -295,9 +304,11 @@ func (util *snowflakeAzureClient) nativeDownloadFile(
295
304
return err
296
305
}
297
306
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
+ })
301
312
if err != nil {
302
313
return err
303
314
}
0 commit comments