SNOW-1501662 Allow JDBC to handle ZSTD compressed http response streams

FIPS/pom.xml

@@ -429,6 +429,10 @@
                        <pattern>org.jsoup</pattern>
                        <shadedPattern>${shadeBase}.org.jsoup</shadedPattern>
                      </relocation>
+                     <relocation>
+                       <pattern>com.github.luben.zstd</pattern>
+                       <shadedPattern>${shadeBase}.com.github.luben.zstd</shadedPattern>
+                     </relocation>
                      <relocation>
                        <pattern>com.nimbusds</pattern>
                        <shadedPattern>${shadeBase}.com.nimbusds</shadedPattern>

FIPS/scripts/check_content.sh

@@ -1,12 +1,12 @@
 #!/bin/bash -e
 
-# scripts used to check if all dependency is shaded into snowflake internal path
+# scripts used to check if all dependencies are shaded into snowflake internal path
 
 set -o pipefail
 
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )"
 
-if jar tvf $DIR/../target/snowflake-jdbc-fips.jar  | awk '{print $8}' | grep -v -E "^(net|com)/snowflake" | grep -v -E "(com|net)/\$" | grep -v -E "^META-INF" | grep -v -E "^mozilla" | grep -v -E "^com/sun/jna" | grep -v com/sun/ | grep -v mime.types; then
+if jar tvf $DIR/../target/snowflake-jdbc-fips.jar  | awk '{print $8}' | grep -v -E "^(net|com)/snowflake" | grep -v -E "(com|net)/\$" | grep -v -E "^META-INF" | grep -v -E "^mozilla" | grep -v -E "^com/sun/jna" | grep -v com/sun/ | grep -v mime.types | grep -v -E "^aix/" | grep -v -E "^darwin/" | grep -v -E "^freebsd/" | grep -v -E "^linux/" | grep -v -E "^win/"; then
   echo "[ERROR] JDBC jar includes class not under the snowflake namespace"
   exit 1
 fi

ci/scripts/check_content.sh

@@ -8,12 +8,12 @@ set -o pipefail
 
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )"
 
-if jar tvf $DIR/../../target/snowflake-jdbc${package_modifier}.jar  | awk '{print $8}' | grep -v -E "^(net|com)/snowflake" | grep -v -E "(com|net)/\$" | grep -v -E "^META-INF" | grep -v -E "^mozilla" | grep -v -E "^com/sun/jna" | grep -v com/sun/ | grep -v mime.types; then
+if jar tvf $DIR/../../target/snowflake-jdbc${package_modifier}.jar | awk '{print $8}' | grep -v -E "^(net|com)/snowflake" | grep -v -E "(com|net)/\$" | grep -v -E "^META-INF" | grep -v -E "^mozilla" | grep -v -E "^com/sun/jna" | grep -v com/sun/ | grep -v mime.types | grep -v -E "^aix/" | grep -v -E "^darwin/" | grep -v -E "^freebsd/" | grep -v -E "^linux/" | grep -v -E "^win/"; then
   echo "[ERROR] JDBC jar includes class not under the snowflake namespace"
   exit 1
 fi
 
-if jar tvf $DIR/../../target/snowflake-jdbc${package_modifier}.jar  | awk '{print $8}' | grep -E "^META-INF/versions/.*.class"  | grep -v -E "^META-INF/versions/.*/(net|com)/snowflake"; then
-  echo "[ERROR] JDBC jar includes multi release classes not under the snowflake namespace"
+if jar tvf $DIR/../../target/snowflake-jdbc${package_modifier}.jar | awk '{print $8}' | grep -E "^META-INF/versions/.*.class" | grep -v -E "^META-INF/versions/.*/(net|com)/snowflake" | grep -v -E "^aix/" | grep -v -E "^darwin/" | grep -v -E "^freebsd/" | grep -v -E "^linux/" | grep -v -E "^win/"; then
+  echo "[ERROR] JDBC jar includes multi-release classes not under the snowflake namespace"
   exit 1
 fi

linkage-checker-exclusion-rules.xml

@@ -19,11 +19,6 @@
         <Source><Package name="org.apache.commons.compress.compressors"/></Source>
         <Reason>Optional</Reason>
     </LinkageError>
-    <LinkageError>
-        <Target><Package name="com.github.luben.zstd"/></Target>
-        <Source><Package name="org.apache.commons.compress.compressors"/></Source>
-        <Reason>Optional</Reason>
-    </LinkageError>
     <LinkageError>
         <Target><Package name="com.google.appengine.api.urlfetch"/></Target>
         <Source><Package name="com.google.api.client.extensions.appengine"/></Source>

parent-pom.xml

@@ -19,6 +19,7 @@
     <apache.commons.text.version>1.10.0</apache.commons.text.version>
     <apache.httpclient.version>4.5.14</apache.httpclient.version>
     <apache.httpcore.version>4.4.16</apache.httpcore.version>
+    <zstd-jni.version>1.5.6-5</zstd-jni.version>
     <arrow.version>17.0.0</arrow.version>
     <asm.version>9.3</asm.version>
     <avro.version>1.8.1</avro.version>
@@ -327,6 +328,11 @@
         <artifactId>httpcore</artifactId>
         <version>${apache.httpcore.version}</version>
       </dependency>
+      <dependency>
+        <groupId>com.github.luben</groupId>
+        <artifactId>zstd-jni</artifactId>
+        <version>${zstd-jni.version}</version>
+      </dependency>
       <dependency>
         <groupId>org.apache.tika</groupId>
         <artifactId>tika-core</artifactId>
@@ -644,6 +650,10 @@
       <groupId>org.apache.httpcomponents</groupId>
       <artifactId>httpcore</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.github.luben</groupId>
