-
Notifications
You must be signed in to change notification settings - Fork 166
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
SNOW-1757822: Allow JDBC to handle ZSTD decompression (#1932)
- Loading branch information
1 parent
3d3f401
commit 83e5849
Showing
8 changed files
with
148 additions
and
37 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 "^com/github/" | 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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
38 changes: 38 additions & 0 deletions
38
src/main/java/net/snowflake/client/jdbc/CompressedStreamFactory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,38 @@ | ||
| package net.snowflake.client.jdbc; | ||
|
|
||
| import static net.snowflake.client.core.Constants.MB; | ||
| import static net.snowflake.common.core.FileCompressionType.GZIP; | ||
| import static net.snowflake.common.core.FileCompressionType.ZSTD; | ||
|
|
||
| import com.github.luben.zstd.ZstdInputStream; | ||
| import java.io.IOException; | ||
| import java.io.InputStream; | ||
| import java.util.zip.GZIPInputStream; | ||
| import net.snowflake.common.core.SqlState; | ||
| import org.apache.http.Header; | ||
|
|
||
| class CompressedStreamFactory { | ||
|
|
||
| private static final int STREAM_BUFFER_SIZE = MB; | ||
|
|
||
| /** | ||
| * Determine the format of the response, if it is not either plain text or gzip, raise an error. | ||
| */ | ||
| public InputStream createBasedOnEncodingHeader(InputStream is, Header encoding) | ||
| throws IOException, SnowflakeSQLException { | ||
| if (encoding != null) { | ||
| if (GZIP.name().equalsIgnoreCase(encoding.getValue())) { | ||
| return new GZIPInputStream(is, STREAM_BUFFER_SIZE); | ||
| } else if (ZSTD.name().equalsIgnoreCase(encoding.getValue())) { | ||
| return new ZstdInputStream(is); | ||
| } else { | ||
| throw new SnowflakeSQLException( | ||
| SqlState.INTERNAL_ERROR, | ||
| ErrorCode.INTERNAL_ERROR.getMessageCode(), | ||
| "Exception: unexpected compression got " + encoding.getValue()); | ||
| } | ||
| } else { | ||
| return DefaultResultStreamProvider.detectGzipAndGetStream(is); | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
80 changes: 80 additions & 0 deletions
80
src/test/java/net/snowflake/client/jdbc/CompressedStreamFactoryTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,80 @@ | ||
| package net.snowflake.client.jdbc; | ||
|
|
||
| import static org.junit.Assert.assertEquals; | ||
| import static org.junit.Assert.assertTrue; | ||
|
|
||
| import com.github.luben.zstd.ZstdInputStream; | ||
| import com.github.luben.zstd.ZstdOutputStream; | ||
| import java.io.ByteArrayInputStream; | ||
| import java.io.ByteArrayOutputStream; | ||
| import java.io.InputStream; | ||
| import java.nio.charset.StandardCharsets; | ||
| import java.util.zip.GZIPInputStream; | ||
| import java.util.zip.GZIPOutputStream; | ||
| import org.apache.commons.io.IOUtils; | ||
| import org.apache.http.Header; | ||
| import org.apache.http.message.BasicHeader; | ||
| import org.junit.Test; | ||
|
|
||
| public class CompressedStreamFactoryTest { | ||
|
|
||
| private final CompressedStreamFactory factory = new CompressedStreamFactory(); | ||
|
|
||
| @Test | ||
| public void testDetectContentEncodingAndGetInputStream_Gzip() throws Exception { | ||
| // Original data to compress and validate | ||
| String originalData = "Some data in GZIP"; | ||
|
|
||
| // Creating encoding header | ||
| Header encodingHeader = new BasicHeader("Content-Encoding", "gzip"); | ||
|
|
||
| // Creating a gzip byte array using GZIPOutputStream | ||
| byte[] gzipData; | ||
| try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); | ||
| GZIPOutputStream gzipOutputStream = new GZIPOutputStream(byteArrayOutputStream)) { | ||
| gzipOutputStream.write(originalData.getBytes(StandardCharsets.UTF_8)); | ||
| gzipOutputStream.close(); // close to flush and finish the compression | ||
| gzipData = byteArrayOutputStream.toByteArray(); | ||
| } | ||
|
|
||
| // Mocking input stream with the gzip data | ||
| InputStream gzipStream = new ByteArrayInputStream(gzipData); | ||
|
|
||
| // Call the private method using reflection | ||
| InputStream resultStream = factory.createBasedOnEncodingHeader(gzipStream, encodingHeader); | ||
|
|
||
| // Decompress and validate the data matches original | ||
| assertTrue(resultStream instanceof GZIPInputStream); | ||
| String decompressedData = IOUtils.toString(resultStream, StandardCharsets.UTF_8); | ||
| assertEquals(originalData, decompressedData); | ||
| } | ||
|
|
||
| @Test | ||
| public void testDetectContentEncodingAndGetInputStream_Zstd() throws Exception { | ||
| // Original data to compress and validate | ||
| String originalData = "Some data in ZSTD"; | ||
|
|
||
| // Creating encoding header | ||
| Header encodingHeader = new BasicHeader("Content-Encoding", "zstd"); | ||
|
|
||
| // Creating a zstd byte array using ZstdOutputStream | ||
| byte[] zstdData; | ||
| try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); | ||
| ZstdOutputStream zstdOutputStream = new ZstdOutputStream(byteArrayOutputStream)) { | ||
| zstdOutputStream.write(originalData.getBytes(StandardCharsets.UTF_8)); | ||
| zstdOutputStream.close(); // close to flush and finish the compression | ||
| zstdData = byteArrayOutputStream.toByteArray(); | ||
| } | ||
|
|
||
| // Mocking input stream with the zstd data | ||
| InputStream zstdStream = new ByteArrayInputStream(zstdData); | ||
|
|
||
| // Call the private method using reflection | ||
| InputStream resultStream = factory.createBasedOnEncodingHeader(zstdStream, encodingHeader); | ||
|
|
||
| // Decompress and validate the data matches original | ||
| assertTrue(resultStream instanceof ZstdInputStream); | ||
| String decompressedData = IOUtils.toString(resultStream, StandardCharsets.UTF_8); | ||
| assertEquals(originalData, decompressedData); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters