fix gcs

snowflakedb · Nov 21, 2024 · 237e56e · 237e56e
1 parent 0652b1f
commit 237e56e
Show file tree

Hide file tree

Showing 8 changed files with 65 additions and 9 deletions.
diff --git a/.run/IcebergBigFilesIT.testMultiplePartUpload.run.xml b/.run/IcebergBigFilesIT.testMultiplePartUpload.run.xml
@@ -0,0 +1,19 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="IcebergBigFilesIT.testMultiplePartUpload" type="JUnit" factoryName="JUnit" nameIsGenerated="true">
+    <module name="snowflake-ingest-sdk" />
+    <extension name="coverage">
+      <pattern>
+        <option name="PATTERN" value="net.snowflake.ingest.streaming.internal.it.*" />
+        <option name="ENABLED" value="true" />
+      </pattern>
+    </extension>
+    <option name="PACKAGE_NAME" value="net.snowflake.ingest.streaming.internal.it" />
+    <option name="MAIN_CLASS_NAME" value="net.snowflake.ingest.streaming.internal.it.IcebergBigFilesIT" />
+    <option name="METHOD_NAME" value="testMultiplePartUpload" />
+    <option name="TEST_OBJECT" value="method" />
+    <option name="VM_PARAMETERS" value="-ea --add-opens=java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED --add-exports=jdk.compiler/com.sun.tools.javac.api=ALL-UNNAMED --add-exports=jdk.compiler/com.sun.tools.javac.code=ALL-UNNAMED --add-exports=jdk.compiler/com.sun.tools.javac.file=ALL-UNNAMED --add-exports=jdk.compiler/com.sun.tools.javac.parser=ALL-UNNAMED --add-exports=jdk.compiler/com.sun.tools.javac.tree=ALL-UNNAMED --add-exports=jdk.compiler/com.sun.tools.javac.util=ALL-UNNAMED --add-opens=java.base/java.lang=ALL-UNNAMED --add-opens=java.base/java.util=ALL-UNNAMED" />
+    <method v="2">
+      <option name="Make" enabled="true" />
+    </method>
+  </configuration>
+</component>
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/InternalStage.java b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStage.java
@@ -176,7 +176,7 @@ private Optional<String> putRemote(String fullFilePath, byte[] data, int retryCo
       }
 
       if (this.useIcebergFileTransferAgent) {
-        return Optional.of(
+        return Optional.ofNullable(
             IcebergFileTransferAgent.uploadWithoutConnection(
                 fileTransferMetadataCopy,
                 inStream,

diff --git a/...in/java/net/snowflake/ingest/streaming/internal/fileTransferAgent/IcebergAzureClient.java b/...in/java/net/snowflake/ingest/streaming/internal/fileTransferAgent/IcebergAzureClient.java
@@ -3,7 +3,11 @@
 import static net.snowflake.client.core.Constants.CLOUD_STORAGE_CREDENTIALS_EXPIRED;
 import static net.snowflake.client.core.HttpUtil.setSessionlessProxyForAzure;
 
-import java.io.*;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.io.InputStream;
 import java.net.SocketTimeoutException;
 import java.net.URI;
 import java.net.URISyntaxException;

diff --git a/...a/net/snowflake/ingest/streaming/internal/fileTransferAgent/IcebergFileTransferAgent.java b/...a/net/snowflake/ingest/streaming/internal/fileTransferAgent/IcebergFileTransferAgent.java
@@ -9,7 +9,8 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.sql.SQLException;
-import java.util.*;
+import java.util.Base64;
+import java.util.Properties;
 import net.snowflake.client.core.HttpClientSettingsKey;
 import net.snowflake.client.core.OCSPMode;
 import net.snowflake.client.jdbc.FileBackedOutputStream;

diff --git a/...main/java/net/snowflake/ingest/streaming/internal/fileTransferAgent/IcebergGCSClient.java b/...main/java/net/snowflake/ingest/streaming/internal/fileTransferAgent/IcebergGCSClient.java
@@ -4,7 +4,11 @@
 import static net.snowflake.client.jdbc.SnowflakeUtil.isBlank;
 
 import com.google.common.base.Strings;
-import java.io.*;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.io.InputStream;
 import java.net.SocketTimeoutException;
 import java.security.InvalidKeyException;
 import java.util.ArrayList;
@@ -19,7 +23,12 @@
 import net.snowflake.client.jdbc.cloud.storage.StageInfo;
 import net.snowflake.client.jdbc.cloud.storage.StorageObjectMetadata;
 import net.snowflake.client.jdbc.internal.google.api.gax.rpc.FixedHeaderProvider;
-import net.snowflake.client.jdbc.internal.google.cloud.storage.*;
+import net.snowflake.client.jdbc.internal.google.cloud.storage.Blob;
+import net.snowflake.client.jdbc.internal.google.cloud.storage.BlobId;
+import net.snowflake.client.jdbc.internal.google.cloud.storage.BlobInfo;
+import net.snowflake.client.jdbc.internal.google.cloud.storage.Storage;
+import net.snowflake.client.jdbc.internal.google.cloud.storage.StorageException;
+import net.snowflake.client.jdbc.internal.google.cloud.storage.StorageOptions;
 import net.snowflake.client.jdbc.internal.snowflake.common.core.SqlState;
 import net.snowflake.client.util.SFPair;
 import net.snowflake.client.util.Stopwatch;
@@ -237,7 +246,10 @@ public String uploadWithPresignedUrlWithoutConnection(
       IOUtils.closeQuietly(is);
     }
 
-    return uploadedBlob.getEtag();
+    // For GCS, return null instead of returning uploadedBlob.getEtag(), as the service does expect
+    // to see blob MD5
+    // and not the ETag.
+    return null;
   }
 
   /**

diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/fileTransferAgent/IcebergS3Client.java b/src/main/java/net/snowflake/ingest/streaming/internal/fileTransferAgent/IcebergS3Client.java
@@ -4,7 +4,11 @@
 import static net.snowflake.client.jdbc.SnowflakeUtil.createDefaultExecutorService;
 import static net.snowflake.client.jdbc.SnowflakeUtil.getRootCause;
 
-import java.io.*;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.io.InputStream;
 import java.net.SocketTimeoutException;
 import java.security.InvalidKeyException;
 import java.security.KeyManagementException;

diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/it/IcebergBigFilesIT.java b/src/test/java/net/snowflake/ingest/streaming/internal/it/IcebergBigFilesIT.java
@@ -213,5 +213,11 @@ public void testMultiplePartUpload() throws Exception {
     assertThat(rs.next()).isTrue();
     assertThat(rs.getLong(1)).isEqualTo(expectedCount.get());
     assertThat(rs.next()).isFalse();
+
+    ResultSet rs2 =
+        conn.createStatement().executeQuery(String.format("select hash_agg(*) from %s", tableName));
+    assertThat(rs2.next()).isTrue();
+    assertThat(rs2.getLong(1)).isNotNull();
+    assertThat(rs2.next()).isFalse();
   }
 }
diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/it/StreamingIngestBigFilesIT.java b/src/test/java/net/snowflake/ingest/streaming/internal/it/StreamingIngestBigFilesIT.java
@@ -7,8 +7,18 @@
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
-import java.util.*;
-import java.util.concurrent.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Random;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import net.snowflake.ingest.TestUtils;
 import net.snowflake.ingest.streaming.OpenChannelRequest;
 import net.snowflake.ingest.streaming.SnowflakeStreamingIngestChannel;