+      <artifactId>zstd-jni</artifactId>
+    </dependency>
     <dependency>
       <groupId>org.apache.tika</groupId>
       <artifactId>tika-core</artifactId>

pom.xml

@@ -947,6 +947,10 @@
                       <pattern>android.annotation</pattern>
                       <shadedPattern>${shadeBase}.android.annotation</shadedPattern>
                     </relocation>
+                    <relocation>
+                      <pattern>com.github.luben.zstd</pattern>
+                      <shadedPattern>${shadeBase}.com.github.luben.zstd</shadedPattern>
+                    </relocation>
                   </relocations>
                   <filters>
                     <filter>

src/main/java/net/snowflake/client/jdbc/DefaultResultStreamProvider.java

@@ -2,6 +2,7 @@
 
 import static net.snowflake.client.core.Constants.MB;
 
+import com.github.luben.zstd.ZstdInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.PushbackInputStream;
@@ -153,6 +154,8 @@ private InputStream detectContentEncodingAndGetInputStream(HttpResponse response
       if ("gzip".equalsIgnoreCase(encoding.getValue())) {
         /* specify buffer size for GZIPInputStream */
         inputStream = new GZIPInputStream(is, STREAM_BUFFER_SIZE);
+      } else if ("zstd".equalsIgnoreCase(encoding.getValue())) {
+        inputStream = new ZstdInputStream(is);
       } else {
         throw new SnowflakeSQLException(
             SqlState.INTERNAL_ERROR,

src/test/java/net/snowflake/client/jdbc/DefaultResultStreamProviderTest.java

@@ -0,0 +1,84 @@
+package net.snowflake.client.jdbc;
+
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.github.luben.zstd.ZstdInputStream;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.InputStream;
+import java.lang.reflect.Method;
+import java.util.zip.GZIPInputStream;
+import java.util.zip.GZIPOutputStream;
+import org.apache.http.Header;
+import org.apache.http.HttpResponse;
+import org.apache.http.entity.InputStreamEntity;
+import org.junit.Before;
+import org.junit.Test;
+
+public class DefaultResultStreamProviderTest {
+
+  private DefaultResultStreamProvider resultStreamProvider;
+  private HttpResponse mockResponse;
+
+  @Before
+  public void setUp() {
+    resultStreamProvider = new DefaultResultStreamProvider();
+    mockResponse = mock(HttpResponse.class);
+  }
+
+  private InputStream invokeDetectContentEncodingAndGetInputStream(
+      HttpResponse response, InputStream inputStream) throws Exception {
+    Method method =
+        DefaultResultStreamProvider.class.getDeclaredMethod(
+            "detectContentEncodingAndGetInputStream", HttpResponse.class, InputStream.class);
+    method.setAccessible(true);
+    return (InputStream) method.invoke(resultStreamProvider, response, inputStream);
+  }
+
+  @Test
+  public void testDetectContentEncodingAndGetInputStream_Gzip() throws Exception {
+    // Mocking gzip content encoding
+    Header encodingHeader = mock(Header.class);
+    when(encodingHeader.getValue()).thenReturn("gzip");
+    when(mockResponse.getFirstHeader("Content-Encoding")).thenReturn(encodingHeader);
+
+    // Creating a valid gzip byte array using a string compressed with GZIPOutputStream
+    byte[] validGzipData;
+    try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+        GZIPOutputStream gzipOutputStream = new GZIPOutputStream(byteArrayOutputStream)) {
+      gzipOutputStream.write("Some data in GZIP".getBytes());
+      gzipOutputStream.finish();
+      validGzipData = byteArrayOutputStream.toByteArray();
+    }
+
+    // Mocking input stream with the valid gzip data
+    InputStream gzipStream = new ByteArrayInputStream(validGzipData);
+
+    // Call the private method using reflection
+    InputStream resultStream =
+        invokeDetectContentEncodingAndGetInputStream(mockResponse, gzipStream);
+
+    assertTrue(resultStream instanceof GZIPInputStream);
+  }
+
+  @Test
+  public void testDetectContentEncodingAndGetInputStream_Zstd() throws Exception {
+    // Mocking zstd content encoding
+    Header encodingHeader = mock(Header.class);
+    when(encodingHeader.getValue()).thenReturn("zstd");
+    when(mockResponse.getFirstHeader("Content-Encoding")).thenReturn(encodingHeader);
+
+    // Mocking input stream with zstd data
+    byte[] zstdData = new byte[] {0x28, (byte) 0xb5, 0x2f}; // Zstd magic numbers
+    InputStream zstdStream = new ByteArrayInputStream(zstdData);
+    when(mockResponse.getEntity()).thenReturn(new InputStreamEntity(zstdStream));
+
+    // Call the private method using reflection
+    InputStream resultStream =
+        invokeDetectContentEncodingAndGetInputStream(mockResponse, zstdStream);
+
+    assertTrue(resultStream instanceof ZstdInputStream);
+  }
+}

thin_public_pom.xml

@@ -64,6 +64,7 @@
     <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
     <slf4j.version>2.0.13</slf4j.version>
     <threeten.version>1.6.9</threeten.version>
+    <zstd-jni.version>1.5.6-5</zstd-jni.version>
   </properties>
 
   <dependencyManagement>
@@ -262,6 +263,11 @@
       <artifactId>jsoup</artifactId>
       <version>${jsoup.version}</version>
     </dependency>
+    <dependency>
+      <groupId>com.github.luben</groupId>
+      <artifactId>zstd-jni</artifactId>
+      <version>${zstd-jni.version}</version>
+    </dependency>
     <dependency>
       <groupId>org.slf4j</groupId>
       <artifactId>slf4j-api</artifactId>