Skip to content

Commit 4d73661

Browse files
sfc-gh-mkubiksfc-gh-mkubik
authoredFeb 3, 2025··
SNOW-1886186 - gather threadExecutor callables and call Future.get() to prevent silent fails (#2035)
·
v3.23.1v3.23.0
1 parent 19db9b3 commit 4d73661

File tree

5 files changed

+194
-132
lines changed

5 files changed

+194
-132
lines changed
 

‎src/main/java/net/snowflake/client/jdbc/ErrorCode.java

+3-1
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,9 @@ public enum ErrorCode {
8585
AUTHENTICATOR_REQUEST_TIMEOUT(200062, SqlState.CONNECTION_EXCEPTION),
8686
INVALID_STRUCT_DATA(200063, SqlState.DATA_EXCEPTION),
8787
DISABLEOCSP_INSECUREMODE_VALUE_MISMATCH(200064, SqlState.INVALID_PARAMETER_VALUE),
88-
TOO_MANY_FILES_TO_DOWNLOAD_AS_STREAM(200065, SqlState.DATA_EXCEPTION);
88+
TOO_MANY_FILES_TO_DOWNLOAD_AS_STREAM(200065, SqlState.DATA_EXCEPTION),
89+
FILE_OPERATION_UPLOAD_ERROR(200066, SqlState.INTERNAL_ERROR),
90+
FILE_OPERATION_DOWNLOAD_ERROR(200067, SqlState.INTERNAL_ERROR);
8991

9092
public static final String errorMessageResource = "net.snowflake.client.jdbc.jdbc_error_messages";
9193

‎src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java

+60-33
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import java.util.HashMap;
3535
import java.util.HashSet;
3636
import java.util.Iterator;
37+
import java.util.LinkedList;
3738
import java.util.List;
3839
import java.util.Locale;
3940
import java.util.Map;
@@ -1667,7 +1668,7 @@ private void uploadStream() throws SnowflakeSQLException {
16671668
queryID,
16681669
ex.getCause(),
16691670
SqlState.INTERNAL_ERROR,
1670-
ErrorCode.INTERNAL_ERROR.getMessageCode());
1671+
ErrorCode.FILE_OPERATION_UPLOAD_ERROR.getMessageCode());
16711672
}
16721673
logger.debug("Done with uploading from a stream");
16731674
} finally {
@@ -1752,6 +1753,7 @@ private void downloadFiles() throws SnowflakeSQLException {
17521753
try {
17531754
threadExecutor = SnowflakeUtil.createDefaultExecutorService("sf-file-download-worker-", 1);
17541755

1756+
List<Future<Void>> downloadFileFutures = new LinkedList<>();
17551757
for (String srcFile : sourceFiles) {
17561758
FileMetadata fileMetadata = fileMetadataMap.get(srcFile);
17571759

@@ -1768,21 +1770,22 @@ private void downloadFiles() throws SnowflakeSQLException {
17681770

17691771
RemoteStoreFileEncryptionMaterial encMat = srcFileToEncMat.get(srcFile);
17701772
String presignedUrl = srcFileToPresignedUrl.get(srcFile);
1771-
threadExecutor.submit(
1772-
getDownloadFileCallable(
1773-
stageInfo,
1774-
srcFile,
1775-
localLocation,
1776-
fileMetadataMap,
1777-
(stageInfo.getStageType() == StageInfo.StageType.LOCAL_FS)
1778-
? null
1779-
: storageFactory.createClient(stageInfo, parallel, encMat, session),
1780-
session,
1781-
command,
1782-
parallel,
1783-
encMat,
1784-
presignedUrl,
1785-
queryID));
1773+
downloadFileFutures.add(
1774+
threadExecutor.submit(
1775+
getDownloadFileCallable(
1776+
stageInfo,
1777+
srcFile,
1778+
localLocation,
1779+
fileMetadataMap,
1780+
(stageInfo.getStageType() == StageInfo.StageType.LOCAL_FS)
1781+
? null
1782+
: storageFactory.createClient(stageInfo, parallel, encMat, session),
1783+
session,
1784+
command,
1785+
parallel,
1786+
encMat,
1787+
presignedUrl,
1788+
queryID)));
17861789

17871790
logger.debug("Submitted download job for: {}", srcFile);
17881791
}
@@ -1792,9 +1795,20 @@ private void downloadFiles() throws SnowflakeSQLException {
17921795
try {
17931796
// wait for all threads to complete without timeout
17941797
threadExecutor.awaitTermination(Long.MAX_VALUE, TimeUnit.DAYS);
1798+
for (Future<Void> downloadFileFuture : downloadFileFutures) {
1799+
if (downloadFileFuture.isDone()) {
1800+
downloadFileFuture.get();
1801+
}
1802+
}
17951803
} catch (InterruptedException ex) {
17961804
throw new SnowflakeSQLLoggedException(
17971805
queryID, session, ErrorCode.INTERRUPTED.getMessageCode(), SqlState.QUERY_CANCELED);
1806+
} catch (ExecutionException ex) {
1807+
throw new SnowflakeSQLException(
1808+
queryID,
1809+
ex.getCause(),
1810+
SqlState.INTERNAL_ERROR,
1811+
ErrorCode.FILE_OPERATION_DOWNLOAD_ERROR.getMessageCode());
17981812
}
17991813
logger.debug("Done with downloading");
18001814
} finally {
@@ -1835,6 +1849,7 @@ private void uploadFiles(Set<String> fileList, int parallel) throws SnowflakeSQL
18351849
threadExecutor =
18361850
SnowflakeUtil.createDefaultExecutorService("sf-file-upload-worker-", parallel);
18371851

1852+
List<Future<Void>> uploadFileFutures = new LinkedList<>();
18381853
for (String srcFile : fileList) {
18391854
FileMetadata fileMetadata = fileMetadataMap.get(srcFile);
18401855

@@ -1862,23 +1877,24 @@ private void uploadFiles(Set<String> fileList, int parallel) throws SnowflakeSQL
18621877
int delay = session.getInjectWaitInPut();
18631878
setUploadDelay(delay);
18641879

1865-
threadExecutor.submit(
1866-
getUploadFileCallable(
1867-
stageInfo,
1868-
srcFile,
1869-
fileMetadata,
1870-
(stageInfo.getStageType() == StageInfo.StageType.LOCAL_FS)
1871-
? null
1872-
: storageFactory.createClient(
1873-
stageInfo, parallel, encryptionMaterial.get(0), session),
1874-
session,
1875-
command,
1876-
null,
1877-
false,
1878-
(parallel > 1 ? 1 : this.parallel),
1879-
srcFileObj,
1880-
encryptionMaterial.get(0),
1881-
queryID));
1880+
uploadFileFutures.add(
1881+
threadExecutor.submit(
1882+
getUploadFileCallable(
1883+
stageInfo,
1884+
srcFile,
1885+
fileMetadata,
1886+
(stageInfo.getStageType() == StageInfo.StageType.LOCAL_FS)
1887+
? null
1888+
: storageFactory.createClient(
1889+
stageInfo, parallel, encryptionMaterial.get(0), session),
1890+
session,
1891+
command,
1892+
null,
1893+
false,
1894+
(parallel > 1 ? 1 : this.parallel),
1895+
srcFileObj,
1896+
encryptionMaterial.get(0),
1897+
queryID)));
18821898

18831899
logger.debug("Submitted copy job for: {}", srcFile);
18841900
}
@@ -1889,9 +1905,20 @@ private void uploadFiles(Set<String> fileList, int parallel) throws SnowflakeSQL
18891905
try {
18901906
// wait for all threads to complete without timeout
18911907
threadExecutor.awaitTermination(Long.MAX_VALUE, TimeUnit.DAYS);
1908+
for (Future<Void> uploadFileFuture : uploadFileFutures) {
1909+
if (uploadFileFuture.isDone()) {
1910+
uploadFileFuture.get();
1911+
}
1912+
}
18921913
} catch (InterruptedException ex) {
18931914
throw new SnowflakeSQLLoggedException(
18941915
queryID, session, ErrorCode.INTERRUPTED.getMessageCode(), SqlState.QUERY_CANCELED);
1916+
} catch (ExecutionException ex) {
1917+
throw new SnowflakeSQLException(
1918+
queryID,
1919+
ex.getCause(),
1920+
SqlState.INTERNAL_ERROR,
1921+
ErrorCode.FILE_OPERATION_UPLOAD_ERROR.getMessageCode());
18951922
}
18961923
logger.debug("Done with uploading");
18971924

‎src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeAzureClient.java

+25-9
Original file line numberDiff line numberDiff line change
@@ -352,13 +352,21 @@ public void download(
352352
// Get the user-defined BLOB metadata
353353
Map<String, String> userDefinedMetadata =
354354
SnowflakeUtil.createCaseInsensitiveMap(blob.getMetadata());
355-
AbstractMap.SimpleEntry<String, String> encryptionData =
356-
parseEncryptionData(userDefinedMetadata.get(AZ_ENCRYPTIONDATAPROP), queryId);
357-
358-
String key = encryptionData.getKey();
359-
String iv = encryptionData.getValue();
360355

361356
if (this.isEncrypting() && this.getEncryptionKeySize() <= 256) {
357+
if (!userDefinedMetadata.containsKey(AZ_ENCRYPTIONDATAPROP)) {
358+
throw new SnowflakeSQLLoggedException(
359+
queryId,
360+
session,
361+
ErrorCode.INTERNAL_ERROR.getMessageCode(),
362+
SqlState.INTERNAL_ERROR,
363+
"Encryption data not found in the metadata of a file being downloaded");
364+
}
365+
AbstractMap.SimpleEntry<String, String> encryptionData =
366+
parseEncryptionData(userDefinedMetadata.get(AZ_ENCRYPTIONDATAPROP), queryId);
367+
368+
String key = encryptionData.getKey();
369+
String iv = encryptionData.getValue();
362370
stopwatch.restart();
363371
if (key == null || iv == null) {
364372
throw new SnowflakeSQLLoggedException(
@@ -452,12 +460,20 @@ public InputStream downloadToStream(
452460
long downloadMillis = stopwatch.elapsedMillis();
453461
Map<String, String> userDefinedMetadata =
454462
SnowflakeUtil.createCaseInsensitiveMap(blob.getMetadata());
455-
AbstractMap.SimpleEntry<String, String> encryptionData =
456-
parseEncryptionData(userDefinedMetadata.get(AZ_ENCRYPTIONDATAPROP), queryId);
457-
String key = encryptionData.getKey();
458-
String iv = encryptionData.getValue();
459463

460464
if (this.isEncrypting() && this.getEncryptionKeySize() <= 256) {
465+
if (!userDefinedMetadata.containsKey(AZ_ENCRYPTIONDATAPROP)) {
466+
throw new SnowflakeSQLLoggedException(
467+
queryId,
468+
session,
469+
ErrorCode.INTERNAL_ERROR.getMessageCode(),
470+
SqlState.INTERNAL_ERROR,
471+
"Encryption data not found in the metadata of a file being downloaded");
472+
}
473+
AbstractMap.SimpleEntry<String, String> encryptionData =
474+
parseEncryptionData(userDefinedMetadata.get(AZ_ENCRYPTIONDATAPROP), queryId);
475+
String key = encryptionData.getKey();
476+
String iv = encryptionData.getValue();
461477
stopwatch.restart();
462478
if (key == null || iv == null) {
463479
throw new SnowflakeSQLLoggedException(

‎src/main/resources/net/snowflake/client/jdbc/jdbc_error_messages.properties

+4-1
Original file line numberDiff line numberDiff line change
@@ -82,4 +82,7 @@ Error message={3}, Extended error info={4}
8282
200061=GCS operation failed: Operation={0}, Error code={1}, Message={2}, Reason={3}
8383
200062=Authentication timed out.
8484
200063=Invalid data - Cannot be parsed and converted to structured type.
85-
85+
200064=The values for 'disableOCSPChecks' and 'insecureMode' must be identical.
86+
200065=Too many files to download as stream
87+
200066=JDBC driver file operation error while performing stage upload.
88+
200067=JDBC driver file operation error while performing stage download.

‎src/test/java/net/snowflake/client/jdbc/FileUploaderLatestIT.java

+102-88
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@
44
package net.snowflake.client.jdbc;
55

66
import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;
7-
import static org.hamcrest.CoreMatchers.instanceOf;
87
import static org.junit.jupiter.api.Assertions.assertEquals;
98
import static org.junit.jupiter.api.Assertions.assertFalse;
9+
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
10+
import static org.junit.jupiter.api.Assertions.assertThrows;
1011
import static org.junit.jupiter.api.Assertions.assertTrue;
1112
import static org.junit.jupiter.api.Assertions.fail;
1213

@@ -174,13 +175,11 @@ public void testGetObjectMetadataFileNotFoundWithGCS() throws Exception {
174175
int idx = location.indexOf('/');
175176
String remoteStageLocation = location.substring(0, idx);
176177
String path = location.substring(idx + 1) + "wrong_file.csv.gz";
177-
client.getObjectMetadata(remoteStageLocation, path);
178-
fail("should raise exception");
179-
} catch (Exception ex) {
180-
assertTrue(
181-
ex instanceof StorageProviderException,
182-
"Wrong type of exception. Message: " + ex.getMessage());
183-
assertTrue(ex.getMessage().matches(".*Blob.*not found in bucket.*"));
178+
StorageProviderException thrown =
179+
assertThrows(
180+
StorageProviderException.class,
181+
() -> client.getObjectMetadata(remoteStageLocation, path));
182+
assertTrue(thrown.getMessage().matches(".*Blob.*not found in bucket.*"));
184183
} finally {
185184
statement.execute("DROP STAGE if exists " + OBJ_META_STAGE);
186185
}
@@ -212,13 +211,11 @@ public void testGetObjectMetadataStorageExceptionWithGCS() throws Exception {
212211
String location = info.getLocation();
213212
int idx = location.indexOf('/');
214213
String remoteStageLocation = location.substring(0, idx);
215-
client.getObjectMetadata(remoteStageLocation, "");
216-
fail("should raise exception");
217-
} catch (Exception ex) {
218-
assertTrue(
219-
ex instanceof StorageProviderException,
220-
"Wrong type of exception. Message: " + ex.getMessage());
221-
assertTrue(ex.getMessage().matches(".*Permission.*denied.*"));
214+
StorageProviderException thrown =
215+
assertThrows(
216+
StorageProviderException.class,
217+
() -> client.getObjectMetadata(remoteStageLocation, ""));
218+
assertTrue(thrown.getMessage().matches(".*Permission.*denied.*"));
222219
} finally {
223220
statement.execute("DROP STAGE if exists " + OBJ_META_STAGE);
224221
}
@@ -248,12 +245,13 @@ public void testNullCommand() throws SQLException {
248245
try {
249246
statement.execute("create or replace stage testStage");
250247
SFSession sfSession = con.unwrap(SnowflakeConnectionV1.class).getSfSession();
251-
SnowflakeFileTransferAgent sfAgent =
252-
new SnowflakeFileTransferAgent(null, sfSession, new SFStatement(sfSession));
253-
} catch (SnowflakeSQLException err) {
254-
assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode());
248+
SnowflakeSQLException thrown =
249+
assertThrows(
250+
SnowflakeSQLException.class,
251+
() -> new SnowflakeFileTransferAgent(null, sfSession, new SFStatement(sfSession)));
255252
assertTrue(
256-
err.getMessage()
253+
thrown
254+
.getMessage()
257255
.contains("JDBC driver internal error: Missing sql for statement execution"));
258256
} finally {
259257
statement.execute("drop stage if exists testStage");
@@ -281,20 +279,24 @@ public void testCompressStreamWithGzipException() throws Exception {
281279

282280
String srcPath = getFullPathFileInResource(TEST_DATA_FILE);
283281
InputStream inputStream = new FileInputStream(srcPath);
284-
SnowflakeFileTransferAgent.uploadWithoutConnection(
285-
SnowflakeFileTransferConfig.Builder.newInstance()
286-
.setSnowflakeFileTransferMetadata(metadata)
287-
.setUploadStream(inputStream)
288-
.setRequireCompress(true)
289-
.setNetworkTimeoutInMilli(0)
290-
.setOcspMode(OCSPMode.FAIL_OPEN)
291-
.setSFSession(sfSession)
292-
.setCommand(PUT_COMMAND)
293-
.build());
294-
} catch (SnowflakeSQLException err) {
295-
assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode());
282+
SnowflakeSQLException thrown =
283+
assertThrows(
284+
SnowflakeSQLException.class,
285+
() ->
286+
SnowflakeFileTransferAgent.uploadWithoutConnection(
287+
SnowflakeFileTransferConfig.Builder.newInstance()
288+
.setSnowflakeFileTransferMetadata(metadata)
289+
.setUploadStream(inputStream)
290+
.setRequireCompress(true)
291+
.setNetworkTimeoutInMilli(0)
292+
.setOcspMode(OCSPMode.FAIL_OPEN)
293+
.setSFSession(sfSession)
294+
.setCommand(PUT_COMMAND)
295+
.build()));
296+
assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), thrown.getErrorCode());
296297
assertTrue(
297-
err.getMessage()
298+
thrown
299+
.getMessage()
298300
.contains("JDBC driver internal error: error encountered for compression"));
299301
} finally {
300302
statement.execute("DROP STAGE if exists testStage");
@@ -325,20 +327,24 @@ public void testCompressStreamWithGzipNoDigestException() throws Exception {
325327
String srcPath = getFullPathFileInResource(TEST_DATA_FILE);
326328

327329
InputStream inputStream = new FileInputStream(srcPath);
328-
SnowflakeFileTransferAgent.uploadWithoutConnection(
329-
SnowflakeFileTransferConfig.Builder.newInstance()
330-
.setSnowflakeFileTransferMetadata(metadata)
331-
.setUploadStream(inputStream)
332-
.setRequireCompress(true)
333-
.setNetworkTimeoutInMilli(0)
334-
.setOcspMode(OCSPMode.FAIL_OPEN)
335-
.setSFSession(sfSession)
336-
.setCommand(PUT_COMMAND)
337-
.build());
338-
} catch (SnowflakeSQLException err) {
339-
assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode());
330+
SnowflakeSQLException thrown =
331+
assertThrows(
332+
SnowflakeSQLException.class,
333+
() ->
334+
SnowflakeFileTransferAgent.uploadWithoutConnection(
335+
SnowflakeFileTransferConfig.Builder.newInstance()
336+
.setSnowflakeFileTransferMetadata(metadata)
337+
.setUploadStream(inputStream)
338+
.setRequireCompress(true)
339+
.setNetworkTimeoutInMilli(0)
340+
.setOcspMode(OCSPMode.FAIL_OPEN)
341+
.setSFSession(sfSession)
342+
.setCommand(PUT_COMMAND)
343+
.build()));
344+
assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), thrown.getErrorCode());
340345
assertTrue(
341-
err.getMessage()
346+
thrown
347+
.getMessage()
342348
.contains("JDBC driver internal error: error encountered for compression"));
343349
} finally {
344350
statement.execute("DROP STAGE if exists testStage");
@@ -369,19 +375,23 @@ public void testUploadWithoutConnectionException() throws Exception {
369375
String srcPath = getFullPathFileInResource(TEST_DATA_FILE);
370376

371377
InputStream inputStream = new FileInputStream(srcPath);
372-
SnowflakeFileTransferAgent.uploadWithoutConnection(
373-
SnowflakeFileTransferConfig.Builder.newInstance()
374-
.setSnowflakeFileTransferMetadata(metadata)
375-
.setUploadStream(inputStream)
376-
.setRequireCompress(true)
377-
.setNetworkTimeoutInMilli(0)
378-
.setOcspMode(OCSPMode.FAIL_OPEN)
379-
.setSFSession(sfSession)
380-
.setCommand(PUT_COMMAND)
381-
.build());
382-
} catch (Exception err) {
378+
Exception thrown =
379+
assertThrows(
380+
Exception.class,
381+
() ->
382+
SnowflakeFileTransferAgent.uploadWithoutConnection(
383+
SnowflakeFileTransferConfig.Builder.newInstance()
384+
.setSnowflakeFileTransferMetadata(metadata)
385+
.setUploadStream(inputStream)
386+
.setRequireCompress(true)
387+
.setNetworkTimeoutInMilli(0)
388+
.setOcspMode(OCSPMode.FAIL_OPEN)
389+
.setSFSession(sfSession)
390+
.setCommand(PUT_COMMAND)
391+
.build()));
383392
assertTrue(
384-
err.getMessage()
393+
thrown
394+
.getMessage()
385395
.contains(
386396
"Exception encountered during file upload: failed to push to remote store"));
387397
} finally {
@@ -401,9 +411,8 @@ public void testInitFileMetadataFileNotFound() throws Exception {
401411
SnowflakeFileTransferAgent sfAgent =
402412
new SnowflakeFileTransferAgent(PUT_COMMAND, sfSession, new SFStatement(sfSession));
403413

404-
sfAgent.execute();
405-
} catch (SnowflakeSQLException err) {
406-
assertEquals(200008, err.getErrorCode());
414+
SnowflakeSQLException thrown = assertThrows(SnowflakeSQLException.class, sfAgent::execute);
415+
assertEquals(ErrorCode.FILE_NOT_FOUND.getMessageCode(), thrown.getErrorCode());
407416
} finally {
408417
statement.execute("DROP STAGE if exists testStage");
409418
}
@@ -422,9 +431,8 @@ public void testInitFileMetadataFileIsDirectory() throws Exception {
422431
String command = "put file://" + srcPath + " @testStage";
423432
SnowflakeFileTransferAgent sfAgent =
424433
new SnowflakeFileTransferAgent(command, sfSession, new SFStatement(sfSession));
425-
sfAgent.execute();
426-
} catch (SnowflakeSQLException err) {
427-
assertEquals(200009, err.getErrorCode());
434+
SnowflakeSQLException thrown = assertThrows(SnowflakeSQLException.class, sfAgent::execute);
435+
assertEquals(ErrorCode.FILE_IS_DIRECTORY.getMessageCode(), thrown.getErrorCode());
428436
} finally {
429437
statement.execute("DROP STAGE if exists testStage");
430438
}
@@ -445,10 +453,10 @@ public void testCompareAndSkipFilesException() throws Exception {
445453
SnowflakeFileTransferAgent sfAgent =
446454
new SnowflakeFileTransferAgent(command, sfSession, new SFStatement(sfSession));
447455

448-
sfAgent.execute();
449-
} catch (SnowflakeSQLException err) {
450-
assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode());
451-
assertTrue(err.getMessage().contains("Error reading:"));
456+
SnowflakeSQLException thrown = assertThrows(SnowflakeSQLException.class, sfAgent::execute);
457+
assertEquals(
458+
(long) ErrorCode.FILE_OPERATION_UPLOAD_ERROR.getMessageCode(), thrown.getErrorCode());
459+
assertInstanceOf(NoSuchAlgorithmException.class, thrown.getCause().getCause());
452460
} finally {
453461
statement.execute("DROP STAGE if exists testStage");
454462
}
@@ -466,12 +474,14 @@ public void testParseCommandException() throws SQLException {
466474
try {
467475
statement.execute("create or replace stage testStage");
468476
SFSession sfSession = con.unwrap(SnowflakeConnectionV1.class).getSfSession();
469-
SnowflakeFileTransferAgent sfAgent =
470-
new SnowflakeFileTransferAgent(PUT_COMMAND, sfSession, new SFStatement(sfSession));
471-
472-
} catch (SnowflakeSQLException err) {
473-
assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), err.getErrorCode());
474-
assertTrue(err.getMessage().contains("Failed to parse the locations"));
477+
SnowflakeSQLException thrown =
478+
assertThrows(
479+
SnowflakeSQLException.class,
480+
() ->
481+
new SnowflakeFileTransferAgent(
482+
PUT_COMMAND, sfSession, new SFStatement(sfSession)));
483+
assertEquals((long) ErrorCode.INTERNAL_ERROR.getMessageCode(), thrown.getErrorCode());
484+
assertTrue(thrown.getMessage().contains("Failed to parse the locations"));
475485
} finally {
476486
statement.execute("DROP STAGE if exists testStage");
477487
}
@@ -530,10 +540,9 @@ public void testListObjectsStorageException() throws Exception {
530540
SnowflakeFileTransferAgent sfAgent =
531541
new SnowflakeFileTransferAgent(command, sfSession, new SFStatement(sfSession));
532542

533-
sfAgent.execute();
534-
} catch (SnowflakeSQLException err) {
535-
assertEquals(200016, err.getErrorCode());
536-
assertTrue(err.getMessage().contains("Encountered exception during listObjects"));
543+
SnowflakeSQLException thrown = assertThrows(SnowflakeSQLException.class, sfAgent::execute);
544+
assertEquals(ErrorCode.IO_ERROR.getMessageCode(), thrown.getErrorCode());
545+
assertTrue(thrown.getMessage().contains("Encountered exception during listObjects"));
537546
} finally {
538547
statement.execute("DROP STAGE if exists testStage");
539548
}
@@ -555,13 +564,19 @@ public void testUploadStreamInterruptedException() throws IOException, SQLExcept
555564
outputStream.flush();
556565

557566
// upload the data to user stage under testUploadStream with name hello.txt
558-
connection
559-
.unwrap(SnowflakeConnection.class)
560-
.uploadStream(
561-
"~", DEST_PREFIX, outputStream.asByteSource().openStream(), "hello.txt", false);
562-
563-
} catch (SnowflakeSQLLoggedException err) {
564-
assertEquals(200003, err.getErrorCode());
567+
SnowflakeSQLException thrown =
568+
assertThrows(
569+
SnowflakeSQLException.class,
570+
() ->
571+
connection
572+
.unwrap(SnowflakeConnection.class)
573+
.uploadStream(
574+
"~",
575+
DEST_PREFIX,
576+
outputStream.asByteSource().openStream(),
577+
"hello.txt",
578+
false));
579+
assertEquals(ErrorCode.INTERRUPTED.getMessageCode(), thrown.getErrorCode());
565580
} finally {
566581
statement.execute("rm @~/" + DEST_PREFIX);
567582
}
@@ -632,9 +647,8 @@ public void testUploadFileCallableFileNotFound() throws Exception {
632647
String command = "PUT file://" + getFullPathFileInResource(TEST_DATA_FILE) + " @testStage";
633648
SnowflakeFileTransferAgent sfAgent =
634649
new SnowflakeFileTransferAgent(command, sfSession, new SFStatement(sfSession));
635-
sfAgent.execute();
636-
} catch (Exception err) {
637-
assertEquals(err.getCause(), instanceOf(FileNotFoundException.class));
650+
Exception thrown = assertThrows(Exception.class, sfAgent::execute);
651+
assertInstanceOf(FileNotFoundException.class, thrown.getCause());
638652
} finally {
639653
statement.execute("DROP STAGE if exists testStage");
640654
}

0 commit comments

Comments
 (0)
Please sign in to comment.