From c6dfbf31be807072e31bfe137f7c7dc7861fad20 Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Mon, 10 Jun 2024 17:31:43 -0700 Subject: [PATCH 1/9] SNOW-1437885 Disable blob interleaving when in Iceberg mode (#763) * Add parameters & disable interleaving mode --- ...SnowflakeStreamingIngestClientFactory.java | 16 ++- ...nowflakeStreamingIngestClientInternal.java | 16 ++- .../ingest/utils/ParameterProvider.java | 45 +++++++-- .../streaming/internal/ChannelCacheTest.java | 2 +- .../streaming/internal/FlushServiceTest.java | 20 +++- .../streaming/internal/OAuthBasicTest.java | 2 +- .../internal/ParameterProviderTest.java | 97 ++++++++++++++----- .../internal/RegisterServiceTest.java | 15 ++- .../SnowflakeStreamingIngestChannelTest.java | 42 +++++--- .../SnowflakeStreamingIngestClientTest.java | 34 ++++++- .../internal/StreamingIngestStageTest.java | 2 +- 11 files changed, 232 insertions(+), 59 deletions(-) diff --git a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java index cd6d78787..89e528693 100644 --- a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java +++ b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java @@ -28,6 +28,10 @@ public static class Builder { // Allows client to override some default parameter values private Map parameterOverrides; + // Indicates whether it's streaming to Iceberg tables. Open channels on regular tables should + // fail in this mode. + private boolean isIcebergMode; + // Indicates whether it's under test mode private boolean isTestMode; @@ -45,6 +49,11 @@ public Builder setParameterOverrides(Map parameterOverrides) { return this; } + public Builder setIsIceberg(boolean isIcebergMode) { + this.isIcebergMode = isIcebergMode; + return this; + } + public Builder setIsTestMode(boolean isTestMode) { this.isTestMode = isTestMode; return this; @@ -58,7 +67,12 @@ public SnowflakeStreamingIngestClient build() { SnowflakeURL accountURL = new SnowflakeURL(prop.getProperty(Constants.ACCOUNT_URL)); return new SnowflakeStreamingIngestClientInternal<>( - this.name, accountURL, prop, this.parameterOverrides, this.isTestMode); + this.name, + accountURL, + prop, + this.parameterOverrides, + this.isIcebergMode, + this.isTestMode); } } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java index 2990b49d8..13cc673ff 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java @@ -121,6 +121,9 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea // Indicates whether the client has closed private volatile boolean isClosed; + // Indicates wheter the client is streaming to Iceberg tables + private final boolean isIcebergMode; + // Indicates whether the client is under test mode private final boolean isTestMode; @@ -152,6 +155,7 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea * @param accountURL Snowflake account url * @param prop connection properties * @param httpClient http client for sending request + * @param isIcebergMode whether we're streaming to iceberg tables * @param isTestMode whether we're under test mode * @param requestBuilder http request builder * @param parameterOverrides parameters we override in case we want to set different values @@ -161,13 +165,15 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea SnowflakeURL accountURL, Properties prop, CloseableHttpClient httpClient, + boolean isIcebergMode, boolean isTestMode, RequestBuilder requestBuilder, Map parameterOverrides) { - this.parameterProvider = new ParameterProvider(parameterOverrides, prop); + this.parameterProvider = new ParameterProvider(parameterOverrides, prop, isIcebergMode); this.name = name; String accountName = accountURL == null ? null : accountURL.getAccount(); + this.isIcebergMode = isIcebergMode; this.isTestMode = isTestMode; this.httpClient = httpClient == null ? HttpUtil.getHttpClient(accountName) : httpClient; this.channelCache = new ChannelCache<>(); @@ -251,6 +257,7 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea * @param accountURL Snowflake account url * @param prop connection properties * @param parameterOverrides map of parameters to override for this client + * @param isIcebergMode whether we're streaming to iceberg tables * @param isTestMode indicates whether it's under test mode */ public SnowflakeStreamingIngestClientInternal( @@ -258,16 +265,17 @@ public SnowflakeStreamingIngestClientInternal( SnowflakeURL accountURL, Properties prop, Map parameterOverrides, + boolean isIcebergMode, boolean isTestMode) { - this(name, accountURL, prop, null, isTestMode, null, parameterOverrides); + this(name, accountURL, prop, null, isIcebergMode, isTestMode, null, parameterOverrides); } /*** Constructor for TEST ONLY * * @param name the name of the client */ - SnowflakeStreamingIngestClientInternal(String name) { - this(name, null, null, null, true, null, new HashMap<>()); + SnowflakeStreamingIngestClientInternal(String name, boolean isIcebergMode) { + this(name, null, null, null, isIcebergMode, true, null, new HashMap<>()); } // TESTING ONLY - inject the request builder diff --git a/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java b/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java index b98972a7d..33f791ca5 100644 --- a/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java +++ b/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java @@ -1,5 +1,7 @@ package net.snowflake.ingest.utils; +import static net.snowflake.ingest.utils.ErrorCode.INVALID_CONFIG_PARAMETER; + import java.util.HashMap; import java.util.Map; import java.util.Properties; @@ -64,6 +66,10 @@ public class ParameterProvider { public static final Constants.BdecParquetCompression BDEC_PARQUET_COMPRESSION_ALGORITHM_DEFAULT = Constants.BdecParquetCompression.GZIP; + /* Iceberg mode parameters: When streaming to Iceberg mode, different default parameters are required because it generates Parquet files instead of BDEC files. */ + public static final int MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT = + 1; // 1 parquet file per blob + /* Parameter that enables using internal Parquet buffers for buffering of rows before serializing. It reduces memory consumption compared to using Java Objects for buffering.*/ public static final boolean ENABLE_PARQUET_INTERNAL_BUFFERING_DEFAULT = false; @@ -80,14 +86,17 @@ public class ParameterProvider { * * @param parameterOverrides Map of parameter name to value * @param props Properties from profile file + * @param isIcebergMode If the provided parameters need to be verified and modified to meet + * Iceberg mode */ - public ParameterProvider(Map parameterOverrides, Properties props) { - this.setParameterMap(parameterOverrides, props); + public ParameterProvider( + Map parameterOverrides, Properties props, boolean isIcebergMode) { + this.setParameterMap(parameterOverrides, props, isIcebergMode); } /** Empty constructor for tests */ - public ParameterProvider() { - this(null, null); + public ParameterProvider(boolean isIcebergMode) { + this(null, null, isIcebergMode); } private void updateValue( @@ -99,6 +108,8 @@ private void updateValue( this.parameterMap.put(key, parameterOverrides.getOrDefault(key, defaultValue)); } else if (props != null) { this.parameterMap.put(key, props.getOrDefault(key, defaultValue)); + } else { + this.parameterMap.put(key, defaultValue); } } @@ -107,8 +118,11 @@ private void updateValue( * * @param parameterOverrides Map of parameter name -> value * @param props Properties file provided to client constructor + * @param isIcebergMode If the provided parameters need to be verified and modified to meet + * Iceberg mode */ - private void setParameterMap(Map parameterOverrides, Properties props) { + private void setParameterMap( + Map parameterOverrides, Properties props, boolean isIcebergMode) { // BUFFER_FLUSH_INTERVAL_IN_MILLIS is deprecated and disallowed if ((parameterOverrides != null && parameterOverrides.containsKey(BUFFER_FLUSH_INTERVAL_IN_MILLIS)) @@ -179,7 +193,9 @@ private void setParameterMap(Map parameterOverrides, Properties this.updateValue( MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST, - MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT, + isIcebergMode + ? MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT + : MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT, parameterOverrides, props); @@ -188,6 +204,13 @@ private void setParameterMap(Map parameterOverrides, Properties BDEC_PARQUET_COMPRESSION_ALGORITHM_DEFAULT, parameterOverrides, props); + + // Required parameter override for Iceberg mode + if (isIcebergMode) { + icebergModeValidation( + MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST, + MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT); + } } /** @return Longest interval in milliseconds between buffer flushes */ @@ -411,4 +434,14 @@ public Constants.BdecParquetCompression getBdecParquetCompressionAlgorithm() { public String toString() { return "ParameterProvider{" + "parameterMap=" + parameterMap + '}'; } + + private void icebergModeValidation(String key, Object expected) { + Object val = this.parameterMap.get(key); + if (!val.equals(expected)) { + throw new SFException( + INVALID_CONFIG_PARAMETER, + String.format( + "The value %s for %s is invalid in Iceberg mode, should be %s.", val, key, expected)); + } + } } diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/ChannelCacheTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/ChannelCacheTest.java index 947908ef9..7c6d797b4 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/ChannelCacheTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/ChannelCacheTest.java @@ -24,7 +24,7 @@ public class ChannelCacheTest { @Before public void setup() { cache = new ChannelCache<>(); - client = new SnowflakeStreamingIngestClientInternal<>("client"); + client = new SnowflakeStreamingIngestClientInternal<>("client", false); channel1 = new SnowflakeStreamingIngestChannelInternal<>( "channel1", diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java index f200c7177..e66306d1f 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java @@ -7,6 +7,7 @@ import static net.snowflake.ingest.utils.Constants.BLOB_NO_HEADER; import static net.snowflake.ingest.utils.Constants.BLOB_TAG_SIZE_IN_BYTES; import static net.snowflake.ingest.utils.Constants.BLOB_VERSION_SIZE_IN_BYTES; +import static net.snowflake.ingest.utils.ParameterProvider.MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT; import static net.snowflake.ingest.utils.ParameterProvider.MAX_CHUNK_SIZE_IN_BYTES_DEFAULT; import com.codahale.metrics.Histogram; @@ -48,11 +49,22 @@ import net.snowflake.ingest.utils.SFException; import org.junit.Assert; import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; import org.mockito.Mockito; +@RunWith(Parameterized.class) public class FlushServiceTest { + + @Parameterized.Parameters(name = "isIcebergMode: {0}") + public static Object[] isIcebergMode() { + return new Object[] {false, true}; + } + + @Parameterized.Parameter public static boolean isIcebergMode; + public FlushServiceTest() { this.testContextFactory = ParquetTestContext.createFactory(); } @@ -86,7 +98,7 @@ private abstract static class TestContext implements AutoCloseable { TestContext() { stage = Mockito.mock(StreamingIngestStage.class); Mockito.when(stage.getClientPrefix()).thenReturn("client_prefix"); - parameterProvider = new ParameterProvider(); + parameterProvider = new ParameterProvider(isIcebergMode); client = Mockito.mock(SnowflakeStreamingIngestClientInternal.class); Mockito.when(client.getParameterProvider()).thenReturn(parameterProvider); channelCache = new ChannelCache<>(); @@ -586,7 +598,9 @@ public void runTestBlobSplitDueToNumberOfChunks(int numberOfRows) throws Excepti Math.ceil( (double) numberOfRows / channelsPerTable - / ParameterProvider.MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT); + / (isIcebergMode + ? MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT + : ParameterProvider.MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT)); final TestContext>> testContext = testContextFactory.create(); @@ -861,7 +875,7 @@ public void testInvalidateChannels() { // Create a new Client in order to not interfere with other tests SnowflakeStreamingIngestClientInternal client = Mockito.mock(SnowflakeStreamingIngestClientInternal.class); - ParameterProvider parameterProvider = new ParameterProvider(); + ParameterProvider parameterProvider = new ParameterProvider(isIcebergMode); ChannelCache channelCache = new ChannelCache<>(); Mockito.when(client.getChannelCache()).thenReturn(channelCache); Mockito.when(client.getParameterProvider()).thenReturn(parameterProvider); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/OAuthBasicTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/OAuthBasicTest.java index 6351f267d..0c0c5bb85 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/OAuthBasicTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/OAuthBasicTest.java @@ -114,7 +114,7 @@ public void testCreateOAuthClient() throws Exception { @Test public void testSetRefreshToken() throws Exception { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("TEST_CLIENT"); + new SnowflakeStreamingIngestClientInternal<>("TEST_CLIENT", false); MockOAuthClient mockOAuthClient = new MockOAuthClient(); OAuthManager oAuthManager = diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java index 86cece9c7..fb73aaa66 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java @@ -1,12 +1,16 @@ package net.snowflake.ingest.streaming.internal; +import static net.snowflake.ingest.utils.ParameterProvider.MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT; + import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Properties; import net.snowflake.ingest.utils.Constants; +import net.snowflake.ingest.utils.ErrorCode; import net.snowflake.ingest.utils.ParameterProvider; +import net.snowflake.ingest.utils.SFException; import org.junit.Assert; import org.junit.Test; @@ -31,7 +35,7 @@ private Map getStartingParameterMap() { public void withValuesSet() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals(1000L, parameterProvider.getCachedMaxClientLagInMs()); Assert.assertEquals(4L, parameterProvider.getBufferFlushCheckIntervalInMs()); @@ -54,7 +58,7 @@ public void withNullProps() { parameterMap.put(ParameterProvider.BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS, 4L); parameterMap.put(ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE, 6); parameterMap.put(ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_BYTES, 1024); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, null); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, null, false); Assert.assertEquals(3000, parameterProvider.getCachedMaxClientLagInMs()); Assert.assertEquals(4, parameterProvider.getBufferFlushCheckIntervalInMs()); @@ -72,7 +76,7 @@ public void withNullParameterMap() { props.put(ParameterProvider.BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS, 4L); props.put(ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE, 6); props.put(ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_BYTES, 1024); - ParameterProvider parameterProvider = new ParameterProvider(null, props); + ParameterProvider parameterProvider = new ParameterProvider(null, props, false); Assert.assertEquals(3000, parameterProvider.getCachedMaxClientLagInMs()); Assert.assertEquals(4, parameterProvider.getBufferFlushCheckIntervalInMs()); @@ -85,7 +89,7 @@ public void withNullParameterMap() { @Test public void withNullInputs() { - ParameterProvider parameterProvider = new ParameterProvider(null, null); + ParameterProvider parameterProvider = new ParameterProvider(null, null, false); Assert.assertEquals( ParameterProvider.MAX_CLIENT_LAG_DEFAULT, parameterProvider.getCachedMaxClientLagInMs()); @@ -105,7 +109,7 @@ public void withNullInputs() { @Test public void withDefaultValues() { - ParameterProvider parameterProvider = new ParameterProvider(); + ParameterProvider parameterProvider = new ParameterProvider(false); Assert.assertEquals( ParameterProvider.MAX_CLIENT_LAG_DEFAULT, parameterProvider.getCachedMaxClientLagInMs()); @@ -137,12 +141,49 @@ public void withDefaultValues() { parameterProvider.getBdecParquetCompressionAlgorithm()); } + @Test + public void withIcebergDefaultValues() { + ParameterProvider parameterProvider = new ParameterProvider(true); + + Assert.assertEquals( + ParameterProvider.MAX_CLIENT_LAG_DEFAULT, parameterProvider.getCachedMaxClientLagInMs()); + Assert.assertEquals( + ParameterProvider.BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS_DEFAULT, + parameterProvider.getBufferFlushCheckIntervalInMs()); + Assert.assertEquals( + ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE_DEFAULT, + parameterProvider.getInsertThrottleThresholdInPercentage()); + Assert.assertEquals( + ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_BYTES_DEFAULT, + parameterProvider.getInsertThrottleThresholdInBytes()); + Assert.assertEquals( + ParameterProvider.INSERT_THROTTLE_INTERVAL_IN_MILLIS_DEFAULT, + parameterProvider.getInsertThrottleIntervalInMs()); + Assert.assertEquals( + ParameterProvider.IO_TIME_CPU_RATIO_DEFAULT, parameterProvider.getIOTimeCpuRatio()); + Assert.assertEquals( + ParameterProvider.BLOB_UPLOAD_MAX_RETRY_COUNT_DEFAULT, + parameterProvider.getBlobUploadMaxRetryCount()); + Assert.assertEquals( + ParameterProvider.MAX_MEMORY_LIMIT_IN_BYTES_DEFAULT, + parameterProvider.getMaxMemoryLimitInBytes()); + Assert.assertEquals( + ParameterProvider.MAX_CHANNEL_SIZE_IN_BYTES_DEFAULT, + parameterProvider.getMaxChannelSizeInBytes()); + Assert.assertEquals( + ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM_DEFAULT, + parameterProvider.getBdecParquetCompressionAlgorithm()); + Assert.assertEquals( + MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT, + parameterProvider.getMaxChunksInBlobAndRegistrationRequest()); + } + @Test public void testMaxClientLagEnabled() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "2 second"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals(2000, parameterProvider.getCachedMaxClientLagInMs()); // call again to trigger caching logic Assert.assertEquals(2000, parameterProvider.getCachedMaxClientLagInMs()); @@ -153,7 +194,7 @@ public void testMaxClientLagEnabledPluralTimeUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "2 seconds"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals(2000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -162,7 +203,7 @@ public void testMaxClientLagEnabledMinuteTimeUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "1 minute"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals(60000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -171,7 +212,7 @@ public void testMaxClientLagEnabledMinuteTimeUnitPluralTimeUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "2 minutes"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals(120000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -179,7 +220,7 @@ public void testMaxClientLagEnabledMinuteTimeUnitPluralTimeUnit() { public void testMaxClientLagEnabledDefaultValue() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals( ParameterProvider.MAX_CLIENT_LAG_DEFAULT, parameterProvider.getCachedMaxClientLagInMs()); } @@ -189,7 +230,7 @@ public void testMaxClientLagEnabledDefaultUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "3000"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals(3000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -198,7 +239,7 @@ public void testMaxClientLagEnabledLongInput() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, 3000L); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals(3000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -207,7 +248,7 @@ public void testMaxClientLagEnabledMissingUnitTimeUnitSupplied() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, " year"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -221,7 +262,7 @@ public void testMaxClientLagEnabledInvalidTimeUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "1 year"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -235,7 +276,7 @@ public void testMaxClientLagEnabledInvalidUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "banana minute"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -249,7 +290,7 @@ public void testMaxClientLagEnabledThresholdBelow() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "0 second"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -263,7 +304,7 @@ public void testMaxClientLagEnabledThresholdAbove() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "11 minutes"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -277,7 +318,7 @@ public void testMaxClientLagEnableEmptyInput() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, ""); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -291,8 +332,20 @@ public void testMaxChunksInBlobAndRegistrationRequest() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put("max_chunks_in_blob_and_registration_request", 1); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + Assert.assertEquals(1, parameterProvider.getMaxChunksInBlobAndRegistrationRequest()); + + parameterProvider = new ParameterProvider(parameterMap, prop, true); Assert.assertEquals(1, parameterProvider.getMaxChunksInBlobAndRegistrationRequest()); + + SFException e = + Assert.assertThrows( + SFException.class, + () -> { + parameterMap.put("max_chunks_in_blob_and_registration_request", 100); + new ParameterProvider(parameterMap, prop, true); + }); + Assert.assertEquals(e.getVendorCode(), ErrorCode.INVALID_CONFIG_PARAMETER.getMessageCode()); } @Test @@ -303,7 +356,7 @@ public void testValidCompressionAlgorithmsAndWithUppercaseLowerCase() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM, v); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals( Constants.BdecParquetCompression.GZIP, parameterProvider.getBdecParquetCompressionAlgorithm()); @@ -314,7 +367,7 @@ public void testValidCompressionAlgorithmsAndWithUppercaseLowerCase() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM, v); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); Assert.assertEquals( Constants.BdecParquetCompression.ZSTD, parameterProvider.getBdecParquetCompressionAlgorithm()); @@ -326,7 +379,7 @@ public void testInvalidCompressionAlgorithm() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM, "invalid_comp"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); try { parameterProvider.getBdecParquetCompressionAlgorithm(); Assert.fail("Should not have succeeded"); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java index 37eb5f96e..4eaea15a4 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java @@ -14,8 +14,17 @@ import org.junit.Assert; import org.junit.Ignore; import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +@RunWith(Parameterized.class) public class RegisterServiceTest { + @Parameterized.Parameters(name = "isIcebergMode: {0}") + public static Object[] isIcebergMode() { + return new Object[] {false, true}; + } + + @Parameterized.Parameter public boolean isIcebergMode; @Test public void testRegisterService() throws ExecutionException, InterruptedException { @@ -45,7 +54,7 @@ public void testRegisterService() throws ExecutionException, InterruptedExceptio @Test public void testRegisterServiceTimeoutException() throws Exception { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client"); + new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); RegisterService rs = new RegisterService<>(client, true); Pair, CompletableFuture> blobFuture1 = @@ -73,7 +82,7 @@ public void testRegisterServiceTimeoutException() throws Exception { @Test public void testRegisterServiceTimeoutException_testRetries() throws Exception { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client"); + new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); RegisterService rs = new RegisterService<>(client, true); Pair, CompletableFuture> blobFuture1 = @@ -107,7 +116,7 @@ public void testRegisterServiceTimeoutException_testRetries() throws Exception { @Test public void testRegisterServiceNonTimeoutException() { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client"); + new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); RegisterService rs = new RegisterService<>(client, true); CompletableFuture future = new CompletableFuture<>(); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java index 87e3f8f11..5beb0662f 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java @@ -44,8 +44,11 @@ import net.snowflake.ingest.utils.Utils; import org.junit.Assert; import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; import org.mockito.Mockito; +@RunWith(Parameterized.class) public class SnowflakeStreamingIngestChannelTest { /** @@ -72,6 +75,13 @@ public long getFreeMemory() { } } + @Parameterized.Parameters(name = "isIcebergMode: {0}") + public static Object[] isIcebergMode() { + return new Object[] {false, true}; + } + + @Parameterized.Parameter public boolean isIcebergMode; + @Test public void testChannelFactoryNullFields() { String name = "CHANNEL"; @@ -81,7 +91,7 @@ public void testChannelFactoryNullFields() { long channelSequencer = 0L; long rowSequencer = 0L; SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client"); + new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); Object[] fields = new Object[] { @@ -123,7 +133,7 @@ public void testChannelFactorySuccess() { long rowSequencer = 0L; SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client"); + new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); SnowflakeStreamingIngestChannelInternal channel = SnowflakeStreamingIngestChannelFactory.builder(name) @@ -158,7 +168,7 @@ public void testChannelFactorySuccess() { @Test public void testChannelValid() { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client"); + new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -208,7 +218,7 @@ public void testChannelValid() { @Test public void testChannelClose() { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client"); + new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -347,6 +357,7 @@ public void testOpenChannelErrorResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -417,6 +428,7 @@ public void testOpenChannelSnowflakeInternalErrorResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -499,6 +511,7 @@ public void testOpenChannelSuccessResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -537,7 +550,9 @@ public void testOpenChannelSuccessResponse() throws Exception { @Test public void testInsertRow() { SnowflakeStreamingIngestClientInternal client; - client = new SnowflakeStreamingIngestClientInternal("client_PARQUET"); + client = + new SnowflakeStreamingIngestClientInternal( + "client_PARQUET", isIcebergMode); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -621,7 +636,8 @@ public void testInsertTooLargeRow() { schema.forEach(x -> row.put(x.getName(), byteArrayOneMb)); SnowflakeStreamingIngestClientInternal client; - client = new SnowflakeStreamingIngestClientInternal("test_client"); + client = + new SnowflakeStreamingIngestClientInternal("test_client", isIcebergMode); // Test channel with on error CONTINUE SnowflakeStreamingIngestChannelInternal channel = @@ -705,7 +721,7 @@ public void testInsertRowThrottling() { memoryInfoProvider.maxMemory = maxMemory; SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client"); + new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -721,7 +737,7 @@ public void testInsertRowThrottling() { OpenChannelRequest.OnErrorOption.CONTINUE, UTC); - ParameterProvider parameterProvider = new ParameterProvider(); + ParameterProvider parameterProvider = new ParameterProvider(isIcebergMode); memoryInfoProvider.freeMemory = maxMemory * (parameterProvider.getInsertThrottleThresholdInPercentage() - 1) / 100; @@ -751,7 +767,7 @@ public void testInsertRowThrottling() { @Test public void testFlush() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -787,7 +803,7 @@ public void testFlush() throws Exception { @Test public void testClose() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); SnowflakeStreamingIngestChannel channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -821,7 +837,7 @@ public void testClose() throws Exception { @Test public void testDropOnClose() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -858,7 +874,7 @@ public void testDropOnClose() throws Exception { @Test public void testDropOnCloseInvalidChannel() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -891,7 +907,7 @@ public void testDropOnCloseInvalidChannel() throws Exception { public void testGetLatestCommittedOffsetToken() { String offsetToken = "10"; SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); SnowflakeStreamingIngestChannel channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java index 1693e1520..4e66d8a15 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java @@ -66,9 +66,12 @@ import org.junit.Before; import org.junit.Ignore; import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; import org.mockito.ArgumentMatchers; import org.mockito.Mockito; +@RunWith(Parameterized.class) public class SnowflakeStreamingIngestClientTest { private static final ObjectMapper objectMapper = new ObjectMapper(); @@ -79,6 +82,14 @@ public class SnowflakeStreamingIngestClientTest { SnowflakeStreamingIngestChannelInternal channel3; SnowflakeStreamingIngestChannelInternal channel4; + // TODO: Add IcebergMode = True after Streaming to Iceberg is supported. + @Parameterized.Parameters(name = "isIcebergMode: {0}") + public static Object[] isIcebergMode() { + return new Object[] {false}; + } + + @Parameterized.Parameter public boolean isIcebergMode; + @Before public void setup() { objectMapper.setVisibility(PropertyAccessor.GETTER, JsonAutoDetect.Visibility.ANY); @@ -166,6 +177,7 @@ public void testConstructorParameters() throws Exception { SnowflakeStreamingIngestClientFactory.builder("client") .setProperties(prop) .setParameterOverrides(parameterMap) + .setIsIceberg(isIcebergMode) .setIsTestMode(true) .build(); @@ -193,6 +205,7 @@ public void testClientFactoryWithJmxMetrics() throws Exception { .setProperties(prop) .setParameterOverrides( Collections.singletonMap(ENABLE_SNOWPIPE_STREAMING_METRICS, true)) + .setIsIceberg(isIcebergMode) .build(); Assert.assertEquals("client", client.getName()); @@ -344,6 +357,7 @@ public void testGetChannelsStatusWithRequest() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -404,6 +418,7 @@ public void testDropChannel() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -449,6 +464,7 @@ public void testGetChannelsStatusWithRequestError() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -677,6 +693,7 @@ public void testGetRetryBlobs() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -718,6 +735,7 @@ public void testRegisterBlobErrorResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -765,6 +783,7 @@ public void testRegisterBlobSnowflakeInternalErrorResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -821,6 +840,7 @@ public void testRegisterBlobSuccessResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -905,6 +925,7 @@ public void testRegisterBlobsRetries() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -934,6 +955,7 @@ public void testRegisterBlobChunkLimit() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null)); @@ -1106,6 +1128,7 @@ public void testRegisterBlobsRetriesSucceeds() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -1180,6 +1203,7 @@ public void testRegisterBlobResponseWithInvalidChannel() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); @@ -1230,7 +1254,7 @@ public void testRegisterBlobResponseWithInvalidChannel() throws Exception { @Test public void testFlush() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); ChannelsStatusResponse response = new ChannelsStatusResponse(); response.setStatusCode(0L); response.setMessage("Success"); @@ -1252,7 +1276,7 @@ public void testFlush() throws Exception { @Test public void testClose() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); ChannelsStatusResponse response = new ChannelsStatusResponse(); response.setStatusCode(0L); response.setMessage("Success"); @@ -1286,7 +1310,7 @@ public void testClose() throws Exception { @Test public void testCloseWithError() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); CompletableFuture future = new CompletableFuture<>(); future.completeExceptionally(new Exception("Simulating Error")); @@ -1324,7 +1348,7 @@ public void testCloseWithError() throws Exception { @Test public void testVerifyChannelsAreFullyCommittedSuccess() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel1", @@ -1372,6 +1396,7 @@ public void testFlushServiceException() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, parameterMap); @@ -1408,6 +1433,7 @@ public void testGetLatestCommittedOffsetTokens() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, + isIcebergMode, true, requestBuilder, null); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStageTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStageTest.java index 1ba9f98df..d12c6231c 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStageTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStageTest.java @@ -276,7 +276,7 @@ public void testRefreshSnowflakeMetadataRemote() throws Exception { Mockito.when(mockResponse.getEntity()).thenReturn(entity); Mockito.when(mockClient.execute(Mockito.any())).thenReturn(mockResponse); - ParameterProvider parameterProvider = new ParameterProvider(); + ParameterProvider parameterProvider = new ParameterProvider(false); StreamingIngestStage stage = new StreamingIngestStage(true, "role", mockClient, mockBuilder, "clientName", 1); From a9aa6820909387e7ee81d5783aa18b11f9f0a193 Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Fri, 21 Jun 2024 15:18:53 -0700 Subject: [PATCH 2/9] SNOW-1483230 Disable blob encryption / Add class for storage metadata & configure response (#773) --- .../streaming/internal/BlobBuilder.java | 64 +++++---- .../streaming/internal/ConfigureResponse.java | 66 +++++++++ .../streaming/internal/FileLocationInfo.java | 132 ++++++++++++++++++ .../streaming/internal/FlushService.java | 9 +- .../internal/InternalParameterProvider.java | 20 +++ .../internal/OpenChannelResponse.java | 10 ++ ...nowflakeStreamingIngestClientInternal.java | 17 ++- .../internal/StreamingIngestStage.java | 94 ++++++------- .../ingest/utils/ParameterProvider.java | 122 +++++++++------- .../streaming/internal/FlushServiceTest.java | 7 + .../InternalParameterProviderTest.java | 26 ++++ .../internal/ParameterProviderTest.java | 3 +- 12 files changed, 442 insertions(+), 128 deletions(-) create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ConfigureResponse.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/FileLocationInfo.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/InternalParameterProvider.java create mode 100644 src/test/java/net/snowflake/ingest/streaming/internal/InternalParameterProviderTest.java diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java b/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java index b88090e01..75a623f69 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java @@ -61,10 +61,14 @@ class BlobBuilder { * @param blobData All the data for one blob. Assumes that all ChannelData in the inner List * belongs to the same table. Will error if this is not the case * @param bdecVersion version of blob + * @param encrypt If the output chunk is encrypted or not * @return {@link Blob} data */ static Blob constructBlobAndMetadata( - String filePath, List>> blobData, Constants.BdecVersion bdecVersion) + String filePath, + List>> blobData, + Constants.BdecVersion bdecVersion, + boolean encrypt) throws IOException, NoSuchPaddingException, NoSuchAlgorithmException, InvalidAlgorithmParameterException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException { @@ -83,25 +87,32 @@ static Blob constructBlobAndMetadata( flusher.serialize(channelsDataPerTable, filePath); if (!serializedChunk.channelsMetadataList.isEmpty()) { - ByteArrayOutputStream chunkData = serializedChunk.chunkData; - Pair paddedChunk = - padChunk(chunkData, Constants.ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES); - byte[] paddedChunkData = paddedChunk.getFirst(); - int paddedChunkLength = paddedChunk.getSecond(); + byte[] compressedChunkData; + int chunkLength; - // Encrypt the compressed chunk data, the encryption key is derived using the key from - // server with the full blob path. - // We need to maintain IV as a block counter for the whole file, even interleaved, - // to align with decryption on the Snowflake query path. - // TODO: address alignment for the header SNOW-557866 - long iv = curDataSize / Constants.ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES; - byte[] encryptedCompressedChunkData = - Cryptor.encrypt( - paddedChunkData, firstChannelFlushContext.getEncryptionKey(), filePath, iv); + if (encrypt) { + Pair paddedChunk = + padChunk(serializedChunk.chunkData, Constants.ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES); + byte[] paddedChunkData = paddedChunk.getFirst(); + chunkLength = paddedChunk.getSecond(); + + // Encrypt the compressed chunk data, the encryption key is derived using the key from + // server with the full blob path. + // We need to maintain IV as a block counter for the whole file, even interleaved, + // to align with decryption on the Snowflake query path. + // TODO: address alignment for the header SNOW-557866 + long iv = curDataSize / Constants.ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES; + compressedChunkData = + Cryptor.encrypt( + paddedChunkData, firstChannelFlushContext.getEncryptionKey(), filePath, iv); + } else { + compressedChunkData = serializedChunk.chunkData.toByteArray(); + chunkLength = compressedChunkData.length; + } // Compute the md5 of the chunk data - String md5 = computeMD5(encryptedCompressedChunkData, paddedChunkLength); - int encryptedCompressedChunkDataSize = encryptedCompressedChunkData.length; + String md5 = computeMD5(compressedChunkData, chunkLength); + int compressedChunkDataSize = compressedChunkData.length; // Create chunk metadata long startOffset = curDataSize; @@ -111,9 +122,9 @@ static Blob constructBlobAndMetadata( // The start offset will be updated later in BlobBuilder#build to include the blob // header .setChunkStartOffset(startOffset) - // The paddedChunkLength is used because it is the actual data size used for + // The chunkLength is used because it is the actual data size used for // decompression and md5 calculation on server side. - .setChunkLength(paddedChunkLength) + .setChunkLength(chunkLength) .setUncompressedChunkLength((int) serializedChunk.chunkEstimatedUncompressedSize) .setChannelList(serializedChunk.channelsMetadataList) .setChunkMD5(md5) @@ -127,21 +138,22 @@ static Blob constructBlobAndMetadata( // Add chunk metadata and data to the list chunksMetadataList.add(chunkMetadata); - chunksDataList.add(encryptedCompressedChunkData); - curDataSize += encryptedCompressedChunkDataSize; - crc.update(encryptedCompressedChunkData, 0, encryptedCompressedChunkDataSize); + chunksDataList.add(compressedChunkData); + curDataSize += compressedChunkDataSize; + crc.update(compressedChunkData, 0, compressedChunkDataSize); logger.logInfo( "Finish building chunk in blob={}, table={}, rowCount={}, startOffset={}," - + " estimatedUncompressedSize={}, paddedChunkLength={}, encryptedCompressedSize={}," - + " bdecVersion={}", + + " estimatedUncompressedSize={}, chunkLength={}, compressedSize={}," + + " encryption={}, bdecVersion={}", filePath, firstChannelFlushContext.getFullyQualifiedTableName(), serializedChunk.rowCount, startOffset, serializedChunk.chunkEstimatedUncompressedSize, - paddedChunkLength, - encryptedCompressedChunkDataSize, + chunkLength, + compressedChunkDataSize, + encrypt, bdecVersion); } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureResponse.java new file mode 100644 index 000000000..6fe73fbd5 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureResponse.java @@ -0,0 +1,66 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import com.fasterxml.jackson.annotation.JsonProperty; + +/** Class used to deserialize responses from configure endpoint */ +class ConfigureResponse extends StreamingIngestResponse { + @JsonProperty("prefix") + private String prefix; + + @JsonProperty("status_code") + private Long statusCode; + + @JsonProperty("message") + private String message; + + @JsonProperty("stage_location") + private FileLocationInfo stageLocation; + + @JsonProperty("deployment_id") + private Long deploymentId; + + String getPrefix() { + return prefix; + } + + void setPrefix(String prefix) { + this.prefix = prefix; + } + + @Override + Long getStatusCode() { + return statusCode; + } + + void setStatusCode(Long statusCode) { + this.statusCode = statusCode; + } + + String getMessage() { + return message; + } + + void setMessage(String message) { + this.message = message; + } + + FileLocationInfo getStageLocation() { + return stageLocation; + } + + void setStageLocation(FileLocationInfo stageLocation) { + this.stageLocation = stageLocation; + } + + Long getDeploymentId() { + return deploymentId; + } + + void setDeploymentId(Long deploymentId) { + this.deploymentId = deploymentId; + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FileLocationInfo.java b/src/main/java/net/snowflake/ingest/streaming/internal/FileLocationInfo.java new file mode 100644 index 000000000..4af5862c4 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/FileLocationInfo.java @@ -0,0 +1,132 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.Map; + +/** Class used to deserialized volume information response by server */ +class FileLocationInfo { + @JsonProperty("locationType") + private String locationType; // The stage type + + @JsonProperty("location") + private String location; // The container or bucket + + @JsonProperty("path") + private String path; // path of the target file + + @JsonProperty("creds") + private Map credentials; // the credentials required for the stage + + @JsonProperty("region") + private String region; // AWS/S3/GCS region (S3/GCS only) + + @JsonProperty("endPoint") + private String endPoint; // The Azure Storage endpoint (Azure only) + + @JsonProperty("storageAccount") + private String storageAccount; // The Azure Storage account (Azure only) + + @JsonProperty("presignedUrl") + private String presignedUrl; // GCS gives us back a presigned URL instead of a cred + + @JsonProperty("isClientSideEncrypted") + private boolean isClientSideEncrypted; // whether to encrypt/decrypt files on the stage + + @JsonProperty("useS3RegionalUrl") + private boolean useS3RegionalUrl; // whether to use s3 regional URL (AWS Only) + + @JsonProperty("volumeHash") + private String volumeHash; // a unique id for volume assigned by server + + String getLocationType() { + return locationType; + } + + void setLocationType(String locationType) { + this.locationType = locationType; + } + + String getLocation() { + return location; + } + + void setLocation(String location) { + this.location = location; + } + + String getPath() { + return path; + } + + void setPath(String path) { + this.path = path; + } + + Map getCredentials() { + return credentials; + } + + void setCredentials(Map credentials) { + this.credentials = credentials; + } + + String getRegion() { + return region; + } + + void setRegion(String region) { + this.region = region; + } + + String getEndPoint() { + return endPoint; + } + + void setEndPoint(String endPoint) { + this.endPoint = endPoint; + } + + String getStorageAccount() { + return storageAccount; + } + + void setStorageAccount(String storageAccount) { + this.storageAccount = storageAccount; + } + + String getPresignedUrl() { + return presignedUrl; + } + + void setPresignedUrl(String presignedUrl) { + this.presignedUrl = presignedUrl; + } + + boolean getIsClientSideEncrypted() { + return this.isClientSideEncrypted; + } + + void setIsClientSideEncrypted(boolean isClientSideEncrypted) { + this.isClientSideEncrypted = isClientSideEncrypted; + } + + boolean getUseS3RegionalUrl() { + return this.useS3RegionalUrl; + } + + void setUseS3RegionalUrl(boolean useS3RegionalUrl) { + this.useS3RegionalUrl = useS3RegionalUrl; + } + + String getVolumeHash() { + return this.volumeHash; + } + + void setVolumeHash(String volumeHash) { + this.volumeHash = volumeHash; + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java index f08196477..05d72433d 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming.internal; @@ -551,7 +551,12 @@ BlobMetadata buildAndUpload(String blobPath, List>> blobData Timer.Context buildContext = Utils.createTimerContext(this.owningClient.buildLatency); // Construct the blob along with the metadata of the blob - BlobBuilder.Blob blob = BlobBuilder.constructBlobAndMetadata(blobPath, blobData, bdecVersion); + BlobBuilder.Blob blob = + BlobBuilder.constructBlobAndMetadata( + blobPath, + blobData, + bdecVersion, + this.owningClient.getInternalParameterProvider().getEnableChunkEncryption()); blob.blobStats.setBuildDurationMs(buildContext); diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/InternalParameterProvider.java b/src/main/java/net/snowflake/ingest/streaming/internal/InternalParameterProvider.java new file mode 100644 index 000000000..46a99b671 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/InternalParameterProvider.java @@ -0,0 +1,20 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +/** A class to provide non-configurable constants depends on Iceberg or non-Iceberg mode */ +class InternalParameterProvider { + private final boolean isIcebergMode; + + InternalParameterProvider(boolean isIcebergMode) { + this.isIcebergMode = isIcebergMode; + } + + boolean getEnableChunkEncryption() { + // When in Iceberg mode, chunk encryption is disabled. Otherwise, it is enabled. Since Iceberg + // mode does not need client-side encryption. + return !isIcebergMode; + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelResponse.java index beb0eb0c4..a06545819 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelResponse.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelResponse.java @@ -21,6 +21,7 @@ class OpenChannelResponse extends StreamingIngestResponse { private List tableColumns; private String encryptionKey; private Long encryptionKeyId; + private FileLocationInfo stageLocation; @JsonProperty("status_code") void setStatusCode(Long statusCode) { @@ -130,4 +131,13 @@ void setEncryptionKeyId(Long encryptionKeyId) { Long getEncryptionKeyId() { return this.encryptionKeyId; } + + @JsonProperty("stage_location") + void setStageLocation(FileLocationInfo stageLocation) { + this.stageLocation = stageLocation; + } + + FileLocationInfo getStageLocation() { + return this.stageLocation; + } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java index 13cc673ff..1de473da6 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming.internal; @@ -103,6 +103,9 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea // Provides constant values that can be set by constructor private final ParameterProvider parameterProvider; + // Provides constant values which is determined by the Iceberg or non-Iceberg mode + private final InternalParameterProvider internalParameterProvider; + // Name of the client private final String name; @@ -170,6 +173,7 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea RequestBuilder requestBuilder, Map parameterOverrides) { this.parameterProvider = new ParameterProvider(parameterOverrides, prop, isIcebergMode); + this.internalParameterProvider = new InternalParameterProvider(isIcebergMode); this.name = name; String accountName = accountURL == null ? null : accountURL.getAccount(); @@ -337,6 +341,7 @@ public SnowflakeStreamingIngestChannelInternal openChannel(OpenChannelRequest payload.put("schema", request.getSchemaName()); payload.put("write_mode", Constants.WriteMode.CLOUD_STORAGE.name()); payload.put("role", this.role); + payload.put("is_iceberg", isIcebergMode); if (request.isOffsetTokenProvided()) { payload.put("offset_token", request.getOffsetToken()); } @@ -421,6 +426,7 @@ public void dropChannel(DropChannelRequest request) { payload.put("database", request.getDBName()); payload.put("schema", request.getSchemaName()); payload.put("role", this.role); + payload.put("is_iceberg", isIcebergMode); Long clientSequencer = null; if (request instanceof DropChannelVersionRequest) { clientSequencer = ((DropChannelVersionRequest) request).getClientSequencer(); @@ -964,6 +970,15 @@ ParameterProvider getParameterProvider() { return parameterProvider; } + /** + * Get InternalParameterProvider with internal parameters + * + * @return {@link InternalParameterProvider} used by the client + */ + InternalParameterProvider getInternalParameterProvider() { + return internalParameterProvider; + } + /** * Set refresh token, this method is for refresh token renewal without requiring to restart * client. This method only works when the authorization type is OAuth diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStage.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStage.java index e8e56f383..9752b311c 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStage.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStage.java @@ -11,6 +11,10 @@ import static net.snowflake.ingest.utils.HttpUtil.generateProxyPropertiesForJDBC; import static net.snowflake.ingest.utils.Utils.getStackTrace; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.common.annotations.VisibleForTesting; import java.io.ByteArrayInputStream; import java.io.File; @@ -22,7 +26,6 @@ import java.util.Optional; import java.util.Properties; import java.util.concurrent.TimeUnit; -import java.util.function.Function; import net.snowflake.client.core.OCSPMode; import net.snowflake.client.jdbc.SnowflakeFileTransferAgent; import net.snowflake.client.jdbc.SnowflakeFileTransferConfig; @@ -31,9 +34,6 @@ import net.snowflake.client.jdbc.cloud.storage.StageInfo; import net.snowflake.client.jdbc.internal.apache.commons.io.FileUtils; import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient; -import net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.JsonNode; -import net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.ObjectMapper; -import net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.node.ObjectNode; import net.snowflake.ingest.connection.IngestResponseException; import net.snowflake.ingest.connection.RequestBuilder; import net.snowflake.ingest.utils.ErrorCode; @@ -44,6 +44,13 @@ /** Handles uploading files to the Snowflake Streaming Ingest Stage */ class StreamingIngestStage { private static final ObjectMapper mapper = new ObjectMapper(); + + // Object mapper for parsing the client/configure response to Jackson version the same as + // jdbc.internal.fasterxml.jackson + private static final net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.ObjectMapper + parseConfigureResponseMapper = + new net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.ObjectMapper(); + private static final long REFRESH_THRESHOLD_IN_MS = TimeUnit.MILLISECONDS.convert(1, TimeUnit.MINUTES); @@ -104,7 +111,6 @@ state to record unknown age. this.clientName = clientName; this.proxyProperties = generateProxyPropertiesForJDBC(); this.maxUploadRetries = maxUploadRetries; - if (!isTestMode) { refreshSnowflakeMetadata(); } @@ -246,30 +252,25 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boole Map payload = new HashMap<>(); payload.put("role", this.role); - Map response = this.makeClientConfigureCall(payload); + ConfigureResponse response = this.makeClientConfigureCall(payload); - JsonNode responseNode = this.parseClientConfigureResponse(response); // Do not change the prefix everytime we have to refresh credentials if (Utils.isNullOrEmpty(this.clientPrefix)) { - this.clientPrefix = createClientPrefix(responseNode); + this.clientPrefix = createClientPrefix(response); } Utils.assertStringNotNullOrEmpty("client prefix", this.clientPrefix); - if (responseNode - .get("data") - .get("stageInfo") - .get("locationType") - .toString() + if (response + .getStageLocation() + .getLocationType() .replaceAll( "^[\"]|[\"]$", "") // Replace the first and last character if they're double quotes .equals(StageInfo.StageType.LOCAL_FS.name())) { this.fileTransferMetadataWithAge = new SnowflakeFileTransferMetadataWithAge( - responseNode - .get("data") - .get("stageInfo") - .get("location") - .toString() + response + .getStageLocation() + .getLocation() .replaceAll( "^[\"]|[\"]$", ""), // Replace the first and last character if they're double quotes @@ -278,7 +279,9 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boole this.fileTransferMetadataWithAge = new SnowflakeFileTransferMetadataWithAge( (SnowflakeFileTransferMetadataV1) - SnowflakeFileTransferAgent.getFileTransferMetadatas(responseNode).get(0), + SnowflakeFileTransferAgent.getFileTransferMetadatas( + parseClientConfigureResponse(response)) + .get(0), Optional.of(System.currentTimeMillis())); } return this.fileTransferMetadataWithAge; @@ -293,10 +296,10 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boole * @param response the client/configure response from the server * @return the client prefix. */ - private String createClientPrefix(final JsonNode response) { - final String prefix = response.get("prefix").textValue(); + private String createClientPrefix(final ConfigureResponse response) { + final String prefix = response.getPrefix(); final String deploymentId = - response.has("deployment_id") ? "_" + response.get("deployment_id").longValue() : ""; + response.getDeploymentId() != null ? "_" + response.getDeploymentId() : ""; return prefix + deploymentId; } @@ -312,33 +315,22 @@ SnowflakeFileTransferMetadataV1 fetchSignedURL(String fileName) Map payload = new HashMap<>(); payload.put("role", this.role); payload.put("file_name", fileName); - Map response = this.makeClientConfigureCall(payload); - - JsonNode responseNode = this.parseClientConfigureResponse(response); + ConfigureResponse response = this.makeClientConfigureCall(payload); SnowflakeFileTransferMetadataV1 metadata = (SnowflakeFileTransferMetadataV1) - SnowflakeFileTransferAgent.getFileTransferMetadatas(responseNode).get(0); + SnowflakeFileTransferAgent.getFileTransferMetadatas( + parseClientConfigureResponse(response)) + .get(0); // Transfer agent trims path for fileName metadata.setPresignedUrlFileName(fileName); return metadata; } - private static class MapStatusGetter implements Function { - public MapStatusGetter() {} - - public Long apply(T input) { - try { - return ((Integer) ((Map) input).get("status_code")).longValue(); - } catch (Exception e) { - throw new SFException(ErrorCode.INTERNAL_ERROR, "failed to get status_code from response"); - } - } - } - - private static final MapStatusGetter statusGetter = new MapStatusGetter(); - - private JsonNode parseClientConfigureResponse(Map response) { + private net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.JsonNode + parseClientConfigureResponse(ConfigureResponse response) + throws JsonProcessingException, + net.snowflake.client.jdbc.internal.fasterxml.jackson.core.JsonProcessingException { JsonNode responseNode = mapper.valueToTree(response); // Currently there are a few mismatches between the client/configure response and what @@ -350,28 +342,30 @@ private JsonNode parseClientConfigureResponse(Map response) { // JDBC expects this field which maps to presignedFileUrlName. We will set this later dataNode.putArray("src_locations").add("placeholder"); - return responseNode; + + // use String as intermediate object to avoid Jackson version mismatch + // TODO: SNOW-1493470 Align Jackson version + String responseString = mapper.writeValueAsString(responseNode); + return parseConfigureResponseMapper.readTree(responseString); } - private Map makeClientConfigureCall(Map payload) + private ConfigureResponse makeClientConfigureCall(Map payload) throws IOException { try { - Map response = + ConfigureResponse response = executeWithRetries( - Map.class, + ConfigureResponse.class, CLIENT_CONFIGURE_ENDPOINT, mapper.writeValueAsString(payload), "client configure", STREAMING_CLIENT_CONFIGURE, httpClient, - requestBuilder, - statusGetter); + requestBuilder); // Check for Snowflake specific response code - if (!response.get("status_code").equals((int) RESPONSE_SUCCESS)) { - throw new SFException( - ErrorCode.CLIENT_CONFIGURE_FAILURE, response.get("message").toString()); + if (response.getStatusCode() != RESPONSE_SUCCESS) { + throw new SFException(ErrorCode.CLIENT_CONFIGURE_FAILURE, response.getMessage()); } return response; } catch (IngestResponseException e) { diff --git a/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java b/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java index 33f791ca5..7d8cb230f 100644 --- a/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java +++ b/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java @@ -1,3 +1,7 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.utils; import static net.snowflake.ingest.utils.ErrorCode.INVALID_CONFIG_PARAMETER; @@ -57,6 +61,7 @@ public class ParameterProvider { // Lag related parameters public static final long MAX_CLIENT_LAG_DEFAULT = 1000; // 1 second + public static final long MAX_CLIENT_LAG_ICEBERG_MODE_DEFAULT = 30000; // 30 seconds static final long MAX_CLIENT_LAG_MS_MIN = TimeUnit.SECONDS.toMillis(1); static final long MAX_CLIENT_LAG_MS_MAX = TimeUnit.MINUTES.toMillis(10); @@ -67,8 +72,7 @@ public class ParameterProvider { Constants.BdecParquetCompression.GZIP; /* Iceberg mode parameters: When streaming to Iceberg mode, different default parameters are required because it generates Parquet files instead of BDEC files. */ - public static final int MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT = - 1; // 1 parquet file per blob + public static final int MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT = 1; /* Parameter that enables using internal Parquet buffers for buffering of rows before serializing. It reduces memory consumption compared to using Java Objects for buffering.*/ @@ -99,8 +103,12 @@ public ParameterProvider(boolean isIcebergMode) { this(null, null, isIcebergMode); } - private void updateValue( - String key, Object defaultValue, Map parameterOverrides, Properties props) { + private void checkAndUpdate( + String key, + Object defaultValue, + Map parameterOverrides, + Properties props, + boolean enforceDefault) { if (parameterOverrides != null && props != null) { this.parameterMap.put( key, parameterOverrides.getOrDefault(key, props.getOrDefault(key, defaultValue))); @@ -111,6 +119,17 @@ private void updateValue( } else { this.parameterMap.put(key, defaultValue); } + + if (enforceDefault) { + if (!this.parameterMap.getOrDefault(key, defaultValue).equals(defaultValue)) { + throw new SFException( + INVALID_CONFIG_PARAMETER, + String.format( + "The value %s for %s is not configurable, should be %s.", + this.parameterMap.get(key), key, defaultValue)); + } + this.parameterMap.put(key, defaultValue); + } } /** @@ -133,84 +152,101 @@ private void setParameterMap( BUFFER_FLUSH_INTERVAL_IN_MILLIS, MAX_CLIENT_LAG)); } - this.updateValue( + this.checkAndUpdate( BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS, BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS_DEFAULT, parameterOverrides, - props); + props, + false); - this.updateValue( + this.checkAndUpdate( INSERT_THROTTLE_INTERVAL_IN_MILLIS, INSERT_THROTTLE_INTERVAL_IN_MILLIS_DEFAULT, parameterOverrides, - props); + props, + false); - this.updateValue( + this.checkAndUpdate( INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE, INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE_DEFAULT, parameterOverrides, - props); + props, + false); - this.updateValue( + this.checkAndUpdate( INSERT_THROTTLE_THRESHOLD_IN_BYTES, INSERT_THROTTLE_THRESHOLD_IN_BYTES_DEFAULT, parameterOverrides, - props); + props, + false); - this.updateValue( + this.checkAndUpdate( ENABLE_SNOWPIPE_STREAMING_METRICS, SNOWPIPE_STREAMING_METRICS_DEFAULT, parameterOverrides, - props); + props, + false); - this.updateValue(BLOB_FORMAT_VERSION, BLOB_FORMAT_VERSION_DEFAULT, parameterOverrides, props); + this.checkAndUpdate( + BLOB_FORMAT_VERSION, BLOB_FORMAT_VERSION_DEFAULT, parameterOverrides, props, false); getBlobFormatVersion(); // to verify parsing the configured value - this.updateValue(IO_TIME_CPU_RATIO, IO_TIME_CPU_RATIO_DEFAULT, parameterOverrides, props); + this.checkAndUpdate( + IO_TIME_CPU_RATIO, IO_TIME_CPU_RATIO_DEFAULT, parameterOverrides, props, false); - this.updateValue( + this.checkAndUpdate( BLOB_UPLOAD_MAX_RETRY_COUNT, BLOB_UPLOAD_MAX_RETRY_COUNT_DEFAULT, parameterOverrides, - props); + props, + false); - this.updateValue( - MAX_MEMORY_LIMIT_IN_BYTES, MAX_MEMORY_LIMIT_IN_BYTES_DEFAULT, parameterOverrides, props); + this.checkAndUpdate( + MAX_MEMORY_LIMIT_IN_BYTES, + MAX_MEMORY_LIMIT_IN_BYTES_DEFAULT, + parameterOverrides, + props, + false); - this.updateValue( + this.checkAndUpdate( ENABLE_PARQUET_INTERNAL_BUFFERING, ENABLE_PARQUET_INTERNAL_BUFFERING_DEFAULT, parameterOverrides, - props); + props, + false); - this.updateValue( - MAX_CHANNEL_SIZE_IN_BYTES, MAX_CHANNEL_SIZE_IN_BYTES_DEFAULT, parameterOverrides, props); + this.checkAndUpdate( + MAX_CHANNEL_SIZE_IN_BYTES, + MAX_CHANNEL_SIZE_IN_BYTES_DEFAULT, + parameterOverrides, + props, + false); - this.updateValue( - MAX_CHUNK_SIZE_IN_BYTES, MAX_CHUNK_SIZE_IN_BYTES_DEFAULT, parameterOverrides, props); + this.checkAndUpdate( + MAX_CHUNK_SIZE_IN_BYTES, MAX_CHUNK_SIZE_IN_BYTES_DEFAULT, parameterOverrides, props, false); - this.updateValue(MAX_CLIENT_LAG, MAX_CLIENT_LAG_DEFAULT, parameterOverrides, props); + this.checkAndUpdate( + MAX_CLIENT_LAG, + isIcebergMode ? MAX_CLIENT_LAG_ICEBERG_MODE_DEFAULT : MAX_CLIENT_LAG_DEFAULT, + parameterOverrides, + props, + false); - this.updateValue( + this.checkAndUpdate( MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST, isIcebergMode ? MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT : MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT, parameterOverrides, - props); + props, + isIcebergMode); - this.updateValue( + this.checkAndUpdate( BDEC_PARQUET_COMPRESSION_ALGORITHM, BDEC_PARQUET_COMPRESSION_ALGORITHM_DEFAULT, parameterOverrides, - props); - - // Required parameter override for Iceberg mode - if (isIcebergMode) { - icebergModeValidation( - MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST, - MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT); - } + props, + false); } /** @return Longest interval in milliseconds between buffer flushes */ @@ -434,14 +470,4 @@ public Constants.BdecParquetCompression getBdecParquetCompressionAlgorithm() { public String toString() { return "ParameterProvider{" + "parameterMap=" + parameterMap + '}'; } - - private void icebergModeValidation(String key, Object expected) { - Object val = this.parameterMap.get(key); - if (!val.equals(expected)) { - throw new SFException( - INVALID_CONFIG_PARAMETER, - String.format( - "The value %s for %s is invalid in Iceberg mode, should be %s.", val, key, expected)); - } - } } diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java index e66306d1f..5a64e05a0 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java @@ -1,3 +1,7 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.streaming.internal; import static net.snowflake.ingest.utils.Constants.BLOB_CHECKSUM_SIZE_IN_BYTES; @@ -91,6 +95,7 @@ private abstract static class TestContext implements AutoCloseable { FlushService flushService; StreamingIngestStage stage; ParameterProvider parameterProvider; + InternalParameterProvider internalParameterProvider; RegisterService registerService; final List> channelData = new ArrayList<>(); @@ -99,8 +104,10 @@ private abstract static class TestContext implements AutoCloseable { stage = Mockito.mock(StreamingIngestStage.class); Mockito.when(stage.getClientPrefix()).thenReturn("client_prefix"); parameterProvider = new ParameterProvider(isIcebergMode); + internalParameterProvider = new InternalParameterProvider(isIcebergMode); client = Mockito.mock(SnowflakeStreamingIngestClientInternal.class); Mockito.when(client.getParameterProvider()).thenReturn(parameterProvider); + Mockito.when(client.getInternalParameterProvider()).thenReturn(internalParameterProvider); channelCache = new ChannelCache<>(); Mockito.when(client.getChannelCache()).thenReturn(channelCache); registerService = Mockito.spy(new RegisterService(client, client.isTestMode())); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/InternalParameterProviderTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/InternalParameterProviderTest.java new file mode 100644 index 000000000..e716b3ddb --- /dev/null +++ b/src/test/java/net/snowflake/ingest/streaming/internal/InternalParameterProviderTest.java @@ -0,0 +1,26 @@ +/* + * Copyright (c) 2024. Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +@RunWith(Parameterized.class) +public class InternalParameterProviderTest { + @Parameterized.Parameters(name = "isIcebergMode: {0}") + public static Object[] isIcebergMode() { + return new Object[] {false, true}; + } + + @Parameterized.Parameter public static boolean isIcebergMode; + + @Test + public void testConstantParameterProvider() { + InternalParameterProvider internalParameterProvider = + new InternalParameterProvider(isIcebergMode); + assert internalParameterProvider.getEnableChunkEncryption() == !isIcebergMode; + } +} diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java index fb73aaa66..c4631b348 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java @@ -146,7 +146,8 @@ public void withIcebergDefaultValues() { ParameterProvider parameterProvider = new ParameterProvider(true); Assert.assertEquals( - ParameterProvider.MAX_CLIENT_LAG_DEFAULT, parameterProvider.getCachedMaxClientLagInMs()); + ParameterProvider.MAX_CLIENT_LAG_ICEBERG_MODE_DEFAULT, + parameterProvider.getCachedMaxClientLagInMs()); Assert.assertEquals( ParameterProvider.BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS_DEFAULT, parameterProvider.getBufferFlushCheckIntervalInMs()); From eac448aff696c60863be1fcdecd5de3fdc6c1104 Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Thu, 11 Jul 2024 13:01:18 -0700 Subject: [PATCH 3/9] SNOW-1497358 Support multiple storage for Iceberg mode (#783) Co-authored-by: Hitesh Madan --- .../connection/ServiceResponseHandler.java | 5 +- .../ingest/streaming/OpenChannelRequest.java | 12 +- .../streaming/internal/ChannelCache.java | 2 +- .../internal/ChannelConfigureRequest.java | 53 +++++ .../internal/ChannelConfigureResponse.java | 46 ++++ .../internal/ChannelFlushContext.java | 10 +- .../internal/ChannelsStatusRequest.java | 23 +- .../internal/ClientConfigureRequest.java | 22 ++ ...onse.java => ClientConfigureResponse.java} | 11 +- .../streaming/internal/ConfigureRequest.java | 38 +++ .../internal/DropChannelRequestInternal.java | 101 ++++++++ .../internal/ExternalVolumeManager.java | 173 ++++++++++++++ .../streaming/internal/FlushService.java | 141 +++-------- .../internal/InternalStageManager.java | 188 +++++++++++++++ .../internal/OpenChannelRequestInternal.java | 106 +++++++++ .../internal/RegisterBlobRequest.java | 48 ++++ .../internal/SnowflakeServiceClient.java | 221 ++++++++++++++++++ ...nowflakeStreamingIngestClientInternal.java | 201 ++++++---------- .../streaming/internal/StorageManager.java | 67 ++++++ .../internal/StreamingIngestRequest.java | 13 ++ .../internal/StreamingIngestResponse.java | 9 +- ...Stage.java => StreamingIngestStorage.java} | 168 ++++++------- .../internal/StreamingIngestUtils.java | 7 +- .../net/snowflake/ingest/utils/Constants.java | 3 +- .../net/snowflake/ingest/utils/ErrorCode.java | 7 +- .../net/snowflake/ingest/utils/Utils.java | 29 ++- .../streaming/internal/FlushServiceTest.java | 102 ++++---- .../SnowflakeStreamingIngestChannelTest.java | 11 +- .../SnowflakeStreamingIngestClientTest.java | 8 +- ...t.java => StreamingIngestStorageTest.java} | 164 +++++++------ .../internal/StreamingIngestUtilsIT.java | 23 +- 31 files changed, 1522 insertions(+), 490 deletions(-) create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java rename src/main/java/net/snowflake/ingest/streaming/internal/{ConfigureResponse.java => ClientConfigureResponse.java} (79%) create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ConfigureRequest.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelRequestInternal.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/RegisterBlobRequest.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java create mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestRequest.java rename src/main/java/net/snowflake/ingest/streaming/internal/{StreamingIngestStage.java => StreamingIngestStorage.java} (75%) rename src/test/java/net/snowflake/ingest/streaming/internal/{StreamingIngestStageTest.java => StreamingIngestStorageTest.java} (83%) diff --git a/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java b/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java index cef6e631c..034b4a6f0 100644 --- a/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java +++ b/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2012-2017 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.connection; @@ -42,7 +42,8 @@ public enum ApiName { STREAMING_DROP_CHANNEL("POST"), STREAMING_CHANNEL_STATUS("POST"), STREAMING_REGISTER_BLOB("POST"), - STREAMING_CLIENT_CONFIGURE("POST"); + STREAMING_CLIENT_CONFIGURE("POST"), + STREAMING_CHANNEL_CONFIGURE("POST"); private final String httpMethod; private ApiName(String httpMethod) { diff --git a/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java b/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java index cc8782dbd..9a7272a16 100644 --- a/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java +++ b/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming; @@ -41,7 +41,6 @@ public enum OnErrorOption { private final ZoneId defaultTimezone; private final String offsetToken; - private final boolean isOffsetTokenProvided; private final OffsetTokenVerificationFunction offsetTokenVerificationFunction; @@ -59,7 +58,6 @@ public static class OpenChannelRequestBuilder { private ZoneId defaultTimezone; private String offsetToken; - private boolean isOffsetTokenProvided = false; private OffsetTokenVerificationFunction offsetTokenVerificationFunction; @@ -95,7 +93,6 @@ public OpenChannelRequestBuilder setDefaultTimezone(ZoneId defaultTimezone) { public OpenChannelRequestBuilder setOffsetToken(String offsetToken) { this.offsetToken = offsetToken; - this.isOffsetTokenProvided = true; return this; } @@ -125,7 +122,6 @@ private OpenChannelRequest(OpenChannelRequestBuilder builder) { this.onErrorOption = builder.onErrorOption; this.defaultTimezone = builder.defaultTimezone; this.offsetToken = builder.offsetToken; - this.isOffsetTokenProvided = builder.isOffsetTokenProvided; this.offsetTokenVerificationFunction = builder.offsetTokenVerificationFunction; } @@ -150,7 +146,7 @@ public ZoneId getDefaultTimezone() { } public String getFullyQualifiedTableName() { - return String.format("%s.%s.%s", this.dbName, this.schemaName, this.tableName); + return Utils.getFullyQualifiedTableName(this.dbName, this.schemaName, this.tableName); } public OnErrorOption getOnErrorOption() { @@ -161,10 +157,6 @@ public String getOffsetToken() { return this.offsetToken; } - public boolean isOffsetTokenProvided() { - return this.isOffsetTokenProvided; - } - public OffsetTokenVerificationFunction getOffsetTokenVerificationFunction() { return this.offsetTokenVerificationFunction; } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java index 989be0fa1..539da6287 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming.internal; diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java new file mode 100644 index 000000000..31bca95ac --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import com.fasterxml.jackson.annotation.JsonProperty; + +/** Class used to serialize the channel configure request. */ +class ChannelConfigureRequest extends ConfigureRequest { + @JsonProperty("database") + private String database; + + @JsonProperty("schema") + private String schema; + + @JsonProperty("table") + private String table; + + /** + * Constructor for channel configure request + * + * @param role Role to be used for the request. + * @param database Database name. + * @param schema Schema name. + * @param table Table name. + */ + ChannelConfigureRequest(String role, String database, String schema, String table) { + setRole(role); + this.database = database; + this.schema = schema; + this.table = table; + } + + String getDatabase() { + return database; + } + + String getSchema() { + return schema; + } + + String getTable() { + return table; + } + + @Override + public String getStringForLogging() { + return String.format( + "ChannelConfigureRequest(role=%s, db=%s, schema=%s, table=%s, file_name=%s)", + getRole(), database, schema, table, getFileName()); + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java new file mode 100644 index 000000000..da65960b4 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** Class used to deserialize responses from channel configure endpoint */ +@JsonIgnoreProperties(ignoreUnknown = true) +class ChannelConfigureResponse extends StreamingIngestResponse { + @JsonProperty("status_code") + private Long statusCode; + + @JsonProperty("message") + private String message; + + @JsonProperty("stage_location") + private FileLocationInfo stageLocation; + + @Override + Long getStatusCode() { + return statusCode; + } + + void setStatusCode(Long statusCode) { + this.statusCode = statusCode; + } + + String getMessage() { + return message; + } + + void setMessage(String message) { + this.message = message; + } + + FileLocationInfo getStageLocation() { + return stageLocation; + } + + void setStageLocation(FileLocationInfo stageLocation) { + this.stageLocation = stageLocation; + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelFlushContext.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelFlushContext.java index 3e5265719..fe9542267 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelFlushContext.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelFlushContext.java @@ -1,9 +1,11 @@ /* - * Copyright (c) 2022 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2022-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming.internal; +import net.snowflake.ingest.utils.Utils; + /** * Channel immutable identification and encryption attributes. * @@ -36,12 +38,12 @@ class ChannelFlushContext { String encryptionKey, Long encryptionKeyId) { this.name = name; - this.fullyQualifiedName = String.format("%s.%s.%s.%s", dbName, schemaName, tableName, name); + this.fullyQualifiedName = + Utils.getFullyQualifiedChannelName(dbName, schemaName, tableName, name); this.dbName = dbName; this.schemaName = schemaName; this.tableName = tableName; - this.fullyQualifiedTableName = - String.format("%s.%s.%s", this.getDbName(), this.getSchemaName(), this.getTableName()); + this.fullyQualifiedTableName = Utils.getFullyQualifiedTableName(dbName, schemaName, tableName); this.channelSequencer = channelSequencer; this.encryptionKey = encryptionKey; this.encryptionKeyId = encryptionKeyId; diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java index b98782ab9..ce8a76ef1 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java @@ -1,14 +1,16 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming.internal; import com.fasterxml.jackson.annotation.JsonProperty; import java.util.List; +import java.util.stream.Collectors; +import net.snowflake.ingest.utils.Utils; /** Class to deserialize a request from a channel status request */ -class ChannelsStatusRequest { +class ChannelsStatusRequest implements StreamingIngestRequest { // Used to deserialize a channel request static class ChannelStatusRequestDTO { @@ -99,4 +101,21 @@ void setChannels(List channels) { List getChannels() { return channels; } + + @Override + public String getStringForLogging() { + return String.format( + "ChannelsStatusRequest(requestId=%s, role=%s, channels=[%s])", + requestId, + role, + channels.stream() + .map( + r -> + Utils.getFullyQualifiedChannelName( + r.getDatabaseName(), + r.getSchemaName(), + r.getTableName(), + r.getChannelName())) + .collect(Collectors.joining(", "))); + } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java new file mode 100644 index 000000000..ca6d72ab3 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java @@ -0,0 +1,22 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +/** Class used to serialize client configure request */ +class ClientConfigureRequest extends ConfigureRequest { + /** + * Constructor for client configure request + * + * @param role Role to be used for the request. + */ + ClientConfigureRequest(String role) { + setRole(role); + } + + @Override + public String getStringForLogging() { + return String.format("ClientConfigureRequest(role=%s, file_name=%s)", getRole(), getFileName()); + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureResponse.java similarity index 79% rename from src/main/java/net/snowflake/ingest/streaming/internal/ConfigureResponse.java rename to src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureResponse.java index 6fe73fbd5..03a1d3576 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureResponse.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureResponse.java @@ -4,10 +4,12 @@ package net.snowflake.ingest.streaming.internal; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import com.fasterxml.jackson.annotation.JsonProperty; /** Class used to deserialize responses from configure endpoint */ -class ConfigureResponse extends StreamingIngestResponse { +@JsonIgnoreProperties(ignoreUnknown = true) +class ClientConfigureResponse extends StreamingIngestResponse { @JsonProperty("prefix") private String prefix; @@ -63,4 +65,11 @@ Long getDeploymentId() { void setDeploymentId(Long deploymentId) { this.deploymentId = deploymentId; } + + String getClientPrefix() { + if (this.deploymentId == null) { + return this.prefix; + } + return this.prefix + "_" + this.deploymentId; + } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureRequest.java new file mode 100644 index 000000000..6ca10f52a --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureRequest.java @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** Abstract class for {@link ChannelConfigureRequest} and {@link ClientConfigureRequest} */ +abstract class ConfigureRequest implements StreamingIngestRequest { + @JsonProperty("role") + private String role; + + // File name for the GCS signed url request + @JsonInclude(JsonInclude.Include.NON_NULL) + @JsonProperty("file_name") + private String fileName; + + String getRole() { + return role; + } + + void setRole(String role) { + this.role = role; + } + + String getFileName() { + return fileName; + } + + void setFileName(String fileName) { + this.fileName = fileName; + } + + @Override + public abstract String getStringForLogging(); +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java new file mode 100644 index 000000000..25599999e --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java @@ -0,0 +1,101 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import net.snowflake.ingest.streaming.DropChannelRequest; +import net.snowflake.ingest.utils.Utils; + +/** Class used to serialize the {@link DropChannelRequest} */ +class DropChannelRequestInternal implements StreamingIngestRequest { + @JsonProperty("request_id") + private String requestId; + + @JsonProperty("role") + private String role; + + @JsonProperty("channel") + private String channel; + + @JsonProperty("table") + private String table; + + @JsonProperty("database") + private String database; + + @JsonProperty("schema") + private String schema; + + @JsonProperty("is_iceberg") + private boolean isIceberg; + + @JsonInclude(JsonInclude.Include.NON_NULL) + @JsonProperty("client_sequencer") + Long clientSequencer; + + DropChannelRequestInternal( + String requestId, + String role, + String database, + String schema, + String table, + String channel, + Long clientSequencer, + boolean isIceberg) { + this.requestId = requestId; + this.role = role; + this.database = database; + this.schema = schema; + this.table = table; + this.channel = channel; + this.clientSequencer = clientSequencer; + this.isIceberg = isIceberg; + } + + String getRequestId() { + return requestId; + } + + String getRole() { + return role; + } + + String getChannel() { + return channel; + } + + String getTable() { + return table; + } + + String getDatabase() { + return database; + } + + String getSchema() { + return schema; + } + + boolean getIsIceberg() { + return isIceberg; + } + + Long getClientSequencer() { + return clientSequencer; + } + + String getFullyQualifiedTableName() { + return Utils.getFullyQualifiedTableName(database, schema, table); + } + + @Override + public String getStringForLogging() { + return String.format( + "DropChannelRequest(requestId=%s, role=%s, db=%s, schema=%s, table=%s, channel=%s," + + " isIceberg=%s, clientSequencer=%s)", + requestId, role, database, schema, table, channel, isIceberg, clientSequencer); + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java new file mode 100644 index 000000000..2b8796bb2 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java @@ -0,0 +1,173 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import java.io.IOException; +import java.util.Calendar; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import net.snowflake.client.jdbc.SnowflakeSQLException; +import net.snowflake.ingest.connection.IngestResponseException; +import net.snowflake.ingest.utils.ErrorCode; +import net.snowflake.ingest.utils.SFException; +import net.snowflake.ingest.utils.Utils; + +class ExternalVolumeLocation { + public final String dbName; + public final String schemaName; + public final String tableName; + + public ExternalVolumeLocation(String dbName, String schemaName, String tableName) { + this.dbName = dbName; + this.schemaName = schemaName; + this.tableName = tableName; + } +} + +/** Class to manage multiple external volumes */ +class ExternalVolumeManager implements StorageManager { + // Reference to the external volume per table + private final Map> externalVolumeMap; + + // name of the owning client + private final String clientName; + + // role of the owning client + private final String role; + + // Reference to the Snowflake service client used for configure calls + private final SnowflakeServiceClient snowflakeServiceClient; + + // Client prefix generated by the Snowflake server + private final String clientPrefix; + + /** + * Constructor for ExternalVolumeManager + * + * @param isTestMode whether the manager in test mode + * @param role the role of the client + * @param clientName the name of the client + * @param snowflakeServiceClient the Snowflake service client used for configure calls + */ + ExternalVolumeManager( + boolean isTestMode, + String role, + String clientName, + SnowflakeServiceClient snowflakeServiceClient) { + this.role = role; + this.clientName = clientName; + this.snowflakeServiceClient = snowflakeServiceClient; + this.externalVolumeMap = new ConcurrentHashMap<>(); + try { + this.clientPrefix = + isTestMode + ? "testPrefix" + : this.snowflakeServiceClient + .clientConfigure(new ClientConfigureRequest(role)) + .getClientPrefix(); + } catch (IngestResponseException | IOException e) { + throw new SFException(e, ErrorCode.CLIENT_CONFIGURE_FAILURE, e.getMessage()); + } + } + + /** + * Given a fully qualified table name, return the target storage by looking up the table name + * + * @param fullyQualifiedTableName the target fully qualified table name + * @return target storage + */ + @Override + public StreamingIngestStorage getStorage( + String fullyQualifiedTableName) { + // Only one chunk per blob in Iceberg mode. + StreamingIngestStorage stage = + this.externalVolumeMap.get(fullyQualifiedTableName); + + if (stage == null) { + throw new SFException( + ErrorCode.INTERNAL_ERROR, + String.format("No external volume found for table %s", fullyQualifiedTableName)); + } + + return stage; + } + + /** + * Add a storage to the manager by looking up the table name from the open channel response + * + * @param dbName the database name + * @param schemaName the schema name + * @param tableName the table name + * @param fileLocationInfo response from open channel + */ + @Override + public void addStorage( + String dbName, String schemaName, String tableName, FileLocationInfo fileLocationInfo) { + String fullyQualifiedTableName = + Utils.getFullyQualifiedTableName(dbName, schemaName, tableName); + + try { + this.externalVolumeMap.put( + fullyQualifiedTableName, + new StreamingIngestStorage( + this, + this.clientName, + fileLocationInfo, + new ExternalVolumeLocation(dbName, schemaName, tableName), + DEFAULT_MAX_UPLOAD_RETRIES)); + } catch (SnowflakeSQLException | IOException err) { + throw new SFException(err, ErrorCode.UNABLE_TO_CONNECT_TO_STORAGE); + } + } + + /** + * Gets the latest file location info (with a renewed short-lived access token) for the specified + * location + * + * @param location A reference to the target location + * @param fileName optional filename for single-file signed URL fetch from server + * @return the new location information + */ + @Override + public FileLocationInfo getRefreshedLocation( + ExternalVolumeLocation location, Optional fileName) { + try { + ChannelConfigureRequest request = + new ChannelConfigureRequest( + this.role, location.dbName, location.schemaName, location.tableName); + fileName.ifPresent(request::setFileName); + ChannelConfigureResponse response = this.snowflakeServiceClient.channelConfigure(request); + return response.getStageLocation(); + } catch (IngestResponseException | IOException e) { + throw new SFException(e, ErrorCode.CLIENT_CONFIGURE_FAILURE, e.getMessage()); + } + } + + // TODO: SNOW-1502887 Blob path generation for iceberg table + @Override + public String generateBlobPath() { + return "snow_dummy_file_name"; + } + + // TODO: SNOW-1502887 Blob path generation for iceberg table + @Override + public void decrementBlobSequencer() {} + + // TODO: SNOW-1502887 Blob path generation for iceberg table + public String getBlobPath(Calendar calendar, String clientPrefix) { + return ""; + } + + /** + * Get the client prefix from first external volume in the map + * + * @return the client prefix + */ + @Override + public String getClientPrefix() { + return this.clientPrefix; + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java index 05d72433d..454a106bb 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java @@ -4,7 +4,6 @@ package net.snowflake.ingest.streaming.internal; -import static net.snowflake.ingest.utils.Constants.BLOB_EXTENSION_TYPE; import static net.snowflake.ingest.utils.Constants.DISABLE_BACKGROUND_FLUSH; import static net.snowflake.ingest.utils.Constants.MAX_BLOB_SIZE_IN_BYTES; import static net.snowflake.ingest.utils.Constants.MAX_THREAD_COUNT; @@ -19,13 +18,11 @@ import java.security.InvalidKeyException; import java.security.NoSuchAlgorithmException; import java.util.ArrayList; -import java.util.Calendar; import java.util.Collections; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.TimeZone; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; @@ -34,11 +31,9 @@ import java.util.concurrent.ThreadFactory; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; import javax.crypto.BadPaddingException; import javax.crypto.IllegalBlockSizeException; import javax.crypto.NoSuchPaddingException; -import net.snowflake.client.jdbc.SnowflakeSQLException; import net.snowflake.client.jdbc.internal.google.common.util.concurrent.ThreadFactoryBuilder; import net.snowflake.ingest.utils.Constants; import net.snowflake.ingest.utils.ErrorCode; @@ -84,9 +79,6 @@ List>> getData() { private static final Logging logger = new Logging(FlushService.class); - // Increasing counter to generate a unique blob name per client - private final AtomicLong counter; - // The client that owns this flush service private final SnowflakeStreamingIngestClientInternal owningClient; @@ -102,8 +94,8 @@ List>> getData() { // Reference to the channel cache private final ChannelCache channelCache; - // Reference to the Streaming Ingest stage - private final StreamingIngestStage targetStage; + // Reference to the Streaming Ingest storage manager + private final StorageManager storageManager; // Reference to register service private final RegisterService registerService; @@ -124,56 +116,22 @@ List>> getData() { private final Constants.BdecVersion bdecVersion; /** - * Constructor for TESTING that takes (usually mocked) StreamingIngestStage + * Default constructor * - * @param client - * @param cache - * @param isTestMode + * @param client the owning client + * @param cache the channel cache + * @param storageManager the storage manager + * @param isTestMode whether the service is running in test mode */ FlushService( SnowflakeStreamingIngestClientInternal client, ChannelCache cache, - StreamingIngestStage targetStage, // For testing + StorageManager storageManager, boolean isTestMode) { this.owningClient = client; this.channelCache = cache; - this.targetStage = targetStage; - this.counter = new AtomicLong(0); - this.registerService = new RegisterService<>(client, isTestMode); - this.isNeedFlush = false; - this.lastFlushTime = System.currentTimeMillis(); - this.isTestMode = isTestMode; - this.latencyTimerContextMap = new ConcurrentHashMap<>(); - this.bdecVersion = this.owningClient.getParameterProvider().getBlobFormatVersion(); - createWorkers(); - } - - /** - * Default constructor - * - * @param client - * @param cache - * @param isTestMode - */ - FlushService( - SnowflakeStreamingIngestClientInternal client, ChannelCache cache, boolean isTestMode) { - this.owningClient = client; - this.channelCache = cache; - try { - this.targetStage = - new StreamingIngestStage( - isTestMode, - client.getRole(), - client.getHttpClient(), - client.getRequestBuilder(), - client.getName(), - DEFAULT_MAX_UPLOAD_RETRIES); - } catch (SnowflakeSQLException | IOException err) { - throw new SFException(err, ErrorCode.UNABLE_TO_CONNECT_TO_STAGE); - } - + this.storageManager = storageManager; this.registerService = new RegisterService<>(client, isTestMode); - this.counter = new AtomicLong(0); this.isNeedFlush = false; this.lastFlushTime = System.currentTimeMillis(); this.isTestMode = isTestMode; @@ -363,7 +321,7 @@ void distributeFlushTasks() { while (itr.hasNext() || !leftoverChannelsDataPerTable.isEmpty()) { List>> blobData = new ArrayList<>(); float totalBufferSizeInBytes = 0F; - final String blobPath = getBlobPath(this.targetStage.getClientPrefix()); + final String blobPath = this.storageManager.generateBlobPath(); // Distribute work at table level, split the blob if reaching the blob size limit or the // channel has different encryption key ids @@ -445,9 +403,9 @@ && shouldStopProcessing( // Kick off a build job if (blobData.isEmpty()) { - // we decrement the counter so that we do not have gaps in the blob names created by this - // client. See method getBlobPath() below. - this.counter.decrementAndGet(); + // we decrement the blob sequencer so that we do not have gaps in the blob names created by + // this client. + this.storageManager.decrementBlobSequencer(); } else { long flushStartMs = System.currentTimeMillis(); if (this.owningClient.flushLatency != null) { @@ -459,7 +417,13 @@ && shouldStopProcessing( CompletableFuture.supplyAsync( () -> { try { - BlobMetadata blobMetadata = buildAndUpload(blobPath, blobData); + // Get the fully qualified table name from the first channel in the blob. + // This only matters when the client is in Iceberg mode. In Iceberg mode, + // all channels in the blob belong to the same table. + String fullyQualifiedTableName = + blobData.get(0).get(0).getChannelContext().getFullyQualifiedTableName(); + BlobMetadata blobMetadata = + buildAndUpload(blobPath, blobData, fullyQualifiedTableName); blobMetadata.getBlobStats().setFlushStartMs(flushStartMs); return blobMetadata; } catch (Throwable e) { @@ -542,9 +506,12 @@ private boolean shouldStopProcessing( * @param blobPath Path of the destination blob in cloud storage * @param blobData All the data for one blob. Assumes that all ChannelData in the inner List * belongs to the same table. Will error if this is not the case + * @param fullyQualifiedTableName the table name of the first channel in the blob, only matters in + * Iceberg mode * @return BlobMetadata for FlushService.upload */ - BlobMetadata buildAndUpload(String blobPath, List>> blobData) + BlobMetadata buildAndUpload( + String blobPath, List>> blobData, String fullyQualifiedTableName) throws IOException, NoSuchAlgorithmException, InvalidAlgorithmParameterException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidKeyException { @@ -560,12 +527,18 @@ BlobMetadata buildAndUpload(String blobPath, List>> blobData blob.blobStats.setBuildDurationMs(buildContext); - return upload(blobPath, blob.blobBytes, blob.chunksMetadataList, blob.blobStats); + return upload( + this.storageManager.getStorage(fullyQualifiedTableName), + blobPath, + blob.blobBytes, + blob.chunksMetadataList, + blob.blobStats); } /** * Upload a blob to Streaming Ingest dedicated stage * + * @param storage the storage to upload the blob * @param blobPath full path of the blob * @param blob blob data * @param metadata a list of chunk metadata @@ -573,13 +546,17 @@ BlobMetadata buildAndUpload(String blobPath, List>> blobData * @return BlobMetadata object used to create the register blob request */ BlobMetadata upload( - String blobPath, byte[] blob, List metadata, BlobStats blobStats) + StreamingIngestStorage storage, + String blobPath, + byte[] blob, + List metadata, + BlobStats blobStats) throws NoSuchAlgorithmException { logger.logInfo("Start uploading blob={}, size={}", blobPath, blob.length); long startTime = System.currentTimeMillis(); Timer.Context uploadContext = Utils.createTimerContext(this.owningClient.uploadLatency); - this.targetStage.put(blobPath, blob); + storage.put(blobPath, blob); if (uploadContext != null) { blobStats.setUploadDurationMs(uploadContext); @@ -636,45 +613,6 @@ void setNeedFlush() { this.isNeedFlush = true; } - /** - * Generate a blob path, which is: "YEAR/MONTH/DAY_OF_MONTH/HOUR_OF_DAY/MINUTE/.BDEC" - * - * @return the generated blob file path - */ - private String getBlobPath(String clientPrefix) { - Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("UTC")); - return getBlobPath(calendar, clientPrefix); - } - - /** For TESTING */ - String getBlobPath(Calendar calendar, String clientPrefix) { - if (isTestMode && clientPrefix == null) { - clientPrefix = "testPrefix"; - } - - Utils.assertStringNotNullOrEmpty("client prefix", clientPrefix); - int year = calendar.get(Calendar.YEAR); - int month = calendar.get(Calendar.MONTH) + 1; // Gregorian calendar starts from 0 - int day = calendar.get(Calendar.DAY_OF_MONTH); - int hour = calendar.get(Calendar.HOUR_OF_DAY); - int minute = calendar.get(Calendar.MINUTE); - long time = TimeUnit.MILLISECONDS.toSeconds(calendar.getTimeInMillis()); - long threadId = Thread.currentThread().getId(); - // Create the blob short name, the clientPrefix contains the deployment id - String blobShortName = - Long.toString(time, 36) - + "_" - + clientPrefix - + "_" - + threadId - + "_" - + this.counter.getAndIncrement() - + "." - + BLOB_EXTENSION_TYPE; - return year + "/" + month + "/" + day + "/" + hour + "/" + minute + "/" + blobShortName; - } - /** * Invalidate all the channels in the blob data * @@ -698,11 +636,6 @@ void invalidateAllChannelsInBlob( })); } - /** Get the server generated unique prefix for this client */ - String getClientPrefix() { - return this.targetStage.getClientPrefix(); - } - /** * Throttle if the number of queued buildAndUpload tasks is bigger than the total number of * available processors diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java new file mode 100644 index 000000000..ea207b141 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java @@ -0,0 +1,188 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import static net.snowflake.ingest.utils.Constants.BLOB_EXTENSION_TYPE; + +import com.google.common.annotations.VisibleForTesting; +import java.io.IOException; +import java.util.Calendar; +import java.util.Optional; +import java.util.TimeZone; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import net.snowflake.client.jdbc.SnowflakeSQLException; +import net.snowflake.ingest.connection.IngestResponseException; +import net.snowflake.ingest.utils.ErrorCode; +import net.snowflake.ingest.utils.SFException; +import net.snowflake.ingest.utils.Utils; + +class InternalStageLocation { + public InternalStageLocation() {} +} + +/** Class to manage single Snowflake internal stage */ +class InternalStageManager implements StorageManager { + // Target stage for the client + private final StreamingIngestStorage targetStage; + + // Increasing counter to generate a unique blob name per client + private final AtomicLong counter; + + // Whether the manager in test mode + private final boolean isTestMode; + + // Snowflake service client used for configure calls + private final SnowflakeServiceClient snowflakeServiceClient; + + // The role of the client + private final String role; + + // Client prefix generated by the Snowflake server + private final String clientPrefix; + + /** + * Constructor for InternalStageManager + * + * @param isTestMode whether the manager in test mode + * @param role the role of the client + * @param clientName the name of the client + * @param snowflakeServiceClient the Snowflake service client to use for configure calls + */ + InternalStageManager( + boolean isTestMode, + String role, + String clientName, + SnowflakeServiceClient snowflakeServiceClient) { + this.snowflakeServiceClient = snowflakeServiceClient; + this.isTestMode = isTestMode; + this.role = role; + this.counter = new AtomicLong(0); + try { + if (!isTestMode) { + ClientConfigureResponse response = + this.snowflakeServiceClient.clientConfigure(new ClientConfigureRequest(role)); + this.clientPrefix = response.getClientPrefix(); + this.targetStage = + new StreamingIngestStorage( + this, + clientName, + response.getStageLocation(), + new InternalStageLocation(), + DEFAULT_MAX_UPLOAD_RETRIES); + } else { + this.clientPrefix = "testPrefix"; + this.targetStage = + new StreamingIngestStorage( + this, + "testClient", + (StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge) null, + new InternalStageLocation(), + DEFAULT_MAX_UPLOAD_RETRIES); + } + } catch (IngestResponseException | IOException e) { + throw new SFException(e, ErrorCode.CLIENT_CONFIGURE_FAILURE, e.getMessage()); + } catch (SnowflakeSQLException e) { + throw new SFException(e, ErrorCode.UNABLE_TO_CONNECT_TO_STORAGE, e.getMessage()); + } + } + + /** + * Get the storage. In this case, the storage is always the target stage as there's only one stage + * in non-iceberg mode. + * + * @param fullyQualifiedTableName the target fully qualified table name + * @return the target storage + */ + @Override + @SuppressWarnings("unused") + public StreamingIngestStorage getStorage( + String fullyQualifiedTableName) { + // There's always only one stage for the client in non-iceberg mode + return targetStage; + } + + /** Add storage to the manager. Do nothing as there's only one stage in non-Iceberg mode. */ + @Override + public void addStorage( + String dbName, String schemaName, String tableName, FileLocationInfo fileLocationInfo) {} + + /** + * Gets the latest file location info (with a renewed short-lived access token) for the specified + * location + * + * @param location A reference to the target location + * @param fileName optional filename for single-file signed URL fetch from server + * @return the new location information + */ + @Override + public FileLocationInfo getRefreshedLocation( + InternalStageLocation location, Optional fileName) { + try { + ClientConfigureRequest request = new ClientConfigureRequest(this.role); + fileName.ifPresent(request::setFileName); + ClientConfigureResponse response = snowflakeServiceClient.clientConfigure(request); + return response.getStageLocation(); + } catch (IngestResponseException | IOException e) { + throw new SFException(e, ErrorCode.CLIENT_CONFIGURE_FAILURE, e.getMessage()); + } + } + + /** + * Generate a blob path, which is: "YEAR/MONTH/DAY_OF_MONTH/HOUR_OF_DAY/MINUTE/.BDEC" + * + * @return the generated blob file path + */ + @Override + public String generateBlobPath() { + Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("UTC")); + return getBlobPath(calendar, this.clientPrefix); + } + + @Override + public void decrementBlobSequencer() { + this.counter.decrementAndGet(); + } + + /** For TESTING */ + @VisibleForTesting + public String getBlobPath(Calendar calendar, String clientPrefix) { + if (this.isTestMode && clientPrefix == null) { + clientPrefix = "testPrefix"; + } + + Utils.assertStringNotNullOrEmpty("client prefix", clientPrefix); + int year = calendar.get(Calendar.YEAR); + int month = calendar.get(Calendar.MONTH) + 1; // Gregorian calendar starts from 0 + int day = calendar.get(Calendar.DAY_OF_MONTH); + int hour = calendar.get(Calendar.HOUR_OF_DAY); + int minute = calendar.get(Calendar.MINUTE); + long time = TimeUnit.MILLISECONDS.toSeconds(calendar.getTimeInMillis()); + long threadId = Thread.currentThread().getId(); + // Create the blob short name, the clientPrefix contains the deployment id + String blobShortName = + Long.toString(time, 36) + + "_" + + clientPrefix + + "_" + + threadId + + "_" + + this.counter.getAndIncrement() + + "." + + BLOB_EXTENSION_TYPE; + return year + "/" + month + "/" + day + "/" + hour + "/" + minute + "/" + blobShortName; + } + + /** + * Get the unique client prefix generated by the Snowflake server + * + * @return the client prefix + */ + @Override + public String getClientPrefix() { + return this.clientPrefix; + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelRequestInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelRequestInternal.java new file mode 100644 index 000000000..ffbb553f3 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelRequestInternal.java @@ -0,0 +1,106 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import net.snowflake.ingest.streaming.OpenChannelRequest; +import net.snowflake.ingest.utils.Constants; + +/** Class used to serialize the {@link OpenChannelRequest} */ +class OpenChannelRequestInternal implements StreamingIngestRequest { + @JsonProperty("request_id") + private String requestId; + + @JsonProperty("role") + private String role; + + @JsonProperty("channel") + private String channel; + + @JsonProperty("table") + private String table; + + @JsonProperty("database") + private String database; + + @JsonProperty("schema") + private String schema; + + @JsonProperty("write_mode") + private String writeMode; + + @JsonProperty("is_iceberg") + private boolean isIceberg; + + @JsonInclude(JsonInclude.Include.NON_NULL) + @JsonProperty("offset_token") + private String offsetToken; + + OpenChannelRequestInternal( + String requestId, + String role, + String database, + String schema, + String table, + String channel, + Constants.WriteMode writeMode, + String offsetToken, + boolean isIceberg) { + this.requestId = requestId; + this.role = role; + this.database = database; + this.schema = schema; + this.table = table; + this.channel = channel; + this.writeMode = writeMode.name(); + this.isIceberg = isIceberg; + this.offsetToken = offsetToken; + } + + String getRequestId() { + return requestId; + } + + String getRole() { + return role; + } + + String getChannel() { + return channel; + } + + String getTable() { + return table; + } + + String getDatabase() { + return database; + } + + String getSchema() { + return schema; + } + + String getWriteMode() { + return writeMode; + } + + boolean getIsIceberg() { + return isIceberg; + } + + String getOffsetToken() { + return offsetToken; + } + + @Override + public String getStringForLogging() { + return String.format( + "OpenChannelRequestInternal(requestId=%s, role=%s, db=%s, schema=%s, table=%s, channel=%s," + + " writeMode=%s, isIceberg=%s)", + requestId, role, database, schema, table, channel, writeMode, isIceberg); + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/RegisterBlobRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/RegisterBlobRequest.java new file mode 100644 index 000000000..18fcb1fb2 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/RegisterBlobRequest.java @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.List; +import java.util.stream.Collectors; + +/** Class used to serialize the blob register request */ +class RegisterBlobRequest implements StreamingIngestRequest { + @JsonProperty("request_id") + private String requestId; + + @JsonProperty("role") + private String role; + + @JsonProperty("blobs") + private List blobs; + + RegisterBlobRequest(String requestId, String role, List blobs) { + this.requestId = requestId; + this.role = role; + this.blobs = blobs; + } + + String getRequestId() { + return requestId; + } + + String getRole() { + return role; + } + + List getBlobs() { + return blobs; + } + + @Override + public String getStringForLogging() { + return String.format( + "RegisterBlobRequest(requestId=%s, role=%s, blobs=[%s])", + requestId, + role, + blobs.stream().map(BlobMetadata::getPath).collect(Collectors.joining(", "))); + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java new file mode 100644 index 000000000..947c86dbb --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java @@ -0,0 +1,221 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CHANNEL_CONFIGURE; +import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CHANNEL_STATUS; +import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CLIENT_CONFIGURE; +import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_DROP_CHANNEL; +import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_OPEN_CHANNEL; +import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_REGISTER_BLOB; +import static net.snowflake.ingest.streaming.internal.StreamingIngestUtils.executeWithRetries; +import static net.snowflake.ingest.utils.Constants.CHANNEL_CONFIGURE_ENDPOINT; +import static net.snowflake.ingest.utils.Constants.CHANNEL_STATUS_ENDPOINT; +import static net.snowflake.ingest.utils.Constants.CLIENT_CONFIGURE_ENDPOINT; +import static net.snowflake.ingest.utils.Constants.DROP_CHANNEL_ENDPOINT; +import static net.snowflake.ingest.utils.Constants.OPEN_CHANNEL_ENDPOINT; +import static net.snowflake.ingest.utils.Constants.REGISTER_BLOB_ENDPOINT; +import static net.snowflake.ingest.utils.Constants.RESPONSE_SUCCESS; + +import java.io.IOException; +import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient; +import net.snowflake.ingest.connection.IngestResponseException; +import net.snowflake.ingest.connection.RequestBuilder; +import net.snowflake.ingest.connection.ServiceResponseHandler; +import net.snowflake.ingest.utils.ErrorCode; +import net.snowflake.ingest.utils.Logging; +import net.snowflake.ingest.utils.SFException; + +/** + * The SnowflakeServiceClient class is responsible for making API requests to the Snowflake service. + */ +class SnowflakeServiceClient { + private static final Logging logger = new Logging(SnowflakeServiceClient.class); + + // HTTP client used for making requests + private final CloseableHttpClient httpClient; + + // Request builder for building streaming API request + private final RequestBuilder requestBuilder; + + /** + * Default constructor + * + * @param httpClient the HTTP client used for making requests + * @param requestBuilder the request builder for building streaming API requests + */ + SnowflakeServiceClient(CloseableHttpClient httpClient, RequestBuilder requestBuilder) { + this.httpClient = httpClient; + this.requestBuilder = requestBuilder; + } + + /** + * Configures the client given a {@link ClientConfigureRequest}. + * + * @param request the client configuration request + * @return the response from the configuration request + */ + ClientConfigureResponse clientConfigure(ClientConfigureRequest request) + throws IngestResponseException, IOException { + ClientConfigureResponse response = + executeApiRequestWithRetries( + ClientConfigureResponse.class, + request, + CLIENT_CONFIGURE_ENDPOINT, + "client configure", + STREAMING_CLIENT_CONFIGURE); + if (response.getStatusCode() != RESPONSE_SUCCESS) { + logger.logDebug( + "Client configure request failed, request={}, message={}", + request.getStringForLogging(), + response.getMessage()); + throw new SFException(ErrorCode.CLIENT_CONFIGURE_FAILURE, response.getMessage()); + } + return response; + } + + /** + * Configures a storage given a {@link ChannelConfigureRequest}. + * + * @param request the channel configuration request + * @return the response from the configuration request + */ + ChannelConfigureResponse channelConfigure(ChannelConfigureRequest request) + throws IngestResponseException, IOException { + ChannelConfigureResponse response = + executeApiRequestWithRetries( + ChannelConfigureResponse.class, + request, + CHANNEL_CONFIGURE_ENDPOINT, + "channel configure", + STREAMING_CHANNEL_CONFIGURE); + + if (response.getStatusCode() != RESPONSE_SUCCESS) { + logger.logDebug( + "Channel configure request failed, request={}, response={}", + request.getStringForLogging(), + response.getMessage()); + throw new SFException(ErrorCode.CHANNEL_CONFIGURE_FAILURE, response.getMessage()); + } + return response; + } + + /** + * Opens a channel given a {@link OpenChannelRequestInternal}. + * + * @param request the open channel request + * @return the response from the open channel request + */ + OpenChannelResponse openChannel(OpenChannelRequestInternal request) + throws IngestResponseException, IOException { + OpenChannelResponse response = + executeApiRequestWithRetries( + OpenChannelResponse.class, + request, + OPEN_CHANNEL_ENDPOINT, + "open channel", + STREAMING_OPEN_CHANNEL); + + if (response.getStatusCode() != RESPONSE_SUCCESS) { + logger.logDebug( + "Open channel request failed, request={}, response={}", + request.getStringForLogging(), + response.getMessage()); + throw new SFException(ErrorCode.OPEN_CHANNEL_FAILURE, response.getMessage()); + } + return response; + } + + /** + * Drops a channel given a {@link DropChannelRequestInternal}. + * + * @param request the drop channel request + * @return the response from the drop channel request + */ + DropChannelResponse dropChannel(DropChannelRequestInternal request) + throws IngestResponseException, IOException { + DropChannelResponse response = + executeApiRequestWithRetries( + DropChannelResponse.class, + request, + DROP_CHANNEL_ENDPOINT, + "drop channel", + STREAMING_DROP_CHANNEL); + + if (response.getStatusCode() != RESPONSE_SUCCESS) { + logger.logDebug( + "Drop channel request failed, request={}, response={}", + request.getStringForLogging(), + response.getMessage()); + throw new SFException(ErrorCode.DROP_CHANNEL_FAILURE, response.getMessage()); + } + return response; + } + + /** + * Gets the status of a channel given a {@link ChannelsStatusRequest}. + * + * @param request the channel status request + * @return the response from the channel status request + */ + ChannelsStatusResponse channelStatus(ChannelsStatusRequest request) + throws IngestResponseException, IOException { + ChannelsStatusResponse response = + executeApiRequestWithRetries( + ChannelsStatusResponse.class, + request, + CHANNEL_STATUS_ENDPOINT, + "channel status", + STREAMING_CHANNEL_STATUS); + + if (response.getStatusCode() != RESPONSE_SUCCESS) { + logger.logDebug( + "Channel status request failed, request={}, response={}", + request.getStringForLogging(), + response.getMessage()); + throw new SFException(ErrorCode.CHANNEL_STATUS_FAILURE, response.getMessage()); + } + return response; + } + + /** + * Registers a blob given a {@link RegisterBlobRequest}. + * + * @param request the register blob request + * @param executionCount the number of times the request has been executed, used for logging + * @return the response from the register blob request + */ + RegisterBlobResponse registerBlob(RegisterBlobRequest request, final int executionCount) + throws IngestResponseException, IOException { + RegisterBlobResponse response = + executeApiRequestWithRetries( + RegisterBlobResponse.class, + request, + REGISTER_BLOB_ENDPOINT, + "register blob", + STREAMING_REGISTER_BLOB); + + if (response.getStatusCode() != RESPONSE_SUCCESS) { + logger.logDebug( + "Register blob request failed, request={}, response={}, executionCount={}", + request.getStringForLogging(), + response.getMessage(), + executionCount); + throw new SFException(ErrorCode.REGISTER_BLOB_FAILURE, response.getMessage()); + } + return response; + } + + private T executeApiRequestWithRetries( + Class responseClass, + StreamingIngestRequest request, + String endpoint, + String operation, + ServiceResponseHandler.ApiName apiName) + throws IngestResponseException, IOException { + return executeWithRetries( + responseClass, endpoint, request, operation, apiName, this.httpClient, this.requestBuilder); + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java index 1de473da6..8583a94f6 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java @@ -4,20 +4,11 @@ package net.snowflake.ingest.streaming.internal; -import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CHANNEL_STATUS; -import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_DROP_CHANNEL; -import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_OPEN_CHANNEL; -import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_REGISTER_BLOB; -import static net.snowflake.ingest.streaming.internal.StreamingIngestUtils.executeWithRetries; import static net.snowflake.ingest.streaming.internal.StreamingIngestUtils.sleepForRetry; -import static net.snowflake.ingest.utils.Constants.CHANNEL_STATUS_ENDPOINT; import static net.snowflake.ingest.utils.Constants.COMMIT_MAX_RETRY_COUNT; import static net.snowflake.ingest.utils.Constants.COMMIT_RETRY_INTERVAL_IN_MS; -import static net.snowflake.ingest.utils.Constants.DROP_CHANNEL_ENDPOINT; import static net.snowflake.ingest.utils.Constants.ENABLE_TELEMETRY_TO_SF; import static net.snowflake.ingest.utils.Constants.MAX_STREAMING_INGEST_API_CHANNEL_RETRY; -import static net.snowflake.ingest.utils.Constants.OPEN_CHANNEL_ENDPOINT; -import static net.snowflake.ingest.utils.Constants.REGISTER_BLOB_ENDPOINT; import static net.snowflake.ingest.utils.Constants.RESPONSE_ERR_ENQUEUE_TABLE_CHUNK_QUEUE_FULL; import static net.snowflake.ingest.utils.Constants.RESPONSE_ERR_GENERAL_EXCEPTION_RETRY_REQUEST; import static net.snowflake.ingest.utils.Constants.RESPONSE_SUCCESS; @@ -121,6 +112,9 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea // Reference to the flush service private final FlushService flushService; + // Reference to storage manager + private final StorageManager storageManager; + // Indicates whether the client has closed private volatile boolean isClosed; @@ -151,6 +145,9 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea // Background thread that uploads telemetry data periodically private ScheduledExecutorService telemetryWorker; + // Snowflake service client to make API calls + private SnowflakeServiceClient snowflakeServiceClient; + /** * Constructor * @@ -238,8 +235,18 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea this.setupMetricsForClient(); } + this.snowflakeServiceClient = new SnowflakeServiceClient(this.httpClient, this.requestBuilder); + + this.storageManager = + isIcebergMode + ? new ExternalVolumeManager( + isTestMode, this.role, this.name, this.snowflakeServiceClient) + : new InternalStageManager( + isTestMode, this.role, this.name, this.snowflakeServiceClient); + try { - this.flushService = new FlushService<>(this, this.channelCache, this.isTestMode); + this.flushService = + new FlushService<>(this, this.channelCache, this.storageManager, this.isTestMode); } catch (Exception e) { // Need to clean up the resources before throwing any exceptions cleanUpResources(); @@ -286,6 +293,7 @@ public SnowflakeStreamingIngestClientInternal( @VisibleForTesting public void injectRequestBuilder(RequestBuilder requestBuilder) { this.requestBuilder = requestBuilder; + this.snowflakeServiceClient = new SnowflakeServiceClient(this.httpClient, this.requestBuilder); } /** @@ -332,40 +340,18 @@ public SnowflakeStreamingIngestChannelInternal openChannel(OpenChannelRequest getName()); try { - Map payload = new HashMap<>(); - payload.put( - "request_id", this.flushService.getClientPrefix() + "_" + counter.getAndIncrement()); - payload.put("channel", request.getChannelName()); - payload.put("table", request.getTableName()); - payload.put("database", request.getDBName()); - payload.put("schema", request.getSchemaName()); - payload.put("write_mode", Constants.WriteMode.CLOUD_STORAGE.name()); - payload.put("role", this.role); - payload.put("is_iceberg", isIcebergMode); - if (request.isOffsetTokenProvided()) { - payload.put("offset_token", request.getOffsetToken()); - } - - OpenChannelResponse response = - executeWithRetries( - OpenChannelResponse.class, - OPEN_CHANNEL_ENDPOINT, - payload, - "open channel", - STREAMING_OPEN_CHANNEL, - httpClient, - requestBuilder); - - // Check for Snowflake specific response code - if (response.getStatusCode() != RESPONSE_SUCCESS) { - logger.logDebug( - "Open channel request failed, channel={}, table={}, client={}, message={}", - request.getChannelName(), - request.getFullyQualifiedTableName(), - getName(), - response.getMessage()); - throw new SFException(ErrorCode.OPEN_CHANNEL_FAILURE, response.getMessage()); - } + OpenChannelRequestInternal openChannelRequest = + new OpenChannelRequestInternal( + this.storageManager.getClientPrefix() + "_" + counter.getAndIncrement(), + this.role, + request.getDBName(), + request.getSchemaName(), + request.getTableName(), + request.getChannelName(), + Constants.WriteMode.CLOUD_STORAGE, + request.getOffsetToken(), + isIcebergMode); + OpenChannelResponse response = snowflakeServiceClient.openChannel(openChannelRequest); logger.logInfo( "Open channel request succeeded, channel={}, table={}, clientSequencer={}," @@ -399,8 +385,15 @@ public SnowflakeStreamingIngestChannelInternal openChannel(OpenChannelRequest // Add channel to the channel cache this.channelCache.addChannel(channel); + // Add storage to the storage manager, only for external volume + this.storageManager.addStorage( + response.getDBName(), + response.getSchemaName(), + response.getTableName(), + response.getStageLocation()); + return channel; - } catch (IOException | IngestResponseException e) { + } catch (IngestResponseException | IOException e) { throw new SFException(e, ErrorCode.OPEN_CHANNEL_FAILURE, e.getMessage()); } } @@ -418,52 +411,29 @@ public void dropChannel(DropChannelRequest request) { getName()); try { - Map payload = new HashMap<>(); - payload.put( - "request_id", this.flushService.getClientPrefix() + "_" + counter.getAndIncrement()); - payload.put("channel", request.getChannelName()); - payload.put("table", request.getTableName()); - payload.put("database", request.getDBName()); - payload.put("schema", request.getSchemaName()); - payload.put("role", this.role); - payload.put("is_iceberg", isIcebergMode); - Long clientSequencer = null; - if (request instanceof DropChannelVersionRequest) { - clientSequencer = ((DropChannelVersionRequest) request).getClientSequencer(); - if (clientSequencer != null) { - payload.put("client_sequencer", clientSequencer); - } - } - - DropChannelResponse response = - executeWithRetries( - DropChannelResponse.class, - DROP_CHANNEL_ENDPOINT, - payload, - "drop channel", - STREAMING_DROP_CHANNEL, - httpClient, - requestBuilder); - - // Check for Snowflake specific response code - if (response.getStatusCode() != RESPONSE_SUCCESS) { - logger.logDebug( - "Drop channel request failed, channel={}, table={}, client={}, message={}", - request.getChannelName(), - request.getFullyQualifiedTableName(), - getName(), - response.getMessage()); - throw new SFException(ErrorCode.DROP_CHANNEL_FAILURE, response.getMessage()); - } + DropChannelRequestInternal dropChannelRequest = + new DropChannelRequestInternal( + this.storageManager.getClientPrefix() + "_" + counter.getAndIncrement(), + this.role, + request.getDBName(), + request.getSchemaName(), + request.getTableName(), + request.getChannelName(), + request instanceof DropChannelVersionRequest + ? ((DropChannelVersionRequest) request).getClientSequencer() + : null, + isIcebergMode); + snowflakeServiceClient.dropChannel(dropChannelRequest); logger.logInfo( "Drop channel request succeeded, channel={}, table={}, clientSequencer={} client={}", request.getChannelName(), request.getFullyQualifiedTableName(), - clientSequencer, + request instanceof DropChannelVersionRequest + ? ((DropChannelVersionRequest) request).getClientSequencer() + : null, getName()); - - } catch (IOException | IngestResponseException e) { + } catch (IngestResponseException | IOException e) { throw new SFException(e, ErrorCode.DROP_CHANNEL_FAILURE, e.getMessage()); } } @@ -507,24 +477,9 @@ ChannelsStatusResponse getChannelsStatus( .collect(Collectors.toList()); request.setChannels(requestDTOs); request.setRole(this.role); - request.setRequestId(this.flushService.getClientPrefix() + "_" + counter.getAndIncrement()); - - String payload = objectMapper.writeValueAsString(request); - - ChannelsStatusResponse response = - executeWithRetries( - ChannelsStatusResponse.class, - CHANNEL_STATUS_ENDPOINT, - payload, - "channel status", - STREAMING_CHANNEL_STATUS, - httpClient, - requestBuilder); - - // Check for Snowflake specific response code - if (response.getStatusCode() != RESPONSE_SUCCESS) { - throw new SFException(ErrorCode.CHANNEL_STATUS_FAILURE, response.getMessage()); - } + request.setRequestId(this.storageManager.getClientPrefix() + "_" + counter.getAndIncrement()); + + ChannelsStatusResponse response = snowflakeServiceClient.channelStatus(request); for (int idx = 0; idx < channels.size(); idx++) { SnowflakeStreamingIngestChannelInternal channel = channels.get(idx); @@ -546,7 +501,7 @@ ChannelsStatusResponse getChannelsStatus( } return response; - } catch (IOException | IngestResponseException e) { + } catch (IngestResponseException | IOException e) { throw new SFException(e, ErrorCode.CHANNEL_STATUS_FAILURE, e.getMessage()); } } @@ -619,35 +574,15 @@ void registerBlobs(List blobs, final int executionCount) { this.name, executionCount); - RegisterBlobResponse response = null; + RegisterBlobResponse response; try { - Map payload = new HashMap<>(); - payload.put( - "request_id", this.flushService.getClientPrefix() + "_" + counter.getAndIncrement()); - payload.put("blobs", blobs); - payload.put("role", this.role); - - response = - executeWithRetries( - RegisterBlobResponse.class, - REGISTER_BLOB_ENDPOINT, - payload, - "register blob", - STREAMING_REGISTER_BLOB, - httpClient, - requestBuilder); - - // Check for Snowflake specific response code - if (response.getStatusCode() != RESPONSE_SUCCESS) { - logger.logDebug( - "Register blob request failed for blob={}, client={}, message={}, executionCount={}", - blobs.stream().map(BlobMetadata::getPath).collect(Collectors.toList()), - this.name, - response.getMessage(), - executionCount); - throw new SFException(ErrorCode.REGISTER_BLOB_FAILURE, response.getMessage()); - } - } catch (IOException | IngestResponseException e) { + RegisterBlobRequest request = + new RegisterBlobRequest( + this.storageManager.getClientPrefix() + "_" + counter.getAndIncrement(), + this.role, + blobs); + response = snowflakeServiceClient.registerBlob(request, executionCount); + } catch (IngestResponseException | IOException e) { throw new SFException(e, ErrorCode.REGISTER_BLOB_FAILURE, e.getMessage()); } @@ -839,7 +774,7 @@ void setNeedFlush() { this.flushService.setNeedFlush(); } - /** Remove the channel in the channel cache if the channel sequencer matches */ + /** Remove the channel in the channel cache if the channel sequencer matches. Update storage */ void removeChannelIfSequencersMatch(SnowflakeStreamingIngestChannelInternal channel) { this.channelCache.removeChannelIfSequencersMatch(channel); } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java new file mode 100644 index 000000000..ea0b8843f --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java @@ -0,0 +1,67 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +import java.util.Optional; + +/** + * Interface to manage {@link StreamingIngestStorage} for {@link FlushService} + * + * @param The type of chunk data + * @param the type of location that's being managed (internal stage / external volume) + */ +interface StorageManager { + // Default max upload retries for streaming ingest storage + int DEFAULT_MAX_UPLOAD_RETRIES = 5; + + /** + * Given a fully qualified table name, return the target storage + * + * @param fullyQualifiedTableName the target fully qualified table name + * @return target stage + */ + StreamingIngestStorage getStorage(String fullyQualifiedTableName); + + /** + * Add a storage to the manager + * + * @param dbName the database name + * @param schemaName the schema name + * @param tableName the table name + * @param fileLocationInfo file location info from configure response + */ + void addStorage( + String dbName, String schemaName, String tableName, FileLocationInfo fileLocationInfo); + + /** + * Gets the latest file location info (with a renewed short-lived access token) for the specified + * location + * + * @param location A reference to the target location + * @param fileName optional filename for single-file signed URL fetch from server + * @return the new location information + */ + FileLocationInfo getRefreshedLocation(TLocation location, Optional fileName); + + /** + * Generate a unique blob path and increment the blob sequencer + * + * @return the blob path + */ + String generateBlobPath(); + + /** + * Decrement the blob sequencer, this method is needed to prevent gap between file name sequencer. + * See {@link StorageManager#generateBlobPath()} for more details. + */ + void decrementBlobSequencer(); + + /** + * Get the unique client prefix generated by the Snowflake server + * + * @return the client prefix + */ + String getClientPrefix(); +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestRequest.java new file mode 100644 index 000000000..3dd30b5e2 --- /dev/null +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestRequest.java @@ -0,0 +1,13 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + +package net.snowflake.ingest.streaming.internal; + +/** + * The StreamingIngestRequest interface is a marker interface used for type safety in the {@link + * SnowflakeServiceClient} for streaming ingest API request. + */ +interface StreamingIngestRequest { + String getStringForLogging(); +} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestResponse.java index 1ec01fceb..6c4df8c6d 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestResponse.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestResponse.java @@ -1,9 +1,16 @@ /* - * Copyright (c) 2022 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2022-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming.internal; +/** + * The StreamingIngestResponse class is an abstract class that represents a response from the + * Snowflake streaming ingest API. This class provides a common structure for all types of responses + * that can be received from the {@link SnowflakeServiceClient}. + */ abstract class StreamingIngestResponse { abstract Long getStatusCode(); + + abstract String getMessage(); } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStage.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java similarity index 75% rename from src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStage.java rename to src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java index 9752b311c..eb9f10826 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStage.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java @@ -1,13 +1,9 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming.internal; -import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CLIENT_CONFIGURE; -import static net.snowflake.ingest.streaming.internal.StreamingIngestUtils.executeWithRetries; -import static net.snowflake.ingest.utils.Constants.CLIENT_CONFIGURE_ENDPOINT; -import static net.snowflake.ingest.utils.Constants.RESPONSE_SUCCESS; import static net.snowflake.ingest.utils.HttpUtil.generateProxyPropertiesForJDBC; import static net.snowflake.ingest.utils.Utils.getStackTrace; @@ -21,8 +17,6 @@ import java.io.IOException; import java.io.InputStream; import java.nio.file.Paths; -import java.util.HashMap; -import java.util.Map; import java.util.Optional; import java.util.Properties; import java.util.concurrent.TimeUnit; @@ -33,16 +27,13 @@ import net.snowflake.client.jdbc.SnowflakeSQLException; import net.snowflake.client.jdbc.cloud.storage.StageInfo; import net.snowflake.client.jdbc.internal.apache.commons.io.FileUtils; -import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient; -import net.snowflake.ingest.connection.IngestResponseException; -import net.snowflake.ingest.connection.RequestBuilder; import net.snowflake.ingest.utils.ErrorCode; import net.snowflake.ingest.utils.Logging; import net.snowflake.ingest.utils.SFException; import net.snowflake.ingest.utils.Utils; -/** Handles uploading files to the Snowflake Streaming Ingest Stage */ -class StreamingIngestStage { +/** Handles uploading files to the Snowflake Streaming Ingest Storage */ +class StreamingIngestStorage { private static final ObjectMapper mapper = new ObjectMapper(); // Object mapper for parsing the client/configure response to Jackson version the same as @@ -54,7 +45,7 @@ class StreamingIngestStage { private static final long REFRESH_THRESHOLD_IN_MS = TimeUnit.MILLISECONDS.convert(1, TimeUnit.MINUTES); - private static final Logging logger = new Logging(StreamingIngestStage.class); + private static final Logging logger = new Logging(StreamingIngestStorage.class); /** * Wrapper class containing SnowflakeFileTransferMetadata and the timestamp at which the metadata @@ -86,59 +77,61 @@ state to record unknown age. } private SnowflakeFileTransferMetadataWithAge fileTransferMetadataWithAge; - private final CloseableHttpClient httpClient; - private final RequestBuilder requestBuilder; - private final String role; + private final StorageManager owningManager; + private final TLocation location; private final String clientName; - private String clientPrefix; private final int maxUploadRetries; // Proxy parameters that we set while calling the Snowflake JDBC to upload the streams private final Properties proxyProperties; - StreamingIngestStage( - boolean isTestMode, - String role, - CloseableHttpClient httpClient, - RequestBuilder requestBuilder, + /** + * Default constructor + * + * @param owningManager the storage manager owning this storage + * @param clientName The client name + * @param fileLocationInfo The file location information from open channel response + * @param location A reference to the target location + * @param maxUploadRetries The maximum number of retries to attempt + */ + StreamingIngestStorage( + StorageManager owningManager, String clientName, + FileLocationInfo fileLocationInfo, + TLocation location, int maxUploadRetries) throws SnowflakeSQLException, IOException { - this.httpClient = httpClient; - this.role = role; - this.requestBuilder = requestBuilder; - this.clientName = clientName; - this.proxyProperties = generateProxyPropertiesForJDBC(); - this.maxUploadRetries = maxUploadRetries; - if (!isTestMode) { - refreshSnowflakeMetadata(); - } + this( + owningManager, + clientName, + (SnowflakeFileTransferMetadataWithAge) null, + location, + maxUploadRetries); + createFileTransferMetadataWithAge(fileLocationInfo); } /** * Constructor for TESTING that takes SnowflakeFileTransferMetadataWithAge as input * - * @param isTestMode must be true - * @param role Snowflake role used by the Client - * @param httpClient http client reference - * @param requestBuilder request builder to build the HTTP request + * @param owningManager the storage manager owning this storage * @param clientName the client name * @param testMetadata SnowflakeFileTransferMetadataWithAge to test with + * @param location A reference to the target location + * @param maxUploadRetries the maximum number of retries to attempt */ - StreamingIngestStage( - boolean isTestMode, - String role, - CloseableHttpClient httpClient, - RequestBuilder requestBuilder, + StreamingIngestStorage( + StorageManager owningManager, String clientName, SnowflakeFileTransferMetadataWithAge testMetadata, - int maxRetryCount) + TLocation location, + int maxUploadRetries) throws SnowflakeSQLException, IOException { - this(isTestMode, role, httpClient, requestBuilder, clientName, maxRetryCount); - if (!isTestMode) { - throw new SFException(ErrorCode.INTERNAL_ERROR); - } + this.owningManager = owningManager; + this.clientName = clientName; + this.maxUploadRetries = maxUploadRetries; + this.proxyProperties = generateProxyPropertiesForJDBC(); + this.location = location; this.fileTransferMetadataWithAge = testMetadata; } @@ -192,7 +185,7 @@ private void putRemote(String fullFilePath, byte[] data, int retryCount) .setUploadStream(inStream) .setRequireCompress(false) .setOcspMode(OCSPMode.FAIL_OPEN) - .setStreamingIngestClientKey(this.clientPrefix) + .setStreamingIngestClientKey(this.owningManager.getClientPrefix()) .setStreamingIngestClientName(this.clientName) .setProxyProperties(this.proxyProperties) .setDestFileName(fullFilePath) @@ -250,26 +243,26 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boole return fileTransferMetadataWithAge; } - Map payload = new HashMap<>(); - payload.put("role", this.role); - ConfigureResponse response = this.makeClientConfigureCall(payload); + FileLocationInfo location = + this.owningManager.getRefreshedLocation(this.location, Optional.empty()); + return createFileTransferMetadataWithAge(location); + } - // Do not change the prefix everytime we have to refresh credentials - if (Utils.isNullOrEmpty(this.clientPrefix)) { - this.clientPrefix = createClientPrefix(response); - } - Utils.assertStringNotNullOrEmpty("client prefix", this.clientPrefix); + private SnowflakeFileTransferMetadataWithAge createFileTransferMetadataWithAge( + FileLocationInfo fileLocationInfo) + throws JsonProcessingException, + net.snowflake.client.jdbc.internal.fasterxml.jackson.core.JsonProcessingException, + SnowflakeSQLException { + Utils.assertStringNotNullOrEmpty("client prefix", this.owningManager.getClientPrefix()); - if (response - .getStageLocation() + if (fileLocationInfo .getLocationType() .replaceAll( "^[\"]|[\"]$", "") // Replace the first and last character if they're double quotes .equals(StageInfo.StageType.LOCAL_FS.name())) { this.fileTransferMetadataWithAge = new SnowflakeFileTransferMetadataWithAge( - response - .getStageLocation() + fileLocationInfo .getLocation() .replaceAll( "^[\"]|[\"]$", @@ -280,7 +273,7 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boole new SnowflakeFileTransferMetadataWithAge( (SnowflakeFileTransferMetadataV1) SnowflakeFileTransferAgent.getFileTransferMetadatas( - parseClientConfigureResponse(response)) + parseFileLocationInfo(fileLocationInfo)) .get(0), Optional.of(System.currentTimeMillis())); } @@ -296,8 +289,8 @@ synchronized SnowflakeFileTransferMetadataWithAge refreshSnowflakeMetadata(boole * @param response the client/configure response from the server * @return the client prefix. */ - private String createClientPrefix(final ConfigureResponse response) { - final String prefix = response.getPrefix(); + private String createClientPrefix(final ClientConfigureResponse response) { + final String prefix = response.getPrefix() == null ? "" : response.getPrefix(); final String deploymentId = response.getDeploymentId() != null ? "_" + response.getDeploymentId() : ""; return prefix + deploymentId; @@ -312,15 +305,12 @@ private String createClientPrefix(final ConfigureResponse response) { SnowflakeFileTransferMetadataV1 fetchSignedURL(String fileName) throws SnowflakeSQLException, IOException { - Map payload = new HashMap<>(); - payload.put("role", this.role); - payload.put("file_name", fileName); - ConfigureResponse response = this.makeClientConfigureCall(payload); + FileLocationInfo location = + this.owningManager.getRefreshedLocation(this.location, Optional.of(fileName)); SnowflakeFileTransferMetadataV1 metadata = (SnowflakeFileTransferMetadataV1) - SnowflakeFileTransferAgent.getFileTransferMetadatas( - parseClientConfigureResponse(response)) + SnowflakeFileTransferAgent.getFileTransferMetadatas(parseFileLocationInfo(location)) .get(0); // Transfer agent trims path for fileName metadata.setPresignedUrlFileName(fileName); @@ -328,51 +318,28 @@ SnowflakeFileTransferMetadataV1 fetchSignedURL(String fileName) } private net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.JsonNode - parseClientConfigureResponse(ConfigureResponse response) + parseFileLocationInfo(FileLocationInfo fileLocationInfo) throws JsonProcessingException, net.snowflake.client.jdbc.internal.fasterxml.jackson.core.JsonProcessingException { - JsonNode responseNode = mapper.valueToTree(response); + JsonNode fileLocationInfoNode = mapper.valueToTree(fileLocationInfo); // Currently there are a few mismatches between the client/configure response and what // SnowflakeFileTransferAgent expects - ObjectNode mutable = (ObjectNode) responseNode; - mutable.putObject("data"); - ObjectNode dataNode = (ObjectNode) mutable.get("data"); - dataNode.set("stageInfo", responseNode.get("stage_location")); + + ObjectNode node = mapper.createObjectNode(); + node.putObject("data"); + ObjectNode dataNode = (ObjectNode) node.get("data"); + dataNode.set("stageInfo", fileLocationInfoNode); // JDBC expects this field which maps to presignedFileUrlName. We will set this later dataNode.putArray("src_locations").add("placeholder"); // use String as intermediate object to avoid Jackson version mismatch // TODO: SNOW-1493470 Align Jackson version - String responseString = mapper.writeValueAsString(responseNode); + String responseString = mapper.writeValueAsString(node); return parseConfigureResponseMapper.readTree(responseString); } - private ConfigureResponse makeClientConfigureCall(Map payload) - throws IOException { - try { - - ConfigureResponse response = - executeWithRetries( - ConfigureResponse.class, - CLIENT_CONFIGURE_ENDPOINT, - mapper.writeValueAsString(payload), - "client configure", - STREAMING_CLIENT_CONFIGURE, - httpClient, - requestBuilder); - - // Check for Snowflake specific response code - if (response.getStatusCode() != RESPONSE_SUCCESS) { - throw new SFException(ErrorCode.CLIENT_CONFIGURE_FAILURE, response.getMessage()); - } - return response; - } catch (IngestResponseException e) { - throw new SFException(e, ErrorCode.CLIENT_CONFIGURE_FAILURE, e.getMessage()); - } - } - /** * Upload file to internal stage * @@ -416,9 +383,4 @@ void putLocal(String fullFilePath, byte[] data) { throw new SFException(ex, ErrorCode.BLOB_UPLOAD_FAILURE); } } - - /** Get the server generated unique prefix for this client */ - String getClientPrefix() { - return this.clientPrefix; - } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtils.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtils.java index 56e960064..ebfffa234 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtils.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtils.java @@ -1,3 +1,7 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.streaming.internal; import static net.snowflake.ingest.utils.Constants.MAX_STREAMING_INGEST_API_CHANNEL_RETRY; @@ -6,7 +10,6 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; -import java.util.Map; import java.util.function.Function; import net.snowflake.client.jdbc.internal.apache.http.client.methods.CloseableHttpResponse; import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpUriRequest; @@ -77,7 +80,7 @@ public static void sleepForRetry(int executionCount) { static T executeWithRetries( Class targetClass, String endpoint, - Map payload, + StreamingIngestRequest payload, String message, ServiceResponseHandler.ApiName apiName, CloseableHttpClient httpClient, diff --git a/src/main/java/net/snowflake/ingest/utils/Constants.java b/src/main/java/net/snowflake/ingest/utils/Constants.java index 3579e4d24..3d09d9a2a 100644 --- a/src/main/java/net/snowflake/ingest/utils/Constants.java +++ b/src/main/java/net/snowflake/ingest/utils/Constants.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.utils; @@ -52,6 +52,7 @@ public class Constants { public static final String BLOB_EXTENSION_TYPE = "bdec"; public static final int MAX_THREAD_COUNT = Integer.MAX_VALUE; public static final String CLIENT_CONFIGURE_ENDPOINT = "/v1/streaming/client/configure/"; + public static final String CHANNEL_CONFIGURE_ENDPOINT = "/v1/streaming/channels/configure/"; public static final int COMMIT_MAX_RETRY_COUNT = 60; public static final int COMMIT_RETRY_INTERVAL_IN_MS = 1000; public static final String ENCRYPTION_ALGORITHM = "AES/CTR/NoPadding"; diff --git a/src/main/java/net/snowflake/ingest/utils/ErrorCode.java b/src/main/java/net/snowflake/ingest/utils/ErrorCode.java index b863717e9..4091e98bf 100644 --- a/src/main/java/net/snowflake/ingest/utils/ErrorCode.java +++ b/src/main/java/net/snowflake/ingest/utils/ErrorCode.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.utils; @@ -26,7 +26,7 @@ public enum ErrorCode { INVALID_ENCRYPTED_KEY("0018"), INVALID_DATA_IN_CHUNK("0019"), IO_ERROR("0020"), - UNABLE_TO_CONNECT_TO_STAGE("0021"), + UNABLE_TO_CONNECT_TO_STORAGE("0021"), KEYPAIR_CREATION_FAILURE("0022"), MD5_HASHING_NOT_AVAILABLE("0023"), CHANNEL_STATUS_FAILURE("0024"), @@ -41,7 +41,8 @@ public enum ErrorCode { OAUTH_REFRESH_TOKEN_ERROR("0033"), INVALID_CONFIG_PARAMETER("0034"), CRYPTO_PROVIDER_ERROR("0035"), - DROP_CHANNEL_FAILURE("0036"); + DROP_CHANNEL_FAILURE("0036"), + CHANNEL_CONFIGURE_FAILURE("0037"); public static final String errorMessageResource = "net.snowflake.ingest.ingest_error_messages"; diff --git a/src/main/java/net/snowflake/ingest/utils/Utils.java b/src/main/java/net/snowflake/ingest/utils/Utils.java index a06df4027..5220625da 100644 --- a/src/main/java/net/snowflake/ingest/utils/Utils.java +++ b/src/main/java/net/snowflake/ingest/utils/Utils.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.utils; @@ -384,4 +384,31 @@ public static String getStackTrace(Throwable e) { } return stackTrace.toString(); } + + /** + * Get the fully qualified table name + * + * @param dbName the database name + * @param schemaName the schema name + * @param tableName the table name + * @return the fully qualified table name + */ + public static String getFullyQualifiedTableName( + String dbName, String schemaName, String tableName) { + return String.format("%s.%s.%s", dbName, schemaName, tableName); + } + + /** + * Get the fully qualified channel name + * + * @param dbName the database name + * @param schemaName the schema name + * @param tableName the table name + * @param channelName the channel name + * @return the fully qualified channel name + */ + public static String getFullyQualifiedChannelName( + String dbName, String schemaName, String tableName, String channelName) { + return String.format("%s.%s.%s.%s", dbName, schemaName, tableName, channelName); + } } diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java index 5a64e05a0..f59f9e076 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java @@ -93,7 +93,8 @@ private abstract static class TestContext implements AutoCloseable { ChannelCache channelCache; final Map> channels = new HashMap<>(); FlushService flushService; - StreamingIngestStage stage; + StorageManager storageManager; + StreamingIngestStorage storage; ParameterProvider parameterProvider; InternalParameterProvider internalParameterProvider; RegisterService registerService; @@ -101,17 +102,23 @@ private abstract static class TestContext implements AutoCloseable { final List> channelData = new ArrayList<>(); TestContext() { - stage = Mockito.mock(StreamingIngestStage.class); - Mockito.when(stage.getClientPrefix()).thenReturn("client_prefix"); + storage = Mockito.mock(StreamingIngestStorage.class); parameterProvider = new ParameterProvider(isIcebergMode); internalParameterProvider = new InternalParameterProvider(isIcebergMode); client = Mockito.mock(SnowflakeStreamingIngestClientInternal.class); Mockito.when(client.getParameterProvider()).thenReturn(parameterProvider); Mockito.when(client.getInternalParameterProvider()).thenReturn(internalParameterProvider); + storageManager = + Mockito.spy( + isIcebergMode + ? new ExternalVolumeManager<>(true, "role", "client", null) + : new InternalStageManager<>(true, "role", "client", null)); + Mockito.doReturn(storage).when(storageManager).getStorage(ArgumentMatchers.any()); + Mockito.when(storageManager.getClientPrefix()).thenReturn("client_prefix"); channelCache = new ChannelCache<>(); Mockito.when(client.getChannelCache()).thenReturn(channelCache); registerService = Mockito.spy(new RegisterService(client, client.isTestMode())); - flushService = Mockito.spy(new FlushService<>(client, channelCache, stage, true)); + flushService = Mockito.spy(new FlushService<>(client, channelCache, storageManager, true)); } ChannelData flushChannel(String name) { @@ -124,7 +131,10 @@ ChannelData flushChannel(String name) { BlobMetadata buildAndUpload() throws Exception { List>> blobData = Collections.singletonList(channelData); - return flushService.buildAndUpload("file_name", blobData); + return flushService.buildAndUpload( + "file_name", + blobData, + blobData.get(0).get(0).getChannelContext().getFullyQualifiedTableName()); } abstract SnowflakeStreamingIngestChannelInternal createChannel( @@ -408,35 +418,41 @@ private static ColumnMetadata createLargeTestTextColumn(String name) { @Test public void testGetFilePath() { TestContext testContext = testContextFactory.create(); - FlushService flushService = testContext.flushService; + StorageManager storageManager = testContext.storageManager; Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("UTC")); String clientPrefix = "honk"; - String outputString = flushService.getBlobPath(calendar, clientPrefix); - Path outputPath = Paths.get(outputString); - Assert.assertTrue(outputPath.getFileName().toString().contains(clientPrefix)); - Assert.assertTrue( - calendar.get(Calendar.MINUTE) - - Integer.parseInt(outputPath.getParent().getFileName().toString()) - <= 1); - Assert.assertEquals( - Integer.toString(calendar.get(Calendar.HOUR_OF_DAY)), - outputPath.getParent().getParent().getFileName().toString()); - Assert.assertEquals( - Integer.toString(calendar.get(Calendar.DAY_OF_MONTH)), - outputPath.getParent().getParent().getParent().getFileName().toString()); - Assert.assertEquals( - Integer.toString(calendar.get(Calendar.MONTH) + 1), - outputPath.getParent().getParent().getParent().getParent().getFileName().toString()); - Assert.assertEquals( - Integer.toString(calendar.get(Calendar.YEAR)), - outputPath - .getParent() - .getParent() - .getParent() - .getParent() - .getParent() - .getFileName() - .toString()); + if (isIcebergMode) { + // TODO: SNOW-1502887 Blob path generation for iceberg table + String outputString = storageManager.generateBlobPath(); + } else { + String outputString = + ((InternalStageManager) storageManager).getBlobPath(calendar, clientPrefix); + Path outputPath = Paths.get(outputString); + Assert.assertTrue(outputPath.getFileName().toString().contains(clientPrefix)); + Assert.assertTrue( + calendar.get(Calendar.MINUTE) + - Integer.parseInt(outputPath.getParent().getFileName().toString()) + <= 1); + Assert.assertEquals( + Integer.toString(calendar.get(Calendar.HOUR_OF_DAY)), + outputPath.getParent().getParent().getFileName().toString()); + Assert.assertEquals( + Integer.toString(calendar.get(Calendar.DAY_OF_MONTH)), + outputPath.getParent().getParent().getParent().getFileName().toString()); + Assert.assertEquals( + Integer.toString(calendar.get(Calendar.MONTH) + 1), + outputPath.getParent().getParent().getParent().getParent().getFileName().toString()); + Assert.assertEquals( + Integer.toString(calendar.get(Calendar.YEAR)), + outputPath + .getParent() + .getParent() + .getParent() + .getParent() + .getParent() + .getFileName() + .toString()); + } } @Test @@ -499,7 +515,8 @@ public void testBlobCreation() throws Exception { // Force = true flushes flushService.flush(true).get(); - Mockito.verify(flushService, Mockito.atLeast(2)).buildAndUpload(Mockito.any(), Mockito.any()); + Mockito.verify(flushService, Mockito.atLeast(2)) + .buildAndUpload(Mockito.any(), Mockito.any(), Mockito.any()); } @Test @@ -548,7 +565,8 @@ public void testBlobSplitDueToDifferentSchema() throws Exception { // Force = true flushes flushService.flush(true).get(); - Mockito.verify(flushService, Mockito.atLeast(2)).buildAndUpload(Mockito.any(), Mockito.any()); + Mockito.verify(flushService, Mockito.atLeast(2)) + .buildAndUpload(Mockito.any(), Mockito.any(), Mockito.any()); } @Test @@ -582,7 +600,8 @@ public void testBlobSplitDueToChunkSizeLimit() throws Exception { // Force = true flushes flushService.flush(true).get(); - Mockito.verify(flushService, Mockito.times(2)).buildAndUpload(Mockito.any(), Mockito.any()); + Mockito.verify(flushService, Mockito.times(2)) + .buildAndUpload(Mockito.any(), Mockito.any(), Mockito.any()); } @Test @@ -624,7 +643,7 @@ public void runTestBlobSplitDueToNumberOfChunks(int numberOfRows) throws Excepti ArgumentCaptor>>>>> blobDataCaptor = ArgumentCaptor.forClass(List.class); Mockito.verify(flushService, Mockito.times(expectedBlobs)) - .buildAndUpload(Mockito.any(), blobDataCaptor.capture()); + .buildAndUpload(Mockito.any(), blobDataCaptor.capture(), Mockito.any()); // 1. list => blobs; 2. list => chunks; 3. list => channels; 4. list => rows, 5. list => columns List>>>>> allUploadedBlobs = @@ -667,7 +686,7 @@ public void testBlobSplitDueToNumberOfChunksWithLeftoverChannels() throws Except ArgumentCaptor>>>>> blobDataCaptor = ArgumentCaptor.forClass(List.class); Mockito.verify(flushService, Mockito.atLeast(2)) - .buildAndUpload(Mockito.any(), blobDataCaptor.capture()); + .buildAndUpload(Mockito.any(), blobDataCaptor.capture(), Mockito.any()); // 1. list => blobs; 2. list => chunks; 3. list => channels; 4. list => rows, 5. list => columns List>>>>> allUploadedBlobs = @@ -785,12 +804,15 @@ public void testBuildAndUpload() throws Exception { .build(); // Check FlushService.upload called with correct arguments + final ArgumentCaptor storageCaptor = + ArgumentCaptor.forClass(StreamingIngestStorage.class); final ArgumentCaptor nameCaptor = ArgumentCaptor.forClass(String.class); final ArgumentCaptor blobCaptor = ArgumentCaptor.forClass(byte[].class); final ArgumentCaptor> metadataCaptor = ArgumentCaptor.forClass(List.class); Mockito.verify(testContext.flushService) .upload( + storageCaptor.capture(), nameCaptor.capture(), blobCaptor.capture(), metadataCaptor.capture(), @@ -933,10 +955,10 @@ public void testInvalidateChannels() { innerData.add(channel1Data); innerData.add(channel2Data); - StreamingIngestStage stage = Mockito.mock(StreamingIngestStage.class); - Mockito.when(stage.getClientPrefix()).thenReturn("client_prefix"); + StorageManager storageManager = + Mockito.spy(new InternalStageManager<>(true, "role", "client", null)); FlushService flushService = - new FlushService<>(client, channelCache, stage, false); + new FlushService<>(client, channelCache, storageManager, false); flushService.invalidateAllChannelsInBlob(blobData, "Invalidated by test"); Assert.assertFalse(channel1.isValid()); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java index 5beb0662f..f577c75c1 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java @@ -1,3 +1,7 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.streaming.internal; import static java.time.ZoneOffset.UTC; @@ -274,7 +278,7 @@ public void testOpenChannelRequestCreationSuccess() { Assert.assertEquals( "STREAMINGINGEST_TEST.PUBLIC.T_STREAMINGINGEST", request.getFullyQualifiedTableName()); - Assert.assertFalse(request.isOffsetTokenProvided()); + Assert.assertNull(request.getOffsetToken()); } @Test @@ -291,7 +295,6 @@ public void testOpenChannelRequesCreationtWithOffsetToken() { Assert.assertEquals( "STREAMINGINGEST_TEST.PUBLIC.T_STREAMINGINGEST", request.getFullyQualifiedTableName()); Assert.assertEquals("TEST_TOKEN", request.getOffsetToken()); - Assert.assertTrue(request.isOffsetTokenProvided()); } @Test @@ -451,6 +454,10 @@ public void testOpenChannelSnowflakeInternalErrorResponse() throws Exception { @Test public void testOpenChannelSuccessResponse() throws Exception { + // TODO: SNOW-1490151 Iceberg testing gaps + if (isIcebergMode) { + return; + } String name = "CHANNEL"; String dbName = "STREAMINGINGEST_TEST"; String schemaName = "PUBLIC"; diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java index 4e66d8a15..6325b3144 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java @@ -1,3 +1,7 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.streaming.internal; import static java.time.ZoneOffset.UTC; @@ -382,7 +386,7 @@ public void testGetChannelsStatusWithRequest() throws Exception { ChannelsStatusRequest.ChannelStatusRequestDTO dto = new ChannelsStatusRequest.ChannelStatusRequestDTO(channel); ChannelsStatusRequest request = new ChannelsStatusRequest(); - request.setRequestId("null_0"); + request.setRequestId("testPrefix_0"); request.setChannels(Collections.singletonList(dto)); ChannelsStatusResponse result = client.getChannelsStatus(Collections.singletonList(channel)); Assert.assertEquals(response.getMessage(), result.getMessage()); @@ -1458,7 +1462,7 @@ public void testGetLatestCommittedOffsetTokens() throws Exception { ChannelsStatusRequest.ChannelStatusRequestDTO dto = new ChannelsStatusRequest.ChannelStatusRequestDTO(channel); ChannelsStatusRequest request = new ChannelsStatusRequest(); - request.setRequestId("null_0"); + request.setRequestId("testPrefix_0"); request.setChannels(Collections.singletonList(dto)); Map result = client.getLatestCommittedOffsetTokens(Collections.singletonList(channel)); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStageTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java similarity index 83% rename from src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStageTest.java rename to src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java index d12c6231c..478934f4b 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStageTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java @@ -1,6 +1,11 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.streaming.internal; import static net.snowflake.client.core.Constants.CLOUD_STORAGE_CREDENTIALS_EXPIRED; +import static net.snowflake.ingest.utils.Constants.CLIENT_CONFIGURE_ENDPOINT; import static net.snowflake.ingest.utils.HttpUtil.HTTP_PROXY_PASSWORD; import static net.snowflake.ingest.utils.HttpUtil.HTTP_PROXY_USER; import static net.snowflake.ingest.utils.HttpUtil.NON_PROXY_HOSTS; @@ -33,6 +38,7 @@ import net.snowflake.client.jdbc.SnowflakeSQLException; import net.snowflake.client.jdbc.cloud.storage.StageInfo; import net.snowflake.client.jdbc.internal.amazonaws.util.IOUtils; +import net.snowflake.client.jdbc.internal.apache.http.HttpEntity; import net.snowflake.client.jdbc.internal.apache.http.StatusLine; import net.snowflake.client.jdbc.internal.apache.http.client.methods.CloseableHttpResponse; import net.snowflake.client.jdbc.internal.apache.http.entity.BasicHttpEntity; @@ -42,7 +48,6 @@ import net.snowflake.client.jdbc.internal.google.common.util.concurrent.ThreadFactoryBuilder; import net.snowflake.ingest.TestUtils; import net.snowflake.ingest.connection.RequestBuilder; -import net.snowflake.ingest.utils.Constants; import net.snowflake.ingest.utils.ParameterProvider; import net.snowflake.ingest.utils.SFException; import org.junit.Assert; @@ -56,7 +61,7 @@ @RunWith(PowerMockRunner.class) @PrepareForTest({TestUtils.class, HttpUtil.class, SnowflakeFileTransferAgent.class}) -public class StreamingIngestStageTest { +public class StreamingIngestStorageTest { private final String prefix = "EXAMPLE_PREFIX"; @@ -114,15 +119,16 @@ public void testPutRemote() throws Exception { byte[] dataBytes = "Hello Upload".getBytes(StandardCharsets.UTF_8); - StreamingIngestStage stage = - new StreamingIngestStage( - true, - "role", - null, - null, + StorageManager storageManager = Mockito.mock(StorageManager.class); + Mockito.when(storageManager.getClientPrefix()).thenReturn("testPrefix"); + + StreamingIngestStorage stage = + new StreamingIngestStorage( + storageManager, "clientName", - new StreamingIngestStage.SnowflakeFileTransferMetadataWithAge( + new StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge( originalMetadata, Optional.of(System.currentTimeMillis())), + null, 1); PowerMockito.mockStatic(SnowflakeFileTransferAgent.class); @@ -156,16 +162,14 @@ public void testPutLocal() throws Exception { String fullFilePath = "testOutput"; String fileName = "putLocalOutput"; - StreamingIngestStage stage = + StreamingIngestStorage stage = Mockito.spy( - new StreamingIngestStage( - true, - "role", - null, + new StreamingIngestStorage( null, "clientName", - new StreamingIngestStage.SnowflakeFileTransferMetadataWithAge( + new StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge( fullFilePath, Optional.of(System.currentTimeMillis())), + null, 1)); Mockito.doReturn(true).when(stage).isLocalFS(); @@ -186,15 +190,16 @@ public void doTestPutRemoteRefreshes() throws Exception { byte[] dataBytes = "Hello Upload".getBytes(StandardCharsets.UTF_8); - StreamingIngestStage stage = - new StreamingIngestStage( - true, - "role", - null, - null, + StorageManager storageManager = Mockito.mock(StorageManager.class); + Mockito.when(storageManager.getClientPrefix()).thenReturn("testPrefix"); + + StreamingIngestStorage stage = + new StreamingIngestStorage( + storageManager, "clientName", - new StreamingIngestStage.SnowflakeFileTransferMetadataWithAge( + new StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge( originalMetadata, Optional.of(System.currentTimeMillis())), + null, maxUploadRetryCount); PowerMockito.mockStatic(SnowflakeFileTransferAgent.class); SnowflakeSQLException e = @@ -240,16 +245,17 @@ public void testPutRemoteGCS() throws Exception { byte[] dataBytes = "Hello Upload".getBytes(StandardCharsets.UTF_8); - StreamingIngestStage stage = + StorageManager storageManager = Mockito.mock(StorageManager.class); + Mockito.when(storageManager.getClientPrefix()).thenReturn("testPrefix"); + + StreamingIngestStorage stage = Mockito.spy( - new StreamingIngestStage( - true, - "role", - null, - null, + new StreamingIngestStorage( + storageManager, "clientName", - new StreamingIngestStage.SnowflakeFileTransferMetadataWithAge( + new StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge( originalMetadata, Optional.of(System.currentTimeMillis())), + null, 1)); PowerMockito.mockStatic(SnowflakeFileTransferAgent.class); SnowflakeFileTransferMetadataV1 metaMock = Mockito.mock(SnowflakeFileTransferMetadataV1.class); @@ -265,22 +271,31 @@ public void testRefreshSnowflakeMetadataRemote() throws Exception { RequestBuilder mockBuilder = Mockito.mock(RequestBuilder.class); CloseableHttpClient mockClient = Mockito.mock(CloseableHttpClient.class); CloseableHttpResponse mockResponse = Mockito.mock(CloseableHttpResponse.class); + SnowflakeStreamingIngestClientInternal mockClientInternal = + Mockito.mock(SnowflakeStreamingIngestClientInternal.class); + Mockito.when(mockClientInternal.getRole()).thenReturn("role"); StatusLine mockStatusLine = Mockito.mock(StatusLine.class); Mockito.when(mockStatusLine.getStatusCode()).thenReturn(200); - BasicHttpEntity entity = new BasicHttpEntity(); - entity.setContent( - new ByteArrayInputStream(exampleRemoteMetaResponse.getBytes(StandardCharsets.UTF_8))); - Mockito.when(mockResponse.getStatusLine()).thenReturn(mockStatusLine); - Mockito.when(mockResponse.getEntity()).thenReturn(entity); + Mockito.when(mockResponse.getEntity()).thenReturn(createHttpEntity(exampleRemoteMetaResponse)); Mockito.when(mockClient.execute(Mockito.any())).thenReturn(mockResponse); + SnowflakeServiceClient snowflakeServiceClient = + new SnowflakeServiceClient(mockClient, mockBuilder); + StorageManager storageManager = + new InternalStageManager(true, "role", "client", snowflakeServiceClient); + ParameterProvider parameterProvider = new ParameterProvider(false); - StreamingIngestStage stage = - new StreamingIngestStage(true, "role", mockClient, mockBuilder, "clientName", 1); + StreamingIngestStorage stage = + new StreamingIngestStorage( + storageManager, + "clientName", + (StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge) null, + null, + 1); - StreamingIngestStage.SnowflakeFileTransferMetadataWithAge metadataWithAge = + StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge metadataWithAge = stage.refreshSnowflakeMetadata(true); final ArgumentCaptor endpointCaptor = ArgumentCaptor.forClass(String.class); @@ -288,7 +303,7 @@ public void testRefreshSnowflakeMetadataRemote() throws Exception { Mockito.verify(mockBuilder) .generateStreamingIngestPostRequest( stringCaptor.capture(), endpointCaptor.capture(), Mockito.any()); - Assert.assertEquals(Constants.CLIENT_CONFIGURE_ENDPOINT, endpointCaptor.getValue()); + Assert.assertEquals(CLIENT_CONFIGURE_ENDPOINT, endpointCaptor.getValue()); Assert.assertTrue(metadataWithAge.timestamp.isPresent()); Assert.assertEquals( StageInfo.StageType.S3, metadataWithAge.fileTransferMetadata.getStageInfo().getStageType()); @@ -299,7 +314,7 @@ public void testRefreshSnowflakeMetadataRemote() throws Exception { Assert.assertEquals( Paths.get("placeholder").toAbsolutePath(), Paths.get(metadataWithAge.fileTransferMetadata.getPresignedUrlFileName()).toAbsolutePath()); - Assert.assertEquals(prefix + "_" + deploymentId, stage.getClientPrefix()); + Assert.assertEquals("testPrefix", storageManager.getClientPrefix()); } @Test @@ -307,19 +322,27 @@ public void testFetchSignedURL() throws Exception { RequestBuilder mockBuilder = Mockito.mock(RequestBuilder.class); CloseableHttpClient mockClient = Mockito.mock(CloseableHttpClient.class); CloseableHttpResponse mockResponse = Mockito.mock(CloseableHttpResponse.class); + SnowflakeStreamingIngestClientInternal mockClientInternal = + Mockito.mock(SnowflakeStreamingIngestClientInternal.class); + Mockito.when(mockClientInternal.getRole()).thenReturn("role"); + SnowflakeServiceClient snowflakeServiceClient = + new SnowflakeServiceClient(mockClient, mockBuilder); + StorageManager storageManager = + new InternalStageManager(true, "role", "client", snowflakeServiceClient); StatusLine mockStatusLine = Mockito.mock(StatusLine.class); Mockito.when(mockStatusLine.getStatusCode()).thenReturn(200); - BasicHttpEntity entity = new BasicHttpEntity(); - entity.setContent( - new ByteArrayInputStream(exampleRemoteMetaResponse.getBytes(StandardCharsets.UTF_8))); - Mockito.when(mockResponse.getStatusLine()).thenReturn(mockStatusLine); - Mockito.when(mockResponse.getEntity()).thenReturn(entity); + Mockito.when(mockResponse.getEntity()).thenReturn(createHttpEntity(exampleRemoteMetaResponse)); Mockito.when(mockClient.execute(Mockito.any())).thenReturn(mockResponse); - StreamingIngestStage stage = - new StreamingIngestStage(true, "role", mockClient, mockBuilder, "clientName", 1); + StreamingIngestStorage stage = + new StreamingIngestStorage( + storageManager, + "clientName", + (StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge) null, + null, + 1); SnowflakeFileTransferMetadataV1 metadata = stage.fetchSignedURL("path/fileName"); @@ -328,7 +351,7 @@ public void testFetchSignedURL() throws Exception { Mockito.verify(mockBuilder) .generateStreamingIngestPostRequest( stringCaptor.capture(), endpointCaptor.capture(), Mockito.any()); - Assert.assertEquals(Constants.CLIENT_CONFIGURE_ENDPOINT, endpointCaptor.getValue()); + Assert.assertEquals(CLIENT_CONFIGURE_ENDPOINT, endpointCaptor.getValue()); Assert.assertEquals(StageInfo.StageType.S3, metadata.getStageInfo().getStageType()); Assert.assertEquals("foo/streaming_ingest/", metadata.getStageInfo().getLocation()); Assert.assertEquals("path/fileName", metadata.getPresignedUrlFileName()); @@ -345,26 +368,26 @@ public void testRefreshSnowflakeMetadataSynchronized() throws Exception { RequestBuilder mockBuilder = Mockito.mock(RequestBuilder.class); CloseableHttpClient mockClient = Mockito.mock(CloseableHttpClient.class); CloseableHttpResponse mockResponse = Mockito.mock(CloseableHttpResponse.class); + SnowflakeStreamingIngestClientInternal mockClientInternal = + Mockito.mock(SnowflakeStreamingIngestClientInternal.class); + Mockito.when(mockClientInternal.getRole()).thenReturn("role"); + SnowflakeServiceClient snowflakeServiceClient = + new SnowflakeServiceClient(mockClient, mockBuilder); + StorageManager storageManager = + new InternalStageManager(true, "role", "client", snowflakeServiceClient); StatusLine mockStatusLine = Mockito.mock(StatusLine.class); Mockito.when(mockStatusLine.getStatusCode()).thenReturn(200); - BasicHttpEntity entity = new BasicHttpEntity(); - entity.setContent( - new ByteArrayInputStream(exampleRemoteMetaResponse.getBytes(StandardCharsets.UTF_8))); - Mockito.when(mockResponse.getStatusLine()).thenReturn(mockStatusLine); - Mockito.when(mockResponse.getEntity()).thenReturn(entity); + Mockito.when(mockResponse.getEntity()).thenReturn(createHttpEntity(exampleRemoteMetaResponse)); Mockito.when(mockClient.execute(Mockito.any())).thenReturn(mockResponse); - StreamingIngestStage stage = - new StreamingIngestStage( - true, - "role", - mockClient, - mockBuilder, + StreamingIngestStorage stage = + new StreamingIngestStorage( + storageManager, "clientName", - new StreamingIngestStage.SnowflakeFileTransferMetadataWithAge( - originalMetadata, Optional.of(0L)), + (StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge) null, + null, 1); ThreadFactory buildUploadThreadFactory = @@ -493,15 +516,16 @@ public void testRefreshMetadataOnFirstPutException() throws Exception { byte[] dataBytes = "Hello Upload".getBytes(StandardCharsets.UTF_8); - StreamingIngestStage stage = - new StreamingIngestStage( - true, - "role", - null, - null, + StorageManager storageManager = Mockito.mock(StorageManager.class); + Mockito.when(storageManager.getClientPrefix()).thenReturn("testPrefix"); + + StreamingIngestStorage stage = + new StreamingIngestStorage( + storageManager, "clientName", - new StreamingIngestStage.SnowflakeFileTransferMetadataWithAge( + new StreamingIngestStorage.SnowflakeFileTransferMetadataWithAge( originalMetadata, Optional.of(System.currentTimeMillis())), + null, maxUploadRetryCount); PowerMockito.mockStatic(SnowflakeFileTransferAgent.class); SnowflakeSQLException e = @@ -546,4 +570,10 @@ public Object answer(org.mockito.invocation.InvocationOnMock invocation) InputStream capturedInput = capturedConfig.getUploadStream(); Assert.assertEquals("Hello Upload", IOUtils.toString(capturedInput)); } + + private HttpEntity createHttpEntity(String content) { + BasicHttpEntity entity = new BasicHttpEntity(); + entity.setContent(new ByteArrayInputStream(content.getBytes(StandardCharsets.UTF_8))); + return entity; + } } diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtilsIT.java b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtilsIT.java index 4e054c209..b40cdad82 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtilsIT.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtilsIT.java @@ -1,3 +1,7 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.streaming.internal; import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CLIENT_CONFIGURE; @@ -5,9 +9,6 @@ import static net.snowflake.ingest.utils.Constants.CLIENT_CONFIGURE_ENDPOINT; import static net.snowflake.ingest.utils.Constants.RESPONSE_SUCCESS; -import com.fasterxml.jackson.databind.ObjectMapper; -import java.util.HashMap; -import java.util.Map; import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient; import net.snowflake.ingest.TestUtils; import net.snowflake.ingest.connection.IngestResponseException; @@ -53,11 +54,11 @@ public void testJWTRetries() throws Exception { "testJWTRetries")); // build payload - Map payload = new HashMap<>(); - if (!TestUtils.getRole().isEmpty() && !TestUtils.getRole().equals("DEFAULT_ROLE")) { - payload.put("role", TestUtils.getRole()); - } - ObjectMapper mapper = new ObjectMapper(); + ClientConfigureRequest request = + new ClientConfigureRequest( + !TestUtils.getRole().isEmpty() && !TestUtils.getRole().equals("DEFAULT_ROLE") + ? TestUtils.getRole() + : null); // request wih invalid token, should get 401 3 times PowerMockito.doReturn("invalid_token").when(spyManager).getToken(); @@ -66,7 +67,7 @@ public void testJWTRetries() throws Exception { executeWithRetries( ChannelsStatusResponse.class, CLIENT_CONFIGURE_ENDPOINT, - mapper.writeValueAsString(payload), + request, "client configure", STREAMING_CLIENT_CONFIGURE, httpClient, @@ -84,7 +85,7 @@ public void testJWTRetries() throws Exception { executeWithRetries( ChannelsStatusResponse.class, CLIENT_CONFIGURE_ENDPOINT, - mapper.writeValueAsString(payload), + request, "client configure", STREAMING_CLIENT_CONFIGURE, httpClient, @@ -101,7 +102,7 @@ public void testJWTRetries() throws Exception { executeWithRetries( ChannelsStatusResponse.class, CLIENT_CONFIGURE_ENDPOINT, - mapper.writeValueAsString(payload), + request, "client configure", STREAMING_CLIENT_CONFIGURE, httpClient, From 6eb0fff22feb8b3a7af5c31b7b3b03375dc3b9af Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Thu, 11 Jul 2024 13:35:32 -0700 Subject: [PATCH 4/9] post merge --- .../internal/ChannelsStatusRequest.java | 3 +-- .../streaming/internal/BlobBuilderTest.java | 16 ++++++++++++---- .../internal/InsertRowsBenchmarkTest.java | 6 +++++- 3 files changed, 18 insertions(+), 7 deletions(-) diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java index 4c6f05ac8..5a17b6e91 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java @@ -92,8 +92,7 @@ List getChannels() { @Override public String getStringForLogging() { return String.format( - "ChannelsStatusRequest(requestId=%s, role=%s, channels=[%s])", - requestId, + "ChannelsStatusRequest(role=%s, channels=[%s])", role, channels.stream() .map( diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/BlobBuilderTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/BlobBuilderTest.java index e220aec79..32a352ba9 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/BlobBuilderTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/BlobBuilderTest.java @@ -1,3 +1,7 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.streaming.internal; import java.io.ByteArrayOutputStream; @@ -23,18 +27,21 @@ public void testSerializationErrors() throws Exception { BlobBuilder.constructBlobAndMetadata( "a.bdec", Collections.singletonList(createChannelDataPerTable(1, false)), - Constants.BdecVersion.THREE); + Constants.BdecVersion.THREE, + true); BlobBuilder.constructBlobAndMetadata( "a.bdec", Collections.singletonList(createChannelDataPerTable(1, true)), - Constants.BdecVersion.THREE); + Constants.BdecVersion.THREE, + true); // Construction fails if metadata contains 0 rows and data 1 row try { BlobBuilder.constructBlobAndMetadata( "a.bdec", Collections.singletonList(createChannelDataPerTable(0, false)), - Constants.BdecVersion.THREE); + Constants.BdecVersion.THREE, + true); Assert.fail("Should not pass enableParquetInternalBuffering=false"); } catch (SFException e) { Assert.assertEquals(ErrorCode.INTERNAL_ERROR.getMessageCode(), e.getVendorCode()); @@ -52,7 +59,8 @@ public void testSerializationErrors() throws Exception { BlobBuilder.constructBlobAndMetadata( "a.bdec", Collections.singletonList(createChannelDataPerTable(0, true)), - Constants.BdecVersion.THREE); + Constants.BdecVersion.THREE, + true); Assert.fail("Should not pass enableParquetInternalBuffering=true"); } catch (SFException e) { Assert.assertEquals(ErrorCode.INTERNAL_ERROR.getMessageCode(), e.getVendorCode()); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/InsertRowsBenchmarkTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/InsertRowsBenchmarkTest.java index 5b28e9c45..d28b22669 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/InsertRowsBenchmarkTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/InsertRowsBenchmarkTest.java @@ -1,3 +1,7 @@ +/* + * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. + */ + package net.snowflake.ingest.streaming.internal; import static java.time.ZoneOffset.UTC; @@ -36,7 +40,7 @@ public class InsertRowsBenchmarkTest { @Setup(Level.Trial) public void setUpBeforeAll() { - client = new SnowflakeStreamingIngestClientInternal("client_PARQUET"); + client = new SnowflakeStreamingIngestClientInternal("client_PARQUET", false); channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", From c1d7116602a1fa703980ad0756d85ee7713991a3 Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Thu, 11 Jul 2024 14:51:17 -0700 Subject: [PATCH 5/9] eliminate iceberg logic --- .../connection/ServiceResponseHandler.java | 5 +- ...SnowflakeStreamingIngestClientFactory.java | 16 +- .../streaming/internal/BlobBuilder.java | 64 +++---- .../streaming/internal/ChannelCache.java | 2 +- .../internal/ChannelConfigureRequest.java | 53 ------ .../internal/ChannelConfigureResponse.java | 46 ----- .../internal/DropChannelRequestInternal.java | 15 +- .../internal/ExternalVolumeManager.java | 173 ------------------ .../streaming/internal/FlushService.java | 7 +- .../internal/InternalParameterProvider.java | 20 -- .../internal/OpenChannelRequestInternal.java | 15 +- .../internal/SnowflakeServiceClient.java | 28 --- ...nowflakeStreamingIngestClientInternal.java | 41 +---- .../internal/StreamingIngestStorage.java | 16 -- .../net/snowflake/ingest/utils/Constants.java | 3 +- .../net/snowflake/ingest/utils/ErrorCode.java | 3 +- .../ingest/utils/ParameterProvider.java | 129 ++++--------- .../streaming/internal/BlobBuilderTest.java | 16 +- .../streaming/internal/ChannelCacheTest.java | 2 +- .../streaming/internal/FlushServiceTest.java | 92 ++++------ .../internal/InsertRowsBenchmarkTest.java | 6 +- .../InternalParameterProviderTest.java | 26 --- .../streaming/internal/OAuthBasicTest.java | 2 +- .../internal/ParameterProviderTest.java | 98 +++------- .../internal/RegisterServiceTest.java | 15 +- .../SnowflakeStreamingIngestChannelTest.java | 46 ++--- .../SnowflakeStreamingIngestClientTest.java | 40 +--- .../internal/StreamingIngestStorageTest.java | 2 - 28 files changed, 161 insertions(+), 820 deletions(-) delete mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java delete mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java delete mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java delete mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/InternalParameterProvider.java delete mode 100644 src/test/java/net/snowflake/ingest/streaming/internal/InternalParameterProviderTest.java diff --git a/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java b/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java index 034b4a6f0..cef6e631c 100644 --- a/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java +++ b/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2012-2017 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.connection; @@ -42,8 +42,7 @@ public enum ApiName { STREAMING_DROP_CHANNEL("POST"), STREAMING_CHANNEL_STATUS("POST"), STREAMING_REGISTER_BLOB("POST"), - STREAMING_CLIENT_CONFIGURE("POST"), - STREAMING_CHANNEL_CONFIGURE("POST"); + STREAMING_CLIENT_CONFIGURE("POST"); private final String httpMethod; private ApiName(String httpMethod) { diff --git a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java index 89e528693..cd6d78787 100644 --- a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java +++ b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java @@ -28,10 +28,6 @@ public static class Builder { // Allows client to override some default parameter values private Map parameterOverrides; - // Indicates whether it's streaming to Iceberg tables. Open channels on regular tables should - // fail in this mode. - private boolean isIcebergMode; - // Indicates whether it's under test mode private boolean isTestMode; @@ -49,11 +45,6 @@ public Builder setParameterOverrides(Map parameterOverrides) { return this; } - public Builder setIsIceberg(boolean isIcebergMode) { - this.isIcebergMode = isIcebergMode; - return this; - } - public Builder setIsTestMode(boolean isTestMode) { this.isTestMode = isTestMode; return this; @@ -67,12 +58,7 @@ public SnowflakeStreamingIngestClient build() { SnowflakeURL accountURL = new SnowflakeURL(prop.getProperty(Constants.ACCOUNT_URL)); return new SnowflakeStreamingIngestClientInternal<>( - this.name, - accountURL, - prop, - this.parameterOverrides, - this.isIcebergMode, - this.isTestMode); + this.name, accountURL, prop, this.parameterOverrides, this.isTestMode); } } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java b/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java index 75a623f69..b88090e01 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java @@ -61,14 +61,10 @@ class BlobBuilder { * @param blobData All the data for one blob. Assumes that all ChannelData in the inner List * belongs to the same table. Will error if this is not the case * @param bdecVersion version of blob - * @param encrypt If the output chunk is encrypted or not * @return {@link Blob} data */ static Blob constructBlobAndMetadata( - String filePath, - List>> blobData, - Constants.BdecVersion bdecVersion, - boolean encrypt) + String filePath, List>> blobData, Constants.BdecVersion bdecVersion) throws IOException, NoSuchPaddingException, NoSuchAlgorithmException, InvalidAlgorithmParameterException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException { @@ -87,32 +83,25 @@ static Blob constructBlobAndMetadata( flusher.serialize(channelsDataPerTable, filePath); if (!serializedChunk.channelsMetadataList.isEmpty()) { - byte[] compressedChunkData; - int chunkLength; + ByteArrayOutputStream chunkData = serializedChunk.chunkData; + Pair paddedChunk = + padChunk(chunkData, Constants.ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES); + byte[] paddedChunkData = paddedChunk.getFirst(); + int paddedChunkLength = paddedChunk.getSecond(); - if (encrypt) { - Pair paddedChunk = - padChunk(serializedChunk.chunkData, Constants.ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES); - byte[] paddedChunkData = paddedChunk.getFirst(); - chunkLength = paddedChunk.getSecond(); - - // Encrypt the compressed chunk data, the encryption key is derived using the key from - // server with the full blob path. - // We need to maintain IV as a block counter for the whole file, even interleaved, - // to align with decryption on the Snowflake query path. - // TODO: address alignment for the header SNOW-557866 - long iv = curDataSize / Constants.ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES; - compressedChunkData = - Cryptor.encrypt( - paddedChunkData, firstChannelFlushContext.getEncryptionKey(), filePath, iv); - } else { - compressedChunkData = serializedChunk.chunkData.toByteArray(); - chunkLength = compressedChunkData.length; - } + // Encrypt the compressed chunk data, the encryption key is derived using the key from + // server with the full blob path. + // We need to maintain IV as a block counter for the whole file, even interleaved, + // to align with decryption on the Snowflake query path. + // TODO: address alignment for the header SNOW-557866 + long iv = curDataSize / Constants.ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES; + byte[] encryptedCompressedChunkData = + Cryptor.encrypt( + paddedChunkData, firstChannelFlushContext.getEncryptionKey(), filePath, iv); // Compute the md5 of the chunk data - String md5 = computeMD5(compressedChunkData, chunkLength); - int compressedChunkDataSize = compressedChunkData.length; + String md5 = computeMD5(encryptedCompressedChunkData, paddedChunkLength); + int encryptedCompressedChunkDataSize = encryptedCompressedChunkData.length; // Create chunk metadata long startOffset = curDataSize; @@ -122,9 +111,9 @@ static Blob constructBlobAndMetadata( // The start offset will be updated later in BlobBuilder#build to include the blob // header .setChunkStartOffset(startOffset) - // The chunkLength is used because it is the actual data size used for + // The paddedChunkLength is used because it is the actual data size used for // decompression and md5 calculation on server side. - .setChunkLength(chunkLength) + .setChunkLength(paddedChunkLength) .setUncompressedChunkLength((int) serializedChunk.chunkEstimatedUncompressedSize) .setChannelList(serializedChunk.channelsMetadataList) .setChunkMD5(md5) @@ -138,22 +127,21 @@ static Blob constructBlobAndMetadata( // Add chunk metadata and data to the list chunksMetadataList.add(chunkMetadata); - chunksDataList.add(compressedChunkData); - curDataSize += compressedChunkDataSize; - crc.update(compressedChunkData, 0, compressedChunkDataSize); + chunksDataList.add(encryptedCompressedChunkData); + curDataSize += encryptedCompressedChunkDataSize; + crc.update(encryptedCompressedChunkData, 0, encryptedCompressedChunkDataSize); logger.logInfo( "Finish building chunk in blob={}, table={}, rowCount={}, startOffset={}," - + " estimatedUncompressedSize={}, chunkLength={}, compressedSize={}," - + " encryption={}, bdecVersion={}", + + " estimatedUncompressedSize={}, paddedChunkLength={}, encryptedCompressedSize={}," + + " bdecVersion={}", filePath, firstChannelFlushContext.getFullyQualifiedTableName(), serializedChunk.rowCount, startOffset, serializedChunk.chunkEstimatedUncompressedSize, - chunkLength, - compressedChunkDataSize, - encrypt, + paddedChunkLength, + encryptedCompressedChunkDataSize, bdecVersion); } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java index 539da6287..989be0fa1 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.streaming.internal; diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java deleted file mode 100644 index 31bca95ac..000000000 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureRequest.java +++ /dev/null @@ -1,53 +0,0 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - -package net.snowflake.ingest.streaming.internal; - -import com.fasterxml.jackson.annotation.JsonProperty; - -/** Class used to serialize the channel configure request. */ -class ChannelConfigureRequest extends ConfigureRequest { - @JsonProperty("database") - private String database; - - @JsonProperty("schema") - private String schema; - - @JsonProperty("table") - private String table; - - /** - * Constructor for channel configure request - * - * @param role Role to be used for the request. - * @param database Database name. - * @param schema Schema name. - * @param table Table name. - */ - ChannelConfigureRequest(String role, String database, String schema, String table) { - setRole(role); - this.database = database; - this.schema = schema; - this.table = table; - } - - String getDatabase() { - return database; - } - - String getSchema() { - return schema; - } - - String getTable() { - return table; - } - - @Override - public String getStringForLogging() { - return String.format( - "ChannelConfigureRequest(role=%s, db=%s, schema=%s, table=%s, file_name=%s)", - getRole(), database, schema, table, getFileName()); - } -} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java deleted file mode 100644 index da65960b4..000000000 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelConfigureResponse.java +++ /dev/null @@ -1,46 +0,0 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - -package net.snowflake.ingest.streaming.internal; - -import com.fasterxml.jackson.annotation.JsonIgnoreProperties; -import com.fasterxml.jackson.annotation.JsonProperty; - -/** Class used to deserialize responses from channel configure endpoint */ -@JsonIgnoreProperties(ignoreUnknown = true) -class ChannelConfigureResponse extends StreamingIngestResponse { - @JsonProperty("status_code") - private Long statusCode; - - @JsonProperty("message") - private String message; - - @JsonProperty("stage_location") - private FileLocationInfo stageLocation; - - @Override - Long getStatusCode() { - return statusCode; - } - - void setStatusCode(Long statusCode) { - this.statusCode = statusCode; - } - - String getMessage() { - return message; - } - - void setMessage(String message) { - this.message = message; - } - - FileLocationInfo getStageLocation() { - return stageLocation; - } - - void setStageLocation(FileLocationInfo stageLocation) { - this.stageLocation = stageLocation; - } -} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java index 25599999e..a6e1e1746 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java @@ -29,9 +29,6 @@ class DropChannelRequestInternal implements StreamingIngestRequest { @JsonProperty("schema") private String schema; - @JsonProperty("is_iceberg") - private boolean isIceberg; - @JsonInclude(JsonInclude.Include.NON_NULL) @JsonProperty("client_sequencer") Long clientSequencer; @@ -43,8 +40,7 @@ class DropChannelRequestInternal implements StreamingIngestRequest { String schema, String table, String channel, - Long clientSequencer, - boolean isIceberg) { + Long clientSequencer) { this.requestId = requestId; this.role = role; this.database = database; @@ -52,7 +48,6 @@ class DropChannelRequestInternal implements StreamingIngestRequest { this.table = table; this.channel = channel; this.clientSequencer = clientSequencer; - this.isIceberg = isIceberg; } String getRequestId() { @@ -79,10 +74,6 @@ String getSchema() { return schema; } - boolean getIsIceberg() { - return isIceberg; - } - Long getClientSequencer() { return clientSequencer; } @@ -95,7 +86,7 @@ String getFullyQualifiedTableName() { public String getStringForLogging() { return String.format( "DropChannelRequest(requestId=%s, role=%s, db=%s, schema=%s, table=%s, channel=%s," - + " isIceberg=%s, clientSequencer=%s)", - requestId, role, database, schema, table, channel, isIceberg, clientSequencer); + + " clientSequencer=%s)", + requestId, role, database, schema, table, channel, clientSequencer); } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java deleted file mode 100644 index 2b8796bb2..000000000 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ExternalVolumeManager.java +++ /dev/null @@ -1,173 +0,0 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - -package net.snowflake.ingest.streaming.internal; - -import java.io.IOException; -import java.util.Calendar; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.ConcurrentHashMap; -import net.snowflake.client.jdbc.SnowflakeSQLException; -import net.snowflake.ingest.connection.IngestResponseException; -import net.snowflake.ingest.utils.ErrorCode; -import net.snowflake.ingest.utils.SFException; -import net.snowflake.ingest.utils.Utils; - -class ExternalVolumeLocation { - public final String dbName; - public final String schemaName; - public final String tableName; - - public ExternalVolumeLocation(String dbName, String schemaName, String tableName) { - this.dbName = dbName; - this.schemaName = schemaName; - this.tableName = tableName; - } -} - -/** Class to manage multiple external volumes */ -class ExternalVolumeManager implements StorageManager { - // Reference to the external volume per table - private final Map> externalVolumeMap; - - // name of the owning client - private final String clientName; - - // role of the owning client - private final String role; - - // Reference to the Snowflake service client used for configure calls - private final SnowflakeServiceClient snowflakeServiceClient; - - // Client prefix generated by the Snowflake server - private final String clientPrefix; - - /** - * Constructor for ExternalVolumeManager - * - * @param isTestMode whether the manager in test mode - * @param role the role of the client - * @param clientName the name of the client - * @param snowflakeServiceClient the Snowflake service client used for configure calls - */ - ExternalVolumeManager( - boolean isTestMode, - String role, - String clientName, - SnowflakeServiceClient snowflakeServiceClient) { - this.role = role; - this.clientName = clientName; - this.snowflakeServiceClient = snowflakeServiceClient; - this.externalVolumeMap = new ConcurrentHashMap<>(); - try { - this.clientPrefix = - isTestMode - ? "testPrefix" - : this.snowflakeServiceClient - .clientConfigure(new ClientConfigureRequest(role)) - .getClientPrefix(); - } catch (IngestResponseException | IOException e) { - throw new SFException(e, ErrorCode.CLIENT_CONFIGURE_FAILURE, e.getMessage()); - } - } - - /** - * Given a fully qualified table name, return the target storage by looking up the table name - * - * @param fullyQualifiedTableName the target fully qualified table name - * @return target storage - */ - @Override - public StreamingIngestStorage getStorage( - String fullyQualifiedTableName) { - // Only one chunk per blob in Iceberg mode. - StreamingIngestStorage stage = - this.externalVolumeMap.get(fullyQualifiedTableName); - - if (stage == null) { - throw new SFException( - ErrorCode.INTERNAL_ERROR, - String.format("No external volume found for table %s", fullyQualifiedTableName)); - } - - return stage; - } - - /** - * Add a storage to the manager by looking up the table name from the open channel response - * - * @param dbName the database name - * @param schemaName the schema name - * @param tableName the table name - * @param fileLocationInfo response from open channel - */ - @Override - public void addStorage( - String dbName, String schemaName, String tableName, FileLocationInfo fileLocationInfo) { - String fullyQualifiedTableName = - Utils.getFullyQualifiedTableName(dbName, schemaName, tableName); - - try { - this.externalVolumeMap.put( - fullyQualifiedTableName, - new StreamingIngestStorage( - this, - this.clientName, - fileLocationInfo, - new ExternalVolumeLocation(dbName, schemaName, tableName), - DEFAULT_MAX_UPLOAD_RETRIES)); - } catch (SnowflakeSQLException | IOException err) { - throw new SFException(err, ErrorCode.UNABLE_TO_CONNECT_TO_STORAGE); - } - } - - /** - * Gets the latest file location info (with a renewed short-lived access token) for the specified - * location - * - * @param location A reference to the target location - * @param fileName optional filename for single-file signed URL fetch from server - * @return the new location information - */ - @Override - public FileLocationInfo getRefreshedLocation( - ExternalVolumeLocation location, Optional fileName) { - try { - ChannelConfigureRequest request = - new ChannelConfigureRequest( - this.role, location.dbName, location.schemaName, location.tableName); - fileName.ifPresent(request::setFileName); - ChannelConfigureResponse response = this.snowflakeServiceClient.channelConfigure(request); - return response.getStageLocation(); - } catch (IngestResponseException | IOException e) { - throw new SFException(e, ErrorCode.CLIENT_CONFIGURE_FAILURE, e.getMessage()); - } - } - - // TODO: SNOW-1502887 Blob path generation for iceberg table - @Override - public String generateBlobPath() { - return "snow_dummy_file_name"; - } - - // TODO: SNOW-1502887 Blob path generation for iceberg table - @Override - public void decrementBlobSequencer() {} - - // TODO: SNOW-1502887 Blob path generation for iceberg table - public String getBlobPath(Calendar calendar, String clientPrefix) { - return ""; - } - - /** - * Get the client prefix from first external volume in the map - * - * @return the client prefix - */ - @Override - public String getClientPrefix() { - return this.clientPrefix; - } -} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java index 7f3aafd56..b34335194 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java @@ -522,12 +522,7 @@ BlobMetadata buildAndUpload( Timer.Context buildContext = Utils.createTimerContext(this.owningClient.buildLatency); // Construct the blob along with the metadata of the blob - BlobBuilder.Blob blob = - BlobBuilder.constructBlobAndMetadata( - blobPath, - blobData, - bdecVersion, - this.owningClient.getInternalParameterProvider().getEnableChunkEncryption()); + BlobBuilder.Blob blob = BlobBuilder.constructBlobAndMetadata(blobPath, blobData, bdecVersion); blob.blobStats.setBuildDurationMs(buildContext); diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/InternalParameterProvider.java b/src/main/java/net/snowflake/ingest/streaming/internal/InternalParameterProvider.java deleted file mode 100644 index 46a99b671..000000000 --- a/src/main/java/net/snowflake/ingest/streaming/internal/InternalParameterProvider.java +++ /dev/null @@ -1,20 +0,0 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - -package net.snowflake.ingest.streaming.internal; - -/** A class to provide non-configurable constants depends on Iceberg or non-Iceberg mode */ -class InternalParameterProvider { - private final boolean isIcebergMode; - - InternalParameterProvider(boolean isIcebergMode) { - this.isIcebergMode = isIcebergMode; - } - - boolean getEnableChunkEncryption() { - // When in Iceberg mode, chunk encryption is disabled. Otherwise, it is enabled. Since Iceberg - // mode does not need client-side encryption. - return !isIcebergMode; - } -} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelRequestInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelRequestInternal.java index ffbb553f3..f58198d31 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelRequestInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelRequestInternal.java @@ -32,9 +32,6 @@ class OpenChannelRequestInternal implements StreamingIngestRequest { @JsonProperty("write_mode") private String writeMode; - @JsonProperty("is_iceberg") - private boolean isIceberg; - @JsonInclude(JsonInclude.Include.NON_NULL) @JsonProperty("offset_token") private String offsetToken; @@ -47,8 +44,7 @@ class OpenChannelRequestInternal implements StreamingIngestRequest { String table, String channel, Constants.WriteMode writeMode, - String offsetToken, - boolean isIceberg) { + String offsetToken) { this.requestId = requestId; this.role = role; this.database = database; @@ -56,7 +52,6 @@ class OpenChannelRequestInternal implements StreamingIngestRequest { this.table = table; this.channel = channel; this.writeMode = writeMode.name(); - this.isIceberg = isIceberg; this.offsetToken = offsetToken; } @@ -88,10 +83,6 @@ String getWriteMode() { return writeMode; } - boolean getIsIceberg() { - return isIceberg; - } - String getOffsetToken() { return offsetToken; } @@ -100,7 +91,7 @@ String getOffsetToken() { public String getStringForLogging() { return String.format( "OpenChannelRequestInternal(requestId=%s, role=%s, db=%s, schema=%s, table=%s, channel=%s," - + " writeMode=%s, isIceberg=%s)", - requestId, role, database, schema, table, channel, writeMode, isIceberg); + + " writeMode=%s)", + requestId, role, database, schema, table, channel, writeMode); } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java index 947c86dbb..904e4f93a 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java @@ -4,14 +4,12 @@ package net.snowflake.ingest.streaming.internal; -import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CHANNEL_CONFIGURE; import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CHANNEL_STATUS; import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_CLIENT_CONFIGURE; import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_DROP_CHANNEL; import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_OPEN_CHANNEL; import static net.snowflake.ingest.connection.ServiceResponseHandler.ApiName.STREAMING_REGISTER_BLOB; import static net.snowflake.ingest.streaming.internal.StreamingIngestUtils.executeWithRetries; -import static net.snowflake.ingest.utils.Constants.CHANNEL_CONFIGURE_ENDPOINT; import static net.snowflake.ingest.utils.Constants.CHANNEL_STATUS_ENDPOINT; import static net.snowflake.ingest.utils.Constants.CLIENT_CONFIGURE_ENDPOINT; import static net.snowflake.ingest.utils.Constants.DROP_CHANNEL_ENDPOINT; @@ -76,32 +74,6 @@ ClientConfigureResponse clientConfigure(ClientConfigureRequest request) return response; } - /** - * Configures a storage given a {@link ChannelConfigureRequest}. - * - * @param request the channel configuration request - * @return the response from the configuration request - */ - ChannelConfigureResponse channelConfigure(ChannelConfigureRequest request) - throws IngestResponseException, IOException { - ChannelConfigureResponse response = - executeApiRequestWithRetries( - ChannelConfigureResponse.class, - request, - CHANNEL_CONFIGURE_ENDPOINT, - "channel configure", - STREAMING_CHANNEL_CONFIGURE); - - if (response.getStatusCode() != RESPONSE_SUCCESS) { - logger.logDebug( - "Channel configure request failed, request={}, response={}", - request.getStringForLogging(), - response.getMessage()); - throw new SFException(ErrorCode.CHANNEL_CONFIGURE_FAILURE, response.getMessage()); - } - return response; - } - /** * Opens a channel given a {@link OpenChannelRequestInternal}. * diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java index 43397389d..aee73b8c4 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java @@ -94,9 +94,6 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea // Provides constant values that can be set by constructor private final ParameterProvider parameterProvider; - // Provides constant values which is determined by the Iceberg or non-Iceberg mode - private final InternalParameterProvider internalParameterProvider; - // Name of the client private final String name; @@ -118,9 +115,6 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea // Indicates whether the client has closed private volatile boolean isClosed; - // Indicates wheter the client is streaming to Iceberg tables - private final boolean isIcebergMode; - // Indicates whether the client is under test mode private final boolean isTestMode; @@ -155,7 +149,6 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea * @param accountURL Snowflake account url * @param prop connection properties * @param httpClient http client for sending request - * @param isIcebergMode whether we're streaming to iceberg tables * @param isTestMode whether we're under test mode * @param requestBuilder http request builder * @param parameterOverrides parameters we override in case we want to set different values @@ -165,16 +158,13 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea SnowflakeURL accountURL, Properties prop, CloseableHttpClient httpClient, - boolean isIcebergMode, boolean isTestMode, RequestBuilder requestBuilder, Map parameterOverrides) { - this.parameterProvider = new ParameterProvider(parameterOverrides, prop, isIcebergMode); - this.internalParameterProvider = new InternalParameterProvider(isIcebergMode); + this.parameterProvider = new ParameterProvider(parameterOverrides, prop); this.name = name; String accountName = accountURL == null ? null : accountURL.getAccount(); - this.isIcebergMode = isIcebergMode; this.isTestMode = isTestMode; this.httpClient = httpClient == null ? HttpUtil.getHttpClient(accountName) : httpClient; this.channelCache = new ChannelCache<>(); @@ -238,11 +228,7 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea this.snowflakeServiceClient = new SnowflakeServiceClient(this.httpClient, this.requestBuilder); this.storageManager = - isIcebergMode - ? new ExternalVolumeManager( - isTestMode, this.role, this.name, this.snowflakeServiceClient) - : new InternalStageManager( - isTestMode, this.role, this.name, this.snowflakeServiceClient); + new InternalStageManager(isTestMode, this.role, this.name, this.snowflakeServiceClient); try { this.flushService = @@ -268,7 +254,6 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea * @param accountURL Snowflake account url * @param prop connection properties * @param parameterOverrides map of parameters to override for this client - * @param isIcebergMode whether we're streaming to iceberg tables * @param isTestMode indicates whether it's under test mode */ public SnowflakeStreamingIngestClientInternal( @@ -276,17 +261,16 @@ public SnowflakeStreamingIngestClientInternal( SnowflakeURL accountURL, Properties prop, Map parameterOverrides, - boolean isIcebergMode, boolean isTestMode) { - this(name, accountURL, prop, null, isIcebergMode, isTestMode, null, parameterOverrides); + this(name, accountURL, prop, null, isTestMode, null, parameterOverrides); } /*** Constructor for TEST ONLY * * @param name the name of the client */ - SnowflakeStreamingIngestClientInternal(String name, boolean isIcebergMode) { - this(name, null, null, null, isIcebergMode, true, null, new HashMap<>()); + SnowflakeStreamingIngestClientInternal(String name) { + this(name, null, null, null, true, null, new HashMap<>()); } // TESTING ONLY - inject the request builder @@ -349,8 +333,7 @@ public SnowflakeStreamingIngestChannelInternal openChannel(OpenChannelRequest request.getTableName(), request.getChannelName(), Constants.WriteMode.CLOUD_STORAGE, - request.getOffsetToken(), - isIcebergMode); + request.getOffsetToken()); OpenChannelResponse response = snowflakeServiceClient.openChannel(openChannelRequest); logger.logInfo( @@ -421,8 +404,7 @@ public void dropChannel(DropChannelRequest request) { request.getChannelName(), request instanceof DropChannelVersionRequest ? ((DropChannelVersionRequest) request).getClientSequencer() - : null, - isIcebergMode); + : null); snowflakeServiceClient.dropChannel(dropChannelRequest); logger.logInfo( @@ -904,15 +886,6 @@ ParameterProvider getParameterProvider() { return parameterProvider; } - /** - * Get InternalParameterProvider with internal parameters - * - * @return {@link InternalParameterProvider} used by the client - */ - InternalParameterProvider getInternalParameterProvider() { - return internalParameterProvider; - } - /** * Set refresh token, this method is for refresh token renewal without requiring to restart * client. This method only works when the authorization type is OAuth diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java index eb9f10826..97c8e4a05 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java @@ -280,22 +280,6 @@ private SnowflakeFileTransferMetadataWithAge createFileTransferMetadataWithAge( return this.fileTransferMetadataWithAge; } - /** - * Creates a client-specific prefix that will be also part of the files registered by this client. - * The prefix will include a server-side generated string and the GlobalID of the deployment the - * client is registering blobs to. The latter (deploymentId) is needed in order to guarantee that - * blob filenames are unique across deployments even with replication enabled. - * - * @param response the client/configure response from the server - * @return the client prefix. - */ - private String createClientPrefix(final ClientConfigureResponse response) { - final String prefix = response.getPrefix() == null ? "" : response.getPrefix(); - final String deploymentId = - response.getDeploymentId() != null ? "_" + response.getDeploymentId() : ""; - return prefix + deploymentId; - } - /** * GCS requires a signed url per file. We need to fetch this from the server for each put * diff --git a/src/main/java/net/snowflake/ingest/utils/Constants.java b/src/main/java/net/snowflake/ingest/utils/Constants.java index 3d09d9a2a..3579e4d24 100644 --- a/src/main/java/net/snowflake/ingest/utils/Constants.java +++ b/src/main/java/net/snowflake/ingest/utils/Constants.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved. + * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved. */ package net.snowflake.ingest.utils; @@ -52,7 +52,6 @@ public class Constants { public static final String BLOB_EXTENSION_TYPE = "bdec"; public static final int MAX_THREAD_COUNT = Integer.MAX_VALUE; public static final String CLIENT_CONFIGURE_ENDPOINT = "/v1/streaming/client/configure/"; - public static final String CHANNEL_CONFIGURE_ENDPOINT = "/v1/streaming/channels/configure/"; public static final int COMMIT_MAX_RETRY_COUNT = 60; public static final int COMMIT_RETRY_INTERVAL_IN_MS = 1000; public static final String ENCRYPTION_ALGORITHM = "AES/CTR/NoPadding"; diff --git a/src/main/java/net/snowflake/ingest/utils/ErrorCode.java b/src/main/java/net/snowflake/ingest/utils/ErrorCode.java index 4091e98bf..565973a96 100644 --- a/src/main/java/net/snowflake/ingest/utils/ErrorCode.java +++ b/src/main/java/net/snowflake/ingest/utils/ErrorCode.java @@ -41,8 +41,7 @@ public enum ErrorCode { OAUTH_REFRESH_TOKEN_ERROR("0033"), INVALID_CONFIG_PARAMETER("0034"), CRYPTO_PROVIDER_ERROR("0035"), - DROP_CHANNEL_FAILURE("0036"), - CHANNEL_CONFIGURE_FAILURE("0037"); + DROP_CHANNEL_FAILURE("0036"); public static final String errorMessageResource = "net.snowflake.ingest.ingest_error_messages"; diff --git a/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java b/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java index 7d8cb230f..b98972a7d 100644 --- a/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java +++ b/src/main/java/net/snowflake/ingest/utils/ParameterProvider.java @@ -1,11 +1,5 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - package net.snowflake.ingest.utils; -import static net.snowflake.ingest.utils.ErrorCode.INVALID_CONFIG_PARAMETER; - import java.util.HashMap; import java.util.Map; import java.util.Properties; @@ -61,7 +55,6 @@ public class ParameterProvider { // Lag related parameters public static final long MAX_CLIENT_LAG_DEFAULT = 1000; // 1 second - public static final long MAX_CLIENT_LAG_ICEBERG_MODE_DEFAULT = 30000; // 30 seconds static final long MAX_CLIENT_LAG_MS_MIN = TimeUnit.SECONDS.toMillis(1); static final long MAX_CLIENT_LAG_MS_MAX = TimeUnit.MINUTES.toMillis(10); @@ -71,9 +64,6 @@ public class ParameterProvider { public static final Constants.BdecParquetCompression BDEC_PARQUET_COMPRESSION_ALGORITHM_DEFAULT = Constants.BdecParquetCompression.GZIP; - /* Iceberg mode parameters: When streaming to Iceberg mode, different default parameters are required because it generates Parquet files instead of BDEC files. */ - public static final int MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT = 1; - /* Parameter that enables using internal Parquet buffers for buffering of rows before serializing. It reduces memory consumption compared to using Java Objects for buffering.*/ public static final boolean ENABLE_PARQUET_INTERNAL_BUFFERING_DEFAULT = false; @@ -90,25 +80,18 @@ public class ParameterProvider { * * @param parameterOverrides Map of parameter name to value * @param props Properties from profile file - * @param isIcebergMode If the provided parameters need to be verified and modified to meet - * Iceberg mode */ - public ParameterProvider( - Map parameterOverrides, Properties props, boolean isIcebergMode) { - this.setParameterMap(parameterOverrides, props, isIcebergMode); + public ParameterProvider(Map parameterOverrides, Properties props) { + this.setParameterMap(parameterOverrides, props); } /** Empty constructor for tests */ - public ParameterProvider(boolean isIcebergMode) { - this(null, null, isIcebergMode); + public ParameterProvider() { + this(null, null); } - private void checkAndUpdate( - String key, - Object defaultValue, - Map parameterOverrides, - Properties props, - boolean enforceDefault) { + private void updateValue( + String key, Object defaultValue, Map parameterOverrides, Properties props) { if (parameterOverrides != null && props != null) { this.parameterMap.put( key, parameterOverrides.getOrDefault(key, props.getOrDefault(key, defaultValue))); @@ -116,19 +99,6 @@ private void checkAndUpdate( this.parameterMap.put(key, parameterOverrides.getOrDefault(key, defaultValue)); } else if (props != null) { this.parameterMap.put(key, props.getOrDefault(key, defaultValue)); - } else { - this.parameterMap.put(key, defaultValue); - } - - if (enforceDefault) { - if (!this.parameterMap.getOrDefault(key, defaultValue).equals(defaultValue)) { - throw new SFException( - INVALID_CONFIG_PARAMETER, - String.format( - "The value %s for %s is not configurable, should be %s.", - this.parameterMap.get(key), key, defaultValue)); - } - this.parameterMap.put(key, defaultValue); } } @@ -137,11 +107,8 @@ private void checkAndUpdate( * * @param parameterOverrides Map of parameter name -> value * @param props Properties file provided to client constructor - * @param isIcebergMode If the provided parameters need to be verified and modified to meet - * Iceberg mode */ - private void setParameterMap( - Map parameterOverrides, Properties props, boolean isIcebergMode) { + private void setParameterMap(Map parameterOverrides, Properties props) { // BUFFER_FLUSH_INTERVAL_IN_MILLIS is deprecated and disallowed if ((parameterOverrides != null && parameterOverrides.containsKey(BUFFER_FLUSH_INTERVAL_IN_MILLIS)) @@ -152,101 +119,75 @@ private void setParameterMap( BUFFER_FLUSH_INTERVAL_IN_MILLIS, MAX_CLIENT_LAG)); } - this.checkAndUpdate( + this.updateValue( BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS, BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS_DEFAULT, parameterOverrides, - props, - false); + props); - this.checkAndUpdate( + this.updateValue( INSERT_THROTTLE_INTERVAL_IN_MILLIS, INSERT_THROTTLE_INTERVAL_IN_MILLIS_DEFAULT, parameterOverrides, - props, - false); + props); - this.checkAndUpdate( + this.updateValue( INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE, INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE_DEFAULT, parameterOverrides, - props, - false); + props); - this.checkAndUpdate( + this.updateValue( INSERT_THROTTLE_THRESHOLD_IN_BYTES, INSERT_THROTTLE_THRESHOLD_IN_BYTES_DEFAULT, parameterOverrides, - props, - false); + props); - this.checkAndUpdate( + this.updateValue( ENABLE_SNOWPIPE_STREAMING_METRICS, SNOWPIPE_STREAMING_METRICS_DEFAULT, parameterOverrides, - props, - false); + props); - this.checkAndUpdate( - BLOB_FORMAT_VERSION, BLOB_FORMAT_VERSION_DEFAULT, parameterOverrides, props, false); + this.updateValue(BLOB_FORMAT_VERSION, BLOB_FORMAT_VERSION_DEFAULT, parameterOverrides, props); getBlobFormatVersion(); // to verify parsing the configured value - this.checkAndUpdate( - IO_TIME_CPU_RATIO, IO_TIME_CPU_RATIO_DEFAULT, parameterOverrides, props, false); + this.updateValue(IO_TIME_CPU_RATIO, IO_TIME_CPU_RATIO_DEFAULT, parameterOverrides, props); - this.checkAndUpdate( + this.updateValue( BLOB_UPLOAD_MAX_RETRY_COUNT, BLOB_UPLOAD_MAX_RETRY_COUNT_DEFAULT, parameterOverrides, - props, - false); + props); - this.checkAndUpdate( - MAX_MEMORY_LIMIT_IN_BYTES, - MAX_MEMORY_LIMIT_IN_BYTES_DEFAULT, - parameterOverrides, - props, - false); + this.updateValue( + MAX_MEMORY_LIMIT_IN_BYTES, MAX_MEMORY_LIMIT_IN_BYTES_DEFAULT, parameterOverrides, props); - this.checkAndUpdate( + this.updateValue( ENABLE_PARQUET_INTERNAL_BUFFERING, ENABLE_PARQUET_INTERNAL_BUFFERING_DEFAULT, parameterOverrides, - props, - false); + props); - this.checkAndUpdate( - MAX_CHANNEL_SIZE_IN_BYTES, - MAX_CHANNEL_SIZE_IN_BYTES_DEFAULT, - parameterOverrides, - props, - false); + this.updateValue( + MAX_CHANNEL_SIZE_IN_BYTES, MAX_CHANNEL_SIZE_IN_BYTES_DEFAULT, parameterOverrides, props); - this.checkAndUpdate( - MAX_CHUNK_SIZE_IN_BYTES, MAX_CHUNK_SIZE_IN_BYTES_DEFAULT, parameterOverrides, props, false); + this.updateValue( + MAX_CHUNK_SIZE_IN_BYTES, MAX_CHUNK_SIZE_IN_BYTES_DEFAULT, parameterOverrides, props); - this.checkAndUpdate( - MAX_CLIENT_LAG, - isIcebergMode ? MAX_CLIENT_LAG_ICEBERG_MODE_DEFAULT : MAX_CLIENT_LAG_DEFAULT, - parameterOverrides, - props, - false); + this.updateValue(MAX_CLIENT_LAG, MAX_CLIENT_LAG_DEFAULT, parameterOverrides, props); - this.checkAndUpdate( + this.updateValue( MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST, - isIcebergMode - ? MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT - : MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT, + MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT, parameterOverrides, - props, - isIcebergMode); + props); - this.checkAndUpdate( + this.updateValue( BDEC_PARQUET_COMPRESSION_ALGORITHM, BDEC_PARQUET_COMPRESSION_ALGORITHM_DEFAULT, parameterOverrides, - props, - false); + props); } /** @return Longest interval in milliseconds between buffer flushes */ diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/BlobBuilderTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/BlobBuilderTest.java index 32a352ba9..e220aec79 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/BlobBuilderTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/BlobBuilderTest.java @@ -1,7 +1,3 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - package net.snowflake.ingest.streaming.internal; import java.io.ByteArrayOutputStream; @@ -27,21 +23,18 @@ public void testSerializationErrors() throws Exception { BlobBuilder.constructBlobAndMetadata( "a.bdec", Collections.singletonList(createChannelDataPerTable(1, false)), - Constants.BdecVersion.THREE, - true); + Constants.BdecVersion.THREE); BlobBuilder.constructBlobAndMetadata( "a.bdec", Collections.singletonList(createChannelDataPerTable(1, true)), - Constants.BdecVersion.THREE, - true); + Constants.BdecVersion.THREE); // Construction fails if metadata contains 0 rows and data 1 row try { BlobBuilder.constructBlobAndMetadata( "a.bdec", Collections.singletonList(createChannelDataPerTable(0, false)), - Constants.BdecVersion.THREE, - true); + Constants.BdecVersion.THREE); Assert.fail("Should not pass enableParquetInternalBuffering=false"); } catch (SFException e) { Assert.assertEquals(ErrorCode.INTERNAL_ERROR.getMessageCode(), e.getVendorCode()); @@ -59,8 +52,7 @@ public void testSerializationErrors() throws Exception { BlobBuilder.constructBlobAndMetadata( "a.bdec", Collections.singletonList(createChannelDataPerTable(0, true)), - Constants.BdecVersion.THREE, - true); + Constants.BdecVersion.THREE); Assert.fail("Should not pass enableParquetInternalBuffering=true"); } catch (SFException e) { Assert.assertEquals(ErrorCode.INTERNAL_ERROR.getMessageCode(), e.getVendorCode()); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/ChannelCacheTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/ChannelCacheTest.java index 7c6d797b4..947908ef9 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/ChannelCacheTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/ChannelCacheTest.java @@ -24,7 +24,7 @@ public class ChannelCacheTest { @Before public void setup() { cache = new ChannelCache<>(); - client = new SnowflakeStreamingIngestClientInternal<>("client", false); + client = new SnowflakeStreamingIngestClientInternal<>("client"); channel1 = new SnowflakeStreamingIngestChannelInternal<>( "channel1", diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java index f59f9e076..afb2be40c 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java @@ -1,7 +1,3 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - package net.snowflake.ingest.streaming.internal; import static net.snowflake.ingest.utils.Constants.BLOB_CHECKSUM_SIZE_IN_BYTES; @@ -11,7 +7,6 @@ import static net.snowflake.ingest.utils.Constants.BLOB_NO_HEADER; import static net.snowflake.ingest.utils.Constants.BLOB_TAG_SIZE_IN_BYTES; import static net.snowflake.ingest.utils.Constants.BLOB_VERSION_SIZE_IN_BYTES; -import static net.snowflake.ingest.utils.ParameterProvider.MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT; import static net.snowflake.ingest.utils.ParameterProvider.MAX_CHUNK_SIZE_IN_BYTES_DEFAULT; import com.codahale.metrics.Histogram; @@ -53,22 +48,11 @@ import net.snowflake.ingest.utils.SFException; import org.junit.Assert; import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; import org.mockito.Mockito; -@RunWith(Parameterized.class) public class FlushServiceTest { - - @Parameterized.Parameters(name = "isIcebergMode: {0}") - public static Object[] isIcebergMode() { - return new Object[] {false, true}; - } - - @Parameterized.Parameter public static boolean isIcebergMode; - public FlushServiceTest() { this.testContextFactory = ParquetTestContext.createFactory(); } @@ -96,23 +80,16 @@ private abstract static class TestContext implements AutoCloseable { StorageManager storageManager; StreamingIngestStorage storage; ParameterProvider parameterProvider; - InternalParameterProvider internalParameterProvider; RegisterService registerService; final List> channelData = new ArrayList<>(); TestContext() { storage = Mockito.mock(StreamingIngestStorage.class); - parameterProvider = new ParameterProvider(isIcebergMode); - internalParameterProvider = new InternalParameterProvider(isIcebergMode); + parameterProvider = new ParameterProvider(); client = Mockito.mock(SnowflakeStreamingIngestClientInternal.class); Mockito.when(client.getParameterProvider()).thenReturn(parameterProvider); - Mockito.when(client.getInternalParameterProvider()).thenReturn(internalParameterProvider); - storageManager = - Mockito.spy( - isIcebergMode - ? new ExternalVolumeManager<>(true, "role", "client", null) - : new InternalStageManager<>(true, "role", "client", null)); + storageManager = Mockito.spy(new InternalStageManager<>(true, "role", "client", null)); Mockito.doReturn(storage).when(storageManager).getStorage(ArgumentMatchers.any()); Mockito.when(storageManager.getClientPrefix()).thenReturn("client_prefix"); channelCache = new ChannelCache<>(); @@ -421,38 +398,33 @@ public void testGetFilePath() { StorageManager storageManager = testContext.storageManager; Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("UTC")); String clientPrefix = "honk"; - if (isIcebergMode) { - // TODO: SNOW-1502887 Blob path generation for iceberg table - String outputString = storageManager.generateBlobPath(); - } else { - String outputString = - ((InternalStageManager) storageManager).getBlobPath(calendar, clientPrefix); - Path outputPath = Paths.get(outputString); - Assert.assertTrue(outputPath.getFileName().toString().contains(clientPrefix)); - Assert.assertTrue( - calendar.get(Calendar.MINUTE) - - Integer.parseInt(outputPath.getParent().getFileName().toString()) - <= 1); - Assert.assertEquals( - Integer.toString(calendar.get(Calendar.HOUR_OF_DAY)), - outputPath.getParent().getParent().getFileName().toString()); - Assert.assertEquals( - Integer.toString(calendar.get(Calendar.DAY_OF_MONTH)), - outputPath.getParent().getParent().getParent().getFileName().toString()); - Assert.assertEquals( - Integer.toString(calendar.get(Calendar.MONTH) + 1), - outputPath.getParent().getParent().getParent().getParent().getFileName().toString()); - Assert.assertEquals( - Integer.toString(calendar.get(Calendar.YEAR)), - outputPath - .getParent() - .getParent() - .getParent() - .getParent() - .getParent() - .getFileName() - .toString()); - } + String outputString = + ((InternalStageManager) storageManager).getBlobPath(calendar, clientPrefix); + Path outputPath = Paths.get(outputString); + Assert.assertTrue(outputPath.getFileName().toString().contains(clientPrefix)); + Assert.assertTrue( + calendar.get(Calendar.MINUTE) + - Integer.parseInt(outputPath.getParent().getFileName().toString()) + <= 1); + Assert.assertEquals( + Integer.toString(calendar.get(Calendar.HOUR_OF_DAY)), + outputPath.getParent().getParent().getFileName().toString()); + Assert.assertEquals( + Integer.toString(calendar.get(Calendar.DAY_OF_MONTH)), + outputPath.getParent().getParent().getParent().getFileName().toString()); + Assert.assertEquals( + Integer.toString(calendar.get(Calendar.MONTH) + 1), + outputPath.getParent().getParent().getParent().getParent().getFileName().toString()); + Assert.assertEquals( + Integer.toString(calendar.get(Calendar.YEAR)), + outputPath + .getParent() + .getParent() + .getParent() + .getParent() + .getParent() + .getFileName() + .toString()); } @Test @@ -624,9 +596,7 @@ public void runTestBlobSplitDueToNumberOfChunks(int numberOfRows) throws Excepti Math.ceil( (double) numberOfRows / channelsPerTable - / (isIcebergMode - ? MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT - : ParameterProvider.MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT)); + / ParameterProvider.MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_DEFAULT); final TestContext>> testContext = testContextFactory.create(); @@ -904,7 +874,7 @@ public void testInvalidateChannels() { // Create a new Client in order to not interfere with other tests SnowflakeStreamingIngestClientInternal client = Mockito.mock(SnowflakeStreamingIngestClientInternal.class); - ParameterProvider parameterProvider = new ParameterProvider(isIcebergMode); + ParameterProvider parameterProvider = new ParameterProvider(); ChannelCache channelCache = new ChannelCache<>(); Mockito.when(client.getChannelCache()).thenReturn(channelCache); Mockito.when(client.getParameterProvider()).thenReturn(parameterProvider); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/InsertRowsBenchmarkTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/InsertRowsBenchmarkTest.java index d28b22669..5b28e9c45 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/InsertRowsBenchmarkTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/InsertRowsBenchmarkTest.java @@ -1,7 +1,3 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - package net.snowflake.ingest.streaming.internal; import static java.time.ZoneOffset.UTC; @@ -40,7 +36,7 @@ public class InsertRowsBenchmarkTest { @Setup(Level.Trial) public void setUpBeforeAll() { - client = new SnowflakeStreamingIngestClientInternal("client_PARQUET", false); + client = new SnowflakeStreamingIngestClientInternal("client_PARQUET"); channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/InternalParameterProviderTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/InternalParameterProviderTest.java deleted file mode 100644 index e716b3ddb..000000000 --- a/src/test/java/net/snowflake/ingest/streaming/internal/InternalParameterProviderTest.java +++ /dev/null @@ -1,26 +0,0 @@ -/* - * Copyright (c) 2024. Snowflake Computing Inc. All rights reserved. - */ - -package net.snowflake.ingest.streaming.internal; - -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; - -@RunWith(Parameterized.class) -public class InternalParameterProviderTest { - @Parameterized.Parameters(name = "isIcebergMode: {0}") - public static Object[] isIcebergMode() { - return new Object[] {false, true}; - } - - @Parameterized.Parameter public static boolean isIcebergMode; - - @Test - public void testConstantParameterProvider() { - InternalParameterProvider internalParameterProvider = - new InternalParameterProvider(isIcebergMode); - assert internalParameterProvider.getEnableChunkEncryption() == !isIcebergMode; - } -} diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/OAuthBasicTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/OAuthBasicTest.java index 0c0c5bb85..6351f267d 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/OAuthBasicTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/OAuthBasicTest.java @@ -114,7 +114,7 @@ public void testCreateOAuthClient() throws Exception { @Test public void testSetRefreshToken() throws Exception { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("TEST_CLIENT", false); + new SnowflakeStreamingIngestClientInternal<>("TEST_CLIENT"); MockOAuthClient mockOAuthClient = new MockOAuthClient(); OAuthManager oAuthManager = diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java index c4631b348..86cece9c7 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/ParameterProviderTest.java @@ -1,16 +1,12 @@ package net.snowflake.ingest.streaming.internal; -import static net.snowflake.ingest.utils.ParameterProvider.MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT; - import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Properties; import net.snowflake.ingest.utils.Constants; -import net.snowflake.ingest.utils.ErrorCode; import net.snowflake.ingest.utils.ParameterProvider; -import net.snowflake.ingest.utils.SFException; import org.junit.Assert; import org.junit.Test; @@ -35,7 +31,7 @@ private Map getStartingParameterMap() { public void withValuesSet() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals(1000L, parameterProvider.getCachedMaxClientLagInMs()); Assert.assertEquals(4L, parameterProvider.getBufferFlushCheckIntervalInMs()); @@ -58,7 +54,7 @@ public void withNullProps() { parameterMap.put(ParameterProvider.BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS, 4L); parameterMap.put(ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE, 6); parameterMap.put(ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_BYTES, 1024); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, null, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, null); Assert.assertEquals(3000, parameterProvider.getCachedMaxClientLagInMs()); Assert.assertEquals(4, parameterProvider.getBufferFlushCheckIntervalInMs()); @@ -76,7 +72,7 @@ public void withNullParameterMap() { props.put(ParameterProvider.BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS, 4L); props.put(ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE, 6); props.put(ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_BYTES, 1024); - ParameterProvider parameterProvider = new ParameterProvider(null, props, false); + ParameterProvider parameterProvider = new ParameterProvider(null, props); Assert.assertEquals(3000, parameterProvider.getCachedMaxClientLagInMs()); Assert.assertEquals(4, parameterProvider.getBufferFlushCheckIntervalInMs()); @@ -89,7 +85,7 @@ public void withNullParameterMap() { @Test public void withNullInputs() { - ParameterProvider parameterProvider = new ParameterProvider(null, null, false); + ParameterProvider parameterProvider = new ParameterProvider(null, null); Assert.assertEquals( ParameterProvider.MAX_CLIENT_LAG_DEFAULT, parameterProvider.getCachedMaxClientLagInMs()); @@ -109,7 +105,7 @@ public void withNullInputs() { @Test public void withDefaultValues() { - ParameterProvider parameterProvider = new ParameterProvider(false); + ParameterProvider parameterProvider = new ParameterProvider(); Assert.assertEquals( ParameterProvider.MAX_CLIENT_LAG_DEFAULT, parameterProvider.getCachedMaxClientLagInMs()); @@ -141,50 +137,12 @@ public void withDefaultValues() { parameterProvider.getBdecParquetCompressionAlgorithm()); } - @Test - public void withIcebergDefaultValues() { - ParameterProvider parameterProvider = new ParameterProvider(true); - - Assert.assertEquals( - ParameterProvider.MAX_CLIENT_LAG_ICEBERG_MODE_DEFAULT, - parameterProvider.getCachedMaxClientLagInMs()); - Assert.assertEquals( - ParameterProvider.BUFFER_FLUSH_CHECK_INTERVAL_IN_MILLIS_DEFAULT, - parameterProvider.getBufferFlushCheckIntervalInMs()); - Assert.assertEquals( - ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_PERCENTAGE_DEFAULT, - parameterProvider.getInsertThrottleThresholdInPercentage()); - Assert.assertEquals( - ParameterProvider.INSERT_THROTTLE_THRESHOLD_IN_BYTES_DEFAULT, - parameterProvider.getInsertThrottleThresholdInBytes()); - Assert.assertEquals( - ParameterProvider.INSERT_THROTTLE_INTERVAL_IN_MILLIS_DEFAULT, - parameterProvider.getInsertThrottleIntervalInMs()); - Assert.assertEquals( - ParameterProvider.IO_TIME_CPU_RATIO_DEFAULT, parameterProvider.getIOTimeCpuRatio()); - Assert.assertEquals( - ParameterProvider.BLOB_UPLOAD_MAX_RETRY_COUNT_DEFAULT, - parameterProvider.getBlobUploadMaxRetryCount()); - Assert.assertEquals( - ParameterProvider.MAX_MEMORY_LIMIT_IN_BYTES_DEFAULT, - parameterProvider.getMaxMemoryLimitInBytes()); - Assert.assertEquals( - ParameterProvider.MAX_CHANNEL_SIZE_IN_BYTES_DEFAULT, - parameterProvider.getMaxChannelSizeInBytes()); - Assert.assertEquals( - ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM_DEFAULT, - parameterProvider.getBdecParquetCompressionAlgorithm()); - Assert.assertEquals( - MAX_CHUNKS_IN_BLOB_AND_REGISTRATION_REQUEST_ICEBERG_MODE_DEFAULT, - parameterProvider.getMaxChunksInBlobAndRegistrationRequest()); - } - @Test public void testMaxClientLagEnabled() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "2 second"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals(2000, parameterProvider.getCachedMaxClientLagInMs()); // call again to trigger caching logic Assert.assertEquals(2000, parameterProvider.getCachedMaxClientLagInMs()); @@ -195,7 +153,7 @@ public void testMaxClientLagEnabledPluralTimeUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "2 seconds"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals(2000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -204,7 +162,7 @@ public void testMaxClientLagEnabledMinuteTimeUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "1 minute"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals(60000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -213,7 +171,7 @@ public void testMaxClientLagEnabledMinuteTimeUnitPluralTimeUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "2 minutes"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals(120000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -221,7 +179,7 @@ public void testMaxClientLagEnabledMinuteTimeUnitPluralTimeUnit() { public void testMaxClientLagEnabledDefaultValue() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals( ParameterProvider.MAX_CLIENT_LAG_DEFAULT, parameterProvider.getCachedMaxClientLagInMs()); } @@ -231,7 +189,7 @@ public void testMaxClientLagEnabledDefaultUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "3000"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals(3000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -240,7 +198,7 @@ public void testMaxClientLagEnabledLongInput() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, 3000L); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals(3000, parameterProvider.getCachedMaxClientLagInMs()); } @@ -249,7 +207,7 @@ public void testMaxClientLagEnabledMissingUnitTimeUnitSupplied() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, " year"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -263,7 +221,7 @@ public void testMaxClientLagEnabledInvalidTimeUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "1 year"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -277,7 +235,7 @@ public void testMaxClientLagEnabledInvalidUnit() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "banana minute"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -291,7 +249,7 @@ public void testMaxClientLagEnabledThresholdBelow() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "0 second"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -305,7 +263,7 @@ public void testMaxClientLagEnabledThresholdAbove() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, "11 minutes"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -319,7 +277,7 @@ public void testMaxClientLagEnableEmptyInput() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.MAX_CLIENT_LAG, ""); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); try { parameterProvider.getCachedMaxClientLagInMs(); Assert.fail("Should not have succeeded"); @@ -333,20 +291,8 @@ public void testMaxChunksInBlobAndRegistrationRequest() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put("max_chunks_in_blob_and_registration_request", 1); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); - Assert.assertEquals(1, parameterProvider.getMaxChunksInBlobAndRegistrationRequest()); - - parameterProvider = new ParameterProvider(parameterMap, prop, true); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals(1, parameterProvider.getMaxChunksInBlobAndRegistrationRequest()); - - SFException e = - Assert.assertThrows( - SFException.class, - () -> { - parameterMap.put("max_chunks_in_blob_and_registration_request", 100); - new ParameterProvider(parameterMap, prop, true); - }); - Assert.assertEquals(e.getVendorCode(), ErrorCode.INVALID_CONFIG_PARAMETER.getMessageCode()); } @Test @@ -357,7 +303,7 @@ public void testValidCompressionAlgorithmsAndWithUppercaseLowerCase() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM, v); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals( Constants.BdecParquetCompression.GZIP, parameterProvider.getBdecParquetCompressionAlgorithm()); @@ -368,7 +314,7 @@ public void testValidCompressionAlgorithmsAndWithUppercaseLowerCase() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM, v); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); Assert.assertEquals( Constants.BdecParquetCompression.ZSTD, parameterProvider.getBdecParquetCompressionAlgorithm()); @@ -380,7 +326,7 @@ public void testInvalidCompressionAlgorithm() { Properties prop = new Properties(); Map parameterMap = getStartingParameterMap(); parameterMap.put(ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM, "invalid_comp"); - ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop, false); + ParameterProvider parameterProvider = new ParameterProvider(parameterMap, prop); try { parameterProvider.getBdecParquetCompressionAlgorithm(); Assert.fail("Should not have succeeded"); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java index 4eaea15a4..37eb5f96e 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java @@ -14,17 +14,8 @@ import org.junit.Assert; import org.junit.Ignore; import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; -@RunWith(Parameterized.class) public class RegisterServiceTest { - @Parameterized.Parameters(name = "isIcebergMode: {0}") - public static Object[] isIcebergMode() { - return new Object[] {false, true}; - } - - @Parameterized.Parameter public boolean isIcebergMode; @Test public void testRegisterService() throws ExecutionException, InterruptedException { @@ -54,7 +45,7 @@ public void testRegisterService() throws ExecutionException, InterruptedExceptio @Test public void testRegisterServiceTimeoutException() throws Exception { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); + new SnowflakeStreamingIngestClientInternal<>("client"); RegisterService rs = new RegisterService<>(client, true); Pair, CompletableFuture> blobFuture1 = @@ -82,7 +73,7 @@ public void testRegisterServiceTimeoutException() throws Exception { @Test public void testRegisterServiceTimeoutException_testRetries() throws Exception { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); + new SnowflakeStreamingIngestClientInternal<>("client"); RegisterService rs = new RegisterService<>(client, true); Pair, CompletableFuture> blobFuture1 = @@ -116,7 +107,7 @@ public void testRegisterServiceTimeoutException_testRetries() throws Exception { @Test public void testRegisterServiceNonTimeoutException() { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); + new SnowflakeStreamingIngestClientInternal<>("client"); RegisterService rs = new RegisterService<>(client, true); CompletableFuture future = new CompletableFuture<>(); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java index dc3285df8..5d8d8d36a 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestChannelTest.java @@ -48,11 +48,8 @@ import net.snowflake.ingest.utils.Utils; import org.junit.Assert; import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; import org.mockito.Mockito; -@RunWith(Parameterized.class) public class SnowflakeStreamingIngestChannelTest { /** @@ -74,13 +71,6 @@ public long getFreeMemory() { } } - @Parameterized.Parameters(name = "isIcebergMode: {0}") - public static Object[] isIcebergMode() { - return new Object[] {false, true}; - } - - @Parameterized.Parameter public boolean isIcebergMode; - @Test public void testChannelFactoryNullFields() { String name = "CHANNEL"; @@ -90,7 +80,7 @@ public void testChannelFactoryNullFields() { long channelSequencer = 0L; long rowSequencer = 0L; SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); + new SnowflakeStreamingIngestClientInternal<>("client"); Object[] fields = new Object[] { @@ -132,7 +122,7 @@ public void testChannelFactorySuccess() { long rowSequencer = 0L; SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); + new SnowflakeStreamingIngestClientInternal<>("client"); SnowflakeStreamingIngestChannelInternal channel = SnowflakeStreamingIngestChannelFactory.builder(name) @@ -167,7 +157,7 @@ public void testChannelFactorySuccess() { @Test public void testChannelValid() { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); + new SnowflakeStreamingIngestClientInternal<>("client"); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -217,7 +207,7 @@ public void testChannelValid() { @Test public void testChannelClose() { SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); + new SnowflakeStreamingIngestClientInternal<>("client"); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -361,7 +351,6 @@ public void testOpenChannelErrorResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -432,7 +421,6 @@ public void testOpenChannelSnowflakeInternalErrorResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -455,10 +443,6 @@ public void testOpenChannelSnowflakeInternalErrorResponse() throws Exception { @Test public void testOpenChannelSuccessResponse() throws Exception { - // TODO: SNOW-1490151 Iceberg testing gaps - if (isIcebergMode) { - return; - } String name = "CHANNEL"; String dbName = "STREAMINGINGEST_TEST"; String schemaName = "PUBLIC"; @@ -519,7 +503,6 @@ public void testOpenChannelSuccessResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -558,9 +541,7 @@ public void testOpenChannelSuccessResponse() throws Exception { @Test public void testInsertRow() { SnowflakeStreamingIngestClientInternal client; - client = - new SnowflakeStreamingIngestClientInternal( - "client_PARQUET", isIcebergMode); + client = new SnowflakeStreamingIngestClientInternal("client_PARQUET"); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -644,8 +625,7 @@ public void testInsertTooLargeRow() { schema.forEach(x -> row.put(x.getName(), byteArrayOneMb)); SnowflakeStreamingIngestClientInternal client; - client = - new SnowflakeStreamingIngestClientInternal("test_client", isIcebergMode); + client = new SnowflakeStreamingIngestClientInternal("test_client"); // Test channel with on error CONTINUE SnowflakeStreamingIngestChannelInternal channel = @@ -729,7 +709,7 @@ public void testInsertRowThrottling() { memoryInfoProvider.maxMemory = maxMemory; SnowflakeStreamingIngestClientInternal client = - new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode); + new SnowflakeStreamingIngestClientInternal<>("client"); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -745,7 +725,7 @@ public void testInsertRowThrottling() { OpenChannelRequest.OnErrorOption.CONTINUE, UTC); - ParameterProvider parameterProvider = new ParameterProvider(isIcebergMode); + ParameterProvider parameterProvider = new ParameterProvider(); memoryInfoProvider.freeMemory = maxMemory * (parameterProvider.getInsertThrottleThresholdInPercentage() - 1) / 100; @@ -775,7 +755,7 @@ public void testInsertRowThrottling() { @Test public void testFlush() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -811,7 +791,7 @@ public void testFlush() throws Exception { @Test public void testClose() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); SnowflakeStreamingIngestChannel channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -845,7 +825,7 @@ public void testClose() throws Exception { @Test public void testDropOnClose() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -882,7 +862,7 @@ public void testDropOnClose() throws Exception { @Test public void testDropOnCloseInvalidChannel() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", @@ -915,7 +895,7 @@ public void testDropOnCloseInvalidChannel() throws Exception { public void testGetLatestCommittedOffsetToken() { String offsetToken = "10"; SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); SnowflakeStreamingIngestChannel channel = new SnowflakeStreamingIngestChannelInternal<>( "channel", diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java index 5375adbec..553efbd31 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java @@ -1,7 +1,3 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - package net.snowflake.ingest.streaming.internal; import static java.time.ZoneOffset.UTC; @@ -70,12 +66,9 @@ import org.junit.Before; import org.junit.Ignore; import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; import org.mockito.ArgumentMatchers; import org.mockito.Mockito; -@RunWith(Parameterized.class) public class SnowflakeStreamingIngestClientTest { private static final ObjectMapper objectMapper = new ObjectMapper(); @@ -86,14 +79,6 @@ public class SnowflakeStreamingIngestClientTest { SnowflakeStreamingIngestChannelInternal channel3; SnowflakeStreamingIngestChannelInternal channel4; - // TODO: Add IcebergMode = True after Streaming to Iceberg is supported. - @Parameterized.Parameters(name = "isIcebergMode: {0}") - public static Object[] isIcebergMode() { - return new Object[] {false}; - } - - @Parameterized.Parameter public boolean isIcebergMode; - @Before public void setup() throws Exception { objectMapper.setVisibility(PropertyAccessor.GETTER, JsonAutoDetect.Visibility.ANY); @@ -113,7 +98,6 @@ public void setup() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -200,7 +184,6 @@ public void testConstructorParameters() throws Exception { SnowflakeStreamingIngestClientFactory.builder("client") .setProperties(prop) .setParameterOverrides(parameterMap) - .setIsIceberg(isIcebergMode) .setIsTestMode(true) .build(); @@ -228,7 +211,6 @@ public void testClientFactoryWithJmxMetrics() throws Exception { .setProperties(prop) .setParameterOverrides( Collections.singletonMap(ENABLE_SNOWPIPE_STREAMING_METRICS, true)) - .setIsIceberg(isIcebergMode) .build(); Assert.assertEquals("client", client.getName()); @@ -380,7 +362,6 @@ public void testGetChannelsStatusWithRequest() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -440,7 +421,6 @@ public void testDropChannel() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -486,7 +466,6 @@ public void testGetChannelsStatusWithRequestError() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -540,7 +519,6 @@ public void testRegisterBlobRequestCreationSuccess() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -732,7 +710,6 @@ public void testGetRetryBlobs() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -774,7 +751,6 @@ public void testRegisterBlobErrorResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -822,7 +798,6 @@ public void testRegisterBlobSnowflakeInternalErrorResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -879,7 +854,6 @@ public void testRegisterBlobSuccessResponse() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -964,7 +938,6 @@ public void testRegisterBlobsRetries() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -994,7 +967,6 @@ public void testRegisterBlobChunkLimit() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null)); @@ -1167,7 +1139,6 @@ public void testRegisterBlobsRetriesSucceeds() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -1242,7 +1213,6 @@ public void testRegisterBlobResponseWithInvalidChannel() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); @@ -1293,7 +1263,7 @@ public void testRegisterBlobResponseWithInvalidChannel() throws Exception { @Test public void testFlush() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); ChannelsStatusResponse response = new ChannelsStatusResponse(); response.setStatusCode(0L); response.setMessage("Success"); @@ -1315,7 +1285,7 @@ public void testFlush() throws Exception { @Test public void testClose() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); ChannelsStatusResponse response = new ChannelsStatusResponse(); response.setStatusCode(0L); response.setMessage("Success"); @@ -1349,7 +1319,7 @@ public void testClose() throws Exception { @Test public void testCloseWithError() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); CompletableFuture future = new CompletableFuture<>(); future.completeExceptionally(new Exception("Simulating Error")); @@ -1387,7 +1357,7 @@ public void testCloseWithError() throws Exception { @Test public void testVerifyChannelsAreFullyCommittedSuccess() throws Exception { SnowflakeStreamingIngestClientInternal client = - Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client", isIcebergMode)); + Mockito.spy(new SnowflakeStreamingIngestClientInternal<>("client")); SnowflakeStreamingIngestChannelInternal channel = new SnowflakeStreamingIngestChannelInternal<>( "channel1", @@ -1435,7 +1405,6 @@ public void testFlushServiceException() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, parameterMap); @@ -1472,7 +1441,6 @@ public void testGetLatestCommittedOffsetTokens() throws Exception { new SnowflakeURL("snowflake.dev.local:8082"), null, httpClient, - isIcebergMode, true, requestBuilder, null); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java index 478934f4b..f320c99aa 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java @@ -48,7 +48,6 @@ import net.snowflake.client.jdbc.internal.google.common.util.concurrent.ThreadFactoryBuilder; import net.snowflake.ingest.TestUtils; import net.snowflake.ingest.connection.RequestBuilder; -import net.snowflake.ingest.utils.ParameterProvider; import net.snowflake.ingest.utils.SFException; import org.junit.Assert; import org.junit.Test; @@ -286,7 +285,6 @@ public void testRefreshSnowflakeMetadataRemote() throws Exception { StorageManager storageManager = new InternalStageManager(true, "role", "client", snowflakeServiceClient); - ParameterProvider parameterProvider = new ParameterProvider(false); StreamingIngestStorage stage = new StreamingIngestStorage( storageManager, From 28dea75e7a8b66fe795fbf63bc2eb663e5b47321 Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Mon, 15 Jul 2024 12:48:44 -0700 Subject: [PATCH 6/9] fix --- .../ingest/connection/OAuthClient.java | 1 + .../ingest/streaming/OpenChannelRequest.java | 8 ++++ .../internal/ClientConfigureRequest.java | 31 ++++++++++++++- .../streaming/internal/ConfigureRequest.java | 38 ------------------- .../internal/OpenChannelResponse.java | 10 ----- .../internal/SnowflakeServiceClient.java | 2 +- ...nowflakeStreamingIngestClientInternal.java | 13 +------ .../internal/StreamingIngestStorage.java | 9 ++++- 8 files changed, 47 insertions(+), 65 deletions(-) delete mode 100644 src/main/java/net/snowflake/ingest/streaming/internal/ConfigureRequest.java diff --git a/src/main/java/net/snowflake/ingest/connection/OAuthClient.java b/src/main/java/net/snowflake/ingest/connection/OAuthClient.java index 61c736a42..a592899ea 100644 --- a/src/main/java/net/snowflake/ingest/connection/OAuthClient.java +++ b/src/main/java/net/snowflake/ingest/connection/OAuthClient.java @@ -93,6 +93,7 @@ public void refreshToken() throws IOException { /** Helper method for making refresh request */ private HttpUriRequest makeRefreshTokenRequest() { + // TODO SNOW-1538108 Use SnowflakeServiceClient to make the request HttpPost post = new HttpPost(oAuthCredential.get().getOAuthTokenEndpoint()); post.addHeader(HttpHeaders.CONTENT_TYPE, OAUTH_CONTENT_TYPE_HEADER); post.addHeader(HttpHeaders.AUTHORIZATION, oAuthCredential.get().getAuthHeader()); diff --git a/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java b/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java index 9a7272a16..4d3ea19aa 100644 --- a/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java +++ b/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java @@ -41,6 +41,7 @@ public enum OnErrorOption { private final ZoneId defaultTimezone; private final String offsetToken; + private final boolean isOffsetTokenProvided; private final OffsetTokenVerificationFunction offsetTokenVerificationFunction; @@ -58,6 +59,7 @@ public static class OpenChannelRequestBuilder { private ZoneId defaultTimezone; private String offsetToken; + private boolean isOffsetTokenProvided = false; private OffsetTokenVerificationFunction offsetTokenVerificationFunction; @@ -93,6 +95,7 @@ public OpenChannelRequestBuilder setDefaultTimezone(ZoneId defaultTimezone) { public OpenChannelRequestBuilder setOffsetToken(String offsetToken) { this.offsetToken = offsetToken; + this.isOffsetTokenProvided = true; return this; } @@ -122,6 +125,7 @@ private OpenChannelRequest(OpenChannelRequestBuilder builder) { this.onErrorOption = builder.onErrorOption; this.defaultTimezone = builder.defaultTimezone; this.offsetToken = builder.offsetToken; + this.isOffsetTokenProvided = builder.isOffsetTokenProvided; this.offsetTokenVerificationFunction = builder.offsetTokenVerificationFunction; } @@ -157,6 +161,10 @@ public String getOffsetToken() { return this.offsetToken; } + public boolean isOffsetTokenProvided() { + return this.isOffsetTokenProvided; + } + public OffsetTokenVerificationFunction getOffsetTokenVerificationFunction() { return this.offsetTokenVerificationFunction; } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java index ca6d72ab3..62c7a3578 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java @@ -4,15 +4,42 @@ package net.snowflake.ingest.streaming.internal; +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + /** Class used to serialize client configure request */ -class ClientConfigureRequest extends ConfigureRequest { +class ClientConfigureRequest implements StreamingIngestRequest { /** * Constructor for client configure request * * @param role Role to be used for the request. */ ClientConfigureRequest(String role) { - setRole(role); + this.role = role; + } + + @JsonProperty("role") + private String role; + + // File name for the GCS signed url request + @JsonInclude(JsonInclude.Include.NON_NULL) + @JsonProperty("file_name") + private String fileName; + + String getRole() { + return role; + } + + void setRole(String role) { + this.role = role; + } + + String getFileName() { + return fileName; + } + + void setFileName(String fileName) { + this.fileName = fileName; } @Override diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureRequest.java deleted file mode 100644 index 6ca10f52a..000000000 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ConfigureRequest.java +++ /dev/null @@ -1,38 +0,0 @@ -/* - * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved. - */ - -package net.snowflake.ingest.streaming.internal; - -import com.fasterxml.jackson.annotation.JsonInclude; -import com.fasterxml.jackson.annotation.JsonProperty; - -/** Abstract class for {@link ChannelConfigureRequest} and {@link ClientConfigureRequest} */ -abstract class ConfigureRequest implements StreamingIngestRequest { - @JsonProperty("role") - private String role; - - // File name for the GCS signed url request - @JsonInclude(JsonInclude.Include.NON_NULL) - @JsonProperty("file_name") - private String fileName; - - String getRole() { - return role; - } - - void setRole(String role) { - this.role = role; - } - - String getFileName() { - return fileName; - } - - void setFileName(String fileName) { - this.fileName = fileName; - } - - @Override - public abstract String getStringForLogging(); -} diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelResponse.java index a06545819..beb0eb0c4 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelResponse.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/OpenChannelResponse.java @@ -21,7 +21,6 @@ class OpenChannelResponse extends StreamingIngestResponse { private List tableColumns; private String encryptionKey; private Long encryptionKeyId; - private FileLocationInfo stageLocation; @JsonProperty("status_code") void setStatusCode(Long statusCode) { @@ -131,13 +130,4 @@ void setEncryptionKeyId(Long encryptionKeyId) { Long getEncryptionKeyId() { return this.encryptionKeyId; } - - @JsonProperty("stage_location") - void setStageLocation(FileLocationInfo stageLocation) { - this.stageLocation = stageLocation; - } - - FileLocationInfo getStageLocation() { - return this.stageLocation; - } } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java index 904e4f93a..d753f4781 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java @@ -132,7 +132,7 @@ DropChannelResponse dropChannel(DropChannelRequestInternal request) * @param request the channel status request * @return the response from the channel status request */ - ChannelsStatusResponse channelStatus(ChannelsStatusRequest request) + ChannelsStatusResponse getChannelStatus(ChannelsStatusRequest request) throws IngestResponseException, IOException { ChannelsStatusResponse response = executeApiRequestWithRetries( diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java index aee73b8c4..762e84659 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java @@ -28,7 +28,6 @@ import com.codahale.metrics.jmx.JmxReporter; import com.codahale.metrics.jvm.MemoryUsageGaugeSet; import com.codahale.metrics.jvm.ThreadStatesGaugeSet; -import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.annotations.VisibleForTesting; import java.io.IOException; import java.net.URI; @@ -85,9 +84,6 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea private static final Logging logger = new Logging(SnowflakeStreamingIngestClientInternal.class); - // Object mapper for all marshalling and unmarshalling - private static final ObjectMapper objectMapper = new ObjectMapper(); - // Counter to generate unique request ids per client private final AtomicLong counter = new AtomicLong(0); @@ -368,13 +364,6 @@ public SnowflakeStreamingIngestChannelInternal openChannel(OpenChannelRequest // Add channel to the channel cache this.channelCache.addChannel(channel); - // Add storage to the storage manager, only for external volume - this.storageManager.addStorage( - response.getDBName(), - response.getSchemaName(), - response.getTableName(), - response.getStageLocation()); - return channel; } catch (IOException | IngestResponseException e) { throw new SFException(e, ErrorCode.OPEN_CHANNEL_FAILURE, e.getMessage()); @@ -460,7 +449,7 @@ ChannelsStatusResponse getChannelsStatus( request.setChannels(requestDTOs); request.setRole(this.role); - ChannelsStatusResponse response = snowflakeServiceClient.channelStatus(request); + ChannelsStatusResponse response = snowflakeServiceClient.getChannelStatus(request); for (int idx = 0; idx < channels.size(); idx++) { SnowflakeStreamingIngestChannelInternal channel = channels.get(idx); diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java index 97c8e4a05..72c0f8368 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java @@ -36,8 +36,13 @@ class StreamingIngestStorage { private static final ObjectMapper mapper = new ObjectMapper(); - // Object mapper for parsing the client/configure response to Jackson version the same as - // jdbc.internal.fasterxml.jackson + /** + * Object mapper for parsing the client/configure response to Jackson version the same as + * jdbc.internal.fasterxml.jackson. We need two different versions of ObjectMapper because {@link + * SnowflakeFileTransferAgent#getFileTransferMetadatas(net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.JsonNode)} + * expects a different version of json object than {@link StreamingIngestResponse}. TODO: + * SNOW-1493470 Align Jackson version + */ private static final net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.ObjectMapper parseConfigureResponseMapper = new net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.ObjectMapper(); From c7fbc8c5920ae7364569aae0ec04081e0f93460b Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Mon, 15 Jul 2024 13:25:40 -0700 Subject: [PATCH 7/9] update comments style --- .../streaming/internal/FileLocationInfo.java | 33 ++++++++++++------- .../internal/InternalStageManager.java | 12 +++---- .../internal/SnowflakeServiceClient.java | 4 +-- .../streaming/internal/StorageManager.java | 2 +- 4 files changed, 31 insertions(+), 20 deletions(-) diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FileLocationInfo.java b/src/main/java/net/snowflake/ingest/streaming/internal/FileLocationInfo.java index 4af5862c4..add98a6fb 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/FileLocationInfo.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/FileLocationInfo.java @@ -9,38 +9,49 @@ /** Class used to deserialized volume information response by server */ class FileLocationInfo { + /** The stage type */ @JsonProperty("locationType") - private String locationType; // The stage type + private String locationType; + /** The container or bucket */ @JsonProperty("location") - private String location; // The container or bucket + private String location; + /** The path of the target file */ @JsonProperty("path") - private String path; // path of the target file + private String path; + /** The credentials required for the stage */ @JsonProperty("creds") - private Map credentials; // the credentials required for the stage + private Map credentials; + /** AWS/S3/GCS region (S3/GCS only) */ @JsonProperty("region") - private String region; // AWS/S3/GCS region (S3/GCS only) + private String region; + /** The Azure Storage endpoint (Azure only) */ @JsonProperty("endPoint") - private String endPoint; // The Azure Storage endpoint (Azure only) + private String endPoint; + /** The Azure Storage account (Azure only) */ @JsonProperty("storageAccount") - private String storageAccount; // The Azure Storage account (Azure only) + private String storageAccount; + /** GCS gives us back a presigned URL instead of a cred */ @JsonProperty("presignedUrl") - private String presignedUrl; // GCS gives us back a presigned URL instead of a cred + private String presignedUrl; + /** Whether to encrypt/decrypt files on the stage */ @JsonProperty("isClientSideEncrypted") - private boolean isClientSideEncrypted; // whether to encrypt/decrypt files on the stage + private boolean isClientSideEncrypted; + /** Whether to use s3 regional URL (AWS Only) */ @JsonProperty("useS3RegionalUrl") - private boolean useS3RegionalUrl; // whether to use s3 regional URL (AWS Only) + private boolean useS3RegionalUrl; + /** A unique id for volume assigned by server */ @JsonProperty("volumeHash") - private String volumeHash; // a unique id for volume assigned by server + private String volumeHash; String getLocationType() { return locationType; diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java index ea207b141..1fb406d10 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java @@ -25,22 +25,22 @@ public InternalStageLocation() {} /** Class to manage single Snowflake internal stage */ class InternalStageManager implements StorageManager { - // Target stage for the client + /** Target stage for the client */ private final StreamingIngestStorage targetStage; - // Increasing counter to generate a unique blob name per client + /** Increasing counter to generate a unique blob name per client */ private final AtomicLong counter; - // Whether the manager in test mode + /** Whether the manager in test mode */ private final boolean isTestMode; - // Snowflake service client used for configure calls + /** Snowflake service client used for configure calls */ private final SnowflakeServiceClient snowflakeServiceClient; - // The role of the client + /** The role of the client */ private final String role; - // Client prefix generated by the Snowflake server + /** Client prefix generated by the Snowflake server */ private final String clientPrefix; /** diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java index d753f4781..58b6fc829 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeServiceClient.java @@ -32,10 +32,10 @@ class SnowflakeServiceClient { private static final Logging logger = new Logging(SnowflakeServiceClient.class); - // HTTP client used for making requests + /** HTTP client used for making requests */ private final CloseableHttpClient httpClient; - // Request builder for building streaming API request + /** Request builder for building streaming API request */ private final RequestBuilder requestBuilder; /** diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java index ea0b8843f..798d33ade 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java @@ -13,7 +13,7 @@ * @param the type of location that's being managed (internal stage / external volume) */ interface StorageManager { - // Default max upload retries for streaming ingest storage + /** Default max upload retries for streaming ingest storage */ int DEFAULT_MAX_UPLOAD_RETRIES = 5; /** From 796a5331196b9b6ba22ea1a07afeb665ae7ace0e Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Tue, 16 Jul 2024 13:50:00 -0700 Subject: [PATCH 8/9] format --- .../ingest/streaming/internal/StreamingIngestStorage.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java index 1bf6190d5..82fee8a66 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java @@ -22,8 +22,6 @@ import java.util.Optional; import java.util.Properties; import java.util.concurrent.TimeUnit; -import java.util.function.Function; -import javax.annotation.Nullable; import net.snowflake.client.core.OCSPMode; import net.snowflake.client.jdbc.SnowflakeFileTransferAgent; import net.snowflake.client.jdbc.SnowflakeFileTransferConfig; From ca1d767367a883866bf012209e41b46ab2e229da Mon Sep 17 00:00:00 2001 From: Alec Huang Date: Tue, 16 Jul 2024 15:37:03 -0700 Subject: [PATCH 9/9] rename interface / undo error code change --- .../internal/ChannelsStatusRequest.java | 2 +- .../internal/ClientConfigureRequest.java | 2 +- .../internal/DropChannelRequestInternal.java | 2 +- .../ingest/streaming/internal/FlushService.java | 4 ++-- ...{StorageManager.java => IStorageManager.java} | 4 ++-- ...Request.java => IStreamingIngestRequest.java} | 2 +- .../streaming/internal/InternalStageManager.java | 4 ++-- .../internal/OpenChannelRequestInternal.java | 2 +- .../streaming/internal/RegisterBlobRequest.java | 2 +- .../internal/SnowflakeServiceClient.java | 2 +- .../SnowflakeStreamingIngestClientInternal.java | 2 +- .../internal/StreamingIngestStorage.java | 6 +++--- .../streaming/internal/StreamingIngestUtils.java | 2 +- .../net/snowflake/ingest/utils/ErrorCode.java | 2 +- .../streaming/internal/FlushServiceTest.java | 6 +++--- .../internal/StreamingIngestStorageTest.java | 16 ++++++++-------- 16 files changed, 30 insertions(+), 30 deletions(-) rename src/main/java/net/snowflake/ingest/streaming/internal/{StorageManager.java => IStorageManager.java} (94%) rename src/main/java/net/snowflake/ingest/streaming/internal/{StreamingIngestRequest.java => IStreamingIngestRequest.java} (90%) diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java index 5a17b6e91..025647f14 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java @@ -10,7 +10,7 @@ import net.snowflake.ingest.utils.Utils; /** Class to deserialize a request from a channel status request */ -class ChannelsStatusRequest implements StreamingIngestRequest { +class ChannelsStatusRequest implements IStreamingIngestRequest { // Used to deserialize a channel request static class ChannelStatusRequestDTO { diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java index 62c7a3578..79b282079 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ClientConfigureRequest.java @@ -8,7 +8,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; /** Class used to serialize client configure request */ -class ClientConfigureRequest implements StreamingIngestRequest { +class ClientConfigureRequest implements IStreamingIngestRequest { /** * Constructor for client configure request * diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java index a6e1e1746..322b53acf 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelRequestInternal.java @@ -10,7 +10,7 @@ import net.snowflake.ingest.utils.Utils; /** Class used to serialize the {@link DropChannelRequest} */ -class DropChannelRequestInternal implements StreamingIngestRequest { +class DropChannelRequestInternal implements IStreamingIngestRequest { @JsonProperty("request_id") private String requestId; diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java index b34335194..954abfc4a 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java @@ -95,7 +95,7 @@ List>> getData() { private final ChannelCache channelCache; // Reference to the Streaming Ingest storage manager - private final StorageManager storageManager; + private final IStorageManager storageManager; // Reference to register service private final RegisterService registerService; @@ -127,7 +127,7 @@ List>> getData() { FlushService( SnowflakeStreamingIngestClientInternal client, ChannelCache cache, - StorageManager storageManager, + IStorageManager storageManager, boolean isTestMode) { this.owningClient = client; this.channelCache = cache; diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/IStorageManager.java similarity index 94% rename from src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java rename to src/main/java/net/snowflake/ingest/streaming/internal/IStorageManager.java index 798d33ade..51f4a82de 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StorageManager.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/IStorageManager.java @@ -12,7 +12,7 @@ * @param The type of chunk data * @param the type of location that's being managed (internal stage / external volume) */ -interface StorageManager { +interface IStorageManager { /** Default max upload retries for streaming ingest storage */ int DEFAULT_MAX_UPLOAD_RETRIES = 5; @@ -54,7 +54,7 @@ void addStorage( /** * Decrement the blob sequencer, this method is needed to prevent gap between file name sequencer. - * See {@link StorageManager#generateBlobPath()} for more details. + * See {@link IStorageManager#generateBlobPath()} for more details. */ void decrementBlobSequencer(); diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/IStreamingIngestRequest.java similarity index 90% rename from src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestRequest.java rename to src/main/java/net/snowflake/ingest/streaming/internal/IStreamingIngestRequest.java index 3dd30b5e2..a4b5e29d1 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestRequest.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/IStreamingIngestRequest.java @@ -8,6 +8,6 @@ * The StreamingIngestRequest interface is a marker interface used for type safety in the {@link * SnowflakeServiceClient} for streaming ingest API request. */ -interface StreamingIngestRequest { +interface IStreamingIngestRequest { String getStringForLogging(); } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java index 144bdc662..d33a80738 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/InternalStageManager.java @@ -24,7 +24,7 @@ public InternalStageLocation() {} } /** Class to manage single Snowflake internal stage */ -class InternalStageManager implements StorageManager { +class InternalStageManager implements IStorageManager { /** Target stage for the client */ private final StreamingIngestStorage targetStage; @@ -94,7 +94,7 @@ class InternalStageManager implements StorageManager T executeApiRequestWithRetries( Class responseClass, - StreamingIngestRequest request, + IStreamingIngestRequest request, String endpoint, String operation, ServiceResponseHandler.ApiName apiName) diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java index 762e84659..6331a4045 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java @@ -106,7 +106,7 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea private final FlushService flushService; // Reference to storage manager - private final StorageManager storageManager; + private final IStorageManager storageManager; // Indicates whether the client has closed private volatile boolean isClosed; diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java index 82fee8a66..242b5cc43 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorage.java @@ -89,7 +89,7 @@ state to record unknown age. } private SnowflakeFileTransferMetadataWithAge fileTransferMetadataWithAge; - private final StorageManager owningManager; + private final IStorageManager owningManager; private final TLocation location; private final String clientName; @@ -108,7 +108,7 @@ state to record unknown age. * @param maxUploadRetries The maximum number of retries to attempt */ StreamingIngestStorage( - StorageManager owningManager, + IStorageManager owningManager, String clientName, FileLocationInfo fileLocationInfo, TLocation location, @@ -133,7 +133,7 @@ state to record unknown age. * @param maxUploadRetries the maximum number of retries to attempt */ StreamingIngestStorage( - StorageManager owningManager, + IStorageManager owningManager, String clientName, SnowflakeFileTransferMetadataWithAge testMetadata, TLocation location, diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtils.java b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtils.java index ebfffa234..538283b4e 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtils.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/StreamingIngestUtils.java @@ -80,7 +80,7 @@ public static void sleepForRetry(int executionCount) { static T executeWithRetries( Class targetClass, String endpoint, - StreamingIngestRequest payload, + IStreamingIngestRequest payload, String message, ServiceResponseHandler.ApiName apiName, CloseableHttpClient httpClient, diff --git a/src/main/java/net/snowflake/ingest/utils/ErrorCode.java b/src/main/java/net/snowflake/ingest/utils/ErrorCode.java index bd5c36c7f..a9aab9c3b 100644 --- a/src/main/java/net/snowflake/ingest/utils/ErrorCode.java +++ b/src/main/java/net/snowflake/ingest/utils/ErrorCode.java @@ -26,7 +26,7 @@ public enum ErrorCode { INVALID_ENCRYPTED_KEY("0018"), INVALID_DATA_IN_CHUNK("0019"), IO_ERROR("0020"), - UNABLE_TO_CONNECT_TO_STORAGE("0021"), + UNABLE_TO_CONNECT_TO_STAGE("0021"), KEYPAIR_CREATION_FAILURE("0022"), MD5_HASHING_NOT_AVAILABLE("0023"), CHANNEL_STATUS_FAILURE("0024"), diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java index afb2be40c..8ac1d2b85 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java @@ -77,7 +77,7 @@ private abstract static class TestContext implements AutoCloseable { ChannelCache channelCache; final Map> channels = new HashMap<>(); FlushService flushService; - StorageManager storageManager; + IStorageManager storageManager; StreamingIngestStorage storage; ParameterProvider parameterProvider; RegisterService registerService; @@ -395,7 +395,7 @@ private static ColumnMetadata createLargeTestTextColumn(String name) { @Test public void testGetFilePath() { TestContext testContext = testContextFactory.create(); - StorageManager storageManager = testContext.storageManager; + IStorageManager storageManager = testContext.storageManager; Calendar calendar = Calendar.getInstance(TimeZone.getTimeZone("UTC")); String clientPrefix = "honk"; String outputString = @@ -925,7 +925,7 @@ public void testInvalidateChannels() { innerData.add(channel1Data); innerData.add(channel2Data); - StorageManager storageManager = + IStorageManager storageManager = Mockito.spy(new InternalStageManager<>(true, "role", "client", null)); FlushService flushService = new FlushService<>(client, channelCache, storageManager, false); diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java index 7b1efb2ef..d4c3f0374 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/StreamingIngestStorageTest.java @@ -130,7 +130,7 @@ public void testPutRemote() throws Exception { byte[] dataBytes = "Hello Upload".getBytes(StandardCharsets.UTF_8); - StorageManager storageManager = Mockito.mock(StorageManager.class); + IStorageManager storageManager = Mockito.mock(IStorageManager.class); Mockito.when(storageManager.getClientPrefix()).thenReturn("testPrefix"); StreamingIngestStorage stage = @@ -201,7 +201,7 @@ public void doTestPutRemoteRefreshes() throws Exception { byte[] dataBytes = "Hello Upload".getBytes(StandardCharsets.UTF_8); - StorageManager storageManager = Mockito.mock(StorageManager.class); + IStorageManager storageManager = Mockito.mock(IStorageManager.class); Mockito.when(storageManager.getClientPrefix()).thenReturn("testPrefix"); StreamingIngestStorage stage = @@ -256,7 +256,7 @@ public void testPutRemoteGCS() throws Exception { byte[] dataBytes = "Hello Upload".getBytes(StandardCharsets.UTF_8); - StorageManager storageManager = Mockito.mock(StorageManager.class); + IStorageManager storageManager = Mockito.mock(IStorageManager.class); Mockito.when(storageManager.getClientPrefix()).thenReturn("testPrefix"); StreamingIngestStorage stage = @@ -294,7 +294,7 @@ public void testRefreshSnowflakeMetadataRemote() throws Exception { SnowflakeServiceClient snowflakeServiceClient = new SnowflakeServiceClient(mockClient, mockBuilder); - StorageManager storageManager = + IStorageManager storageManager = new InternalStageManager<>(true, "role", "client", snowflakeServiceClient); StreamingIngestStorage stage = @@ -353,7 +353,7 @@ public void testRefreshSnowflakeMetadataDeploymentIdMismatch() throws Exception SnowflakeServiceClient snowflakeServiceClient = new SnowflakeServiceClient(mockClient, mockBuilder); - StorageManager storageManager = + IStorageManager storageManager = new InternalStageManager<>(true, "role", "clientName", snowflakeServiceClient); StreamingIngestStorage storage = storageManager.getStorage(""); @@ -385,7 +385,7 @@ public void testFetchSignedURL() throws Exception { Mockito.when(mockClientInternal.getRole()).thenReturn("role"); SnowflakeServiceClient snowflakeServiceClient = new SnowflakeServiceClient(mockClient, mockBuilder); - StorageManager storageManager = + IStorageManager storageManager = new InternalStageManager<>(true, "role", "client", snowflakeServiceClient); StatusLine mockStatusLine = Mockito.mock(StatusLine.class); Mockito.when(mockStatusLine.getStatusCode()).thenReturn(200); @@ -431,7 +431,7 @@ public void testRefreshSnowflakeMetadataSynchronized() throws Exception { Mockito.when(mockClientInternal.getRole()).thenReturn("role"); SnowflakeServiceClient snowflakeServiceClient = new SnowflakeServiceClient(mockClient, mockBuilder); - StorageManager storageManager = + IStorageManager storageManager = new InternalStageManager<>(true, "role", "client", snowflakeServiceClient); StatusLine mockStatusLine = Mockito.mock(StatusLine.class); Mockito.when(mockStatusLine.getStatusCode()).thenReturn(200); @@ -574,7 +574,7 @@ public void testRefreshMetadataOnFirstPutException() throws Exception { byte[] dataBytes = "Hello Upload".getBytes(StandardCharsets.UTF_8); - StorageManager storageManager = Mockito.mock(StorageManager.class); + IStorageManager storageManager = Mockito.mock(IStorageManager.class); Mockito.when(storageManager.getClientPrefix()).thenReturn("testPrefix"); StreamingIngestStorage stage =