files, UUID requestId, boolean showSkippedFiles)
throws URISyntaxException, IOException, IngestResponseException, BackOffException {
- return ingestFiles(files, requestId, showSkippedFiles, null /* Client info is null */);
- }
-
- /**
- * ingestFiles With Client Info - synchronously sends a request to the ingest service to enqueue
- * these files along with clientSequencer and offSetToken.
- *
- * OffsetToken will be atomically persisted on server(Snowflake) side along with files if the
- * clientSequencer added in this request matches with what Snowflake currently has.
- *
- *
If clientSequencers doesnt match, 400 response code is sent back and no files will be added.
- *
- * @param files - list of wrappers around filenames and sizes
- * @param requestId - a requestId that we'll use to label - if null, we generate one for the user
- * @param showSkippedFiles - a flag which returns the files that were skipped when set to true.
- * @param clientInfo - clientSequencer and offsetToken to pass along with files. Can be null.
- * @return an insert response from the server
- * @throws URISyntaxException - if the provided account name was illegal and caused a URI
- * construction failure
- * @throws IOException - if we have some other network failure
- * @throws IngestResponseException - if snowflake encountered error during ingest
- * @throws BackOffException - if we have a 503 response
- */
- public IngestResponse ingestFiles(
- List files,
- UUID requestId,
- boolean showSkippedFiles,
- InsertFilesClientInfo clientInfo)
- throws URISyntaxException, IOException, IngestResponseException, BackOffException {
-
// the request id we want to send with this payload
if (requestId == null || requestId.toString().isEmpty()) {
requestId = UUID.randomUUID();
}
HttpPost httpPostForIngestFile =
- builder.generateInsertRequest(
- requestId, pipe, files, showSkippedFiles, Optional.ofNullable(clientInfo));
+ builder.generateInsertRequest(requestId, pipe, files, showSkippedFiles);
// send the request and get a response....
- HttpResponse response = httpClient.execute(httpPostForIngestFile);
-
- LOGGER.info(
- "Attempting to unmarshall insert response - {}, with clientInfo - {}",
- response,
- clientInfo);
- return ServiceResponseHandler.unmarshallIngestResponse(response, requestId);
+ try (CloseableHttpResponse response = httpClient.execute(httpPostForIngestFile)) {
+ LOGGER.info("Attempting to unmarshall insert response - {}", response);
+ return ServiceResponseHandler.unmarshallIngestResponse(
+ response, requestId, httpClient, httpPostForIngestFile, builder);
+ }
}
/**
@@ -597,10 +612,11 @@ public HistoryResponse getHistory(UUID requestId, Integer recentSeconds, String
builder.generateHistoryRequest(requestId, pipe, recentSeconds, beginMark);
// send the request and get a response...
- HttpResponse response = httpClient.execute(httpGetHistory);
-
- LOGGER.info("Attempting to unmarshall history response - {}", response);
- return ServiceResponseHandler.unmarshallHistoryResponse(response, requestId);
+ try (CloseableHttpResponse response = httpClient.execute(httpGetHistory)) {
+ LOGGER.info("Attempting to unmarshall history response - {}", response);
+ return ServiceResponseHandler.unmarshallHistoryResponse(
+ response, requestId, httpClient, httpGetHistory, builder);
+ }
}
/**
@@ -626,57 +642,14 @@ public HistoryRangeResponse getHistoryRange(
requestId = UUID.randomUUID();
}
- HttpResponse response =
- httpClient.execute(
- builder.generateHistoryRangeRequest(
- requestId, pipe, startTimeInclusive, endTimeExclusive));
-
- LOGGER.info("Attempting to unmarshall history range response - {}", response);
- return ServiceResponseHandler.unmarshallHistoryRangeResponse(response, requestId);
- }
+ HttpGet request =
+ builder.generateHistoryRangeRequest(requestId, pipe, startTimeInclusive, endTimeExclusive);
- /**
- * Register a snowpipe client and returns the client sequencer
- *
- * @param requestId a UUID we use to label the request, if null, one is generated for the user
- * @return
- * @throws URISyntaxException - if the provided account name was illegal and caused a URI
- * construction failure
- * @throws IOException - if we have some other network failure
- * @throws IngestResponseException - if snowflake encountered error during ingest
- * @throws BackOffException - if we have a 503 response
- */
- public ConfigureClientResponse configureClient(UUID requestId)
- throws URISyntaxException, IOException, IngestResponseException, BackOffException {
- if (requestId == null || requestId.toString().isEmpty()) {
- requestId = UUID.randomUUID();
- }
- HttpResponse response =
- httpClient.execute(builder.generateConfigureClientRequest(requestId, pipe));
- LOGGER.info("Attempting to unmarshall configure client response - {}", response);
- return ServiceResponseHandler.unmarshallConfigureClientResponse(response, requestId);
- }
-
- /**
- * Get client status for snowpipe which contains offset token and client sequencer
- *
- * @param requestId a UUID we use to label the request, if null, one is generated for the user
- * @return
- * @throws URISyntaxException - if the provided account name was illegal and caused a URI
- * construction failure
- * @throws IOException - if we have some other network failure
- * @throws IngestResponseException - if snowflake encountered error during ingest
- * @throws BackOffException - if we have a 503 response
- */
- public ClientStatusResponse getClientStatus(UUID requestId)
- throws URISyntaxException, IOException, IngestResponseException, BackOffException {
- if (requestId == null || requestId.toString().isEmpty()) {
- requestId = UUID.randomUUID();
+ try (CloseableHttpResponse response = httpClient.execute(request)) {
+ LOGGER.info("Attempting to unmarshall history range response - {}", response);
+ return ServiceResponseHandler.unmarshallHistoryRangeResponse(
+ response, requestId, httpClient, request, builder);
}
- HttpResponse response =
- httpClient.execute(builder.generateGetClientStatusRequest(requestId, pipe));
- LOGGER.info("Attempting to unmarshall get client status response - {}", response);
- return ServiceResponseHandler.unmarshallGetClientStatus(response, requestId);
}
/**
@@ -687,6 +660,7 @@ public ClientStatusResponse getClientStatus(UUID requestId)
@Override
public void close() {
builder.closeResources();
+ HttpUtil.shutdownHttpConnectionManagerDaemonThread();
}
/* Used for testing */
diff --git a/src/main/java/net/snowflake/ingest/connection/ClientStatusResponse.java b/src/main/java/net/snowflake/ingest/connection/ClientStatusResponse.java
deleted file mode 100644
index 5a3cbdc24..000000000
--- a/src/main/java/net/snowflake/ingest/connection/ClientStatusResponse.java
+++ /dev/null
@@ -1,44 +0,0 @@
-/*
- * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved.
- */
-package net.snowflake.ingest.connection;
-
-import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
-
-/** ClientStatusResponse - response from a get client status request */
-@JsonIgnoreProperties(ignoreUnknown = true)
-public class ClientStatusResponse {
- private Long clientSequencer;
-
- private String offsetToken;
-
- @Override
- public String toString() {
- return "IngestResponse{"
- + "clientSequencer='"
- + clientSequencer
- + '\''
- + ", offsetToken='"
- + offsetToken
- + '\''
- + '}';
- }
-
- /**
- * unique identifier for the client
- *
- * @return clientSequencer
- */
- public Long getClientSequencer() {
- return clientSequencer;
- }
-
- /**
- * offset number for kafka connector
- *
- * @return offsetToken
- */
- public String getOffsetToken() {
- return offsetToken;
- }
-}
diff --git a/src/main/java/net/snowflake/ingest/connection/ConfigureClientResponse.java b/src/main/java/net/snowflake/ingest/connection/ConfigureClientResponse.java
deleted file mode 100644
index 144065c22..000000000
--- a/src/main/java/net/snowflake/ingest/connection/ConfigureClientResponse.java
+++ /dev/null
@@ -1,26 +0,0 @@
-/*
- * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved.
- */
-package net.snowflake.ingest.connection;
-
-import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
-
-/** ConfigureClientResponse - response from a configure client request */
-@JsonIgnoreProperties(ignoreUnknown = true)
-public class ConfigureClientResponse {
- private Long clientSequencer;
-
- @Override
- public String toString() {
- return "IngestResponse{" + "clientSequencer='" + clientSequencer + '\'' + '}';
- }
-
- /**
- * unique identifier for the client
- *
- * @return clientSequencer
- */
- public Long getClientSequencer() {
- return clientSequencer;
- }
-}
diff --git a/src/main/java/net/snowflake/ingest/connection/InsertFilesClientInfo.java b/src/main/java/net/snowflake/ingest/connection/InsertFilesClientInfo.java
deleted file mode 100644
index 175f994d0..000000000
--- a/src/main/java/net/snowflake/ingest/connection/InsertFilesClientInfo.java
+++ /dev/null
@@ -1,75 +0,0 @@
-/*
- * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved.
- */
-package net.snowflake.ingest.connection;
-
-import java.nio.charset.StandardCharsets;
-
-/**
- * Just a wrapper class which is serialised into REST request for insertFiles.
- *
- * This is an optional field which can be passed with a required field "files" in the request
- * body.
- *
- *
Here is how the new request could look like
- *
- *
- * {
- * "files":[
- * {
- * "path":"file1.csv.gz"
- * },
- * {
- * "path":"file2.csv.gz"
- * }
- * ],
- * "clientInfo": {
- * "clientSequencer": 1,
- * "offsetToken": "2"
- * }
- * }
- *
- */
-public class InsertFilesClientInfo {
- // FDB constant
- public static final int MAX_ALLOWED_OFFSET_TOKEN_BYTE_SIZE = 100000;
-
- // client sequencer which the caller thinks it currently has
- private final long clientSequencer;
-
- // offsetToken to atomically commit with files.
- private final String offsetToken;
-
- /** Constructor with both fields as required. */
- public InsertFilesClientInfo(long clientSequencer, String offsetToken) {
- if (clientSequencer < 0) {
- throw new IllegalArgumentException("ClientSequencer should be non negative.");
- }
- if (offsetToken == null || offsetToken.isEmpty()) {
- throw new IllegalArgumentException("OffsetToken can not be null or empty.");
- }
- /* Keeping some buffer */
- if (offsetToken.getBytes(StandardCharsets.UTF_8).length
- > MAX_ALLOWED_OFFSET_TOKEN_BYTE_SIZE - 1000) {
- throw new IllegalArgumentException("OffsetToken size too large.");
- }
- this.clientSequencer = clientSequencer;
- this.offsetToken = offsetToken;
- }
-
- /** Gets client Sequencer associated with this clientInfo record. */
- public long getClientSequencer() {
- return clientSequencer;
- }
-
- /** Gets offsetToken associated with this clientInfo record. */
- public String getOffsetToken() {
- return offsetToken;
- }
-
- /** Only printing clientSequencer */
- @Override
- public String toString() {
- return "InsertFilesClientInfo{" + "clientSequencer=" + clientSequencer + '}';
- }
-}
diff --git a/src/main/java/net/snowflake/ingest/connection/JWTManager.java b/src/main/java/net/snowflake/ingest/connection/JWTManager.java
new file mode 100644
index 000000000..691348a34
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/connection/JWTManager.java
@@ -0,0 +1,180 @@
+/*
+ * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.connection;
+
+import com.nimbusds.jose.JOSEException;
+import com.nimbusds.jose.JWSAlgorithm;
+import com.nimbusds.jose.JWSHeader;
+import com.nimbusds.jose.JWSSigner;
+import com.nimbusds.jose.crypto.RSASSASigner;
+import com.nimbusds.jwt.JWTClaimsSet;
+import com.nimbusds.jwt.SignedJWT;
+import java.security.KeyPair;
+import java.util.Date;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import net.snowflake.ingest.utils.Cryptor;
+
+/**
+ * This class manages creating and automatically renewing the JWT token
+ *
+ * @author obabarinsa
+ */
+public final class JWTManager extends SecurityManager {
+ // the token lifetime is 59 minutes
+ private static final float LIFETIME_IN_MINUTES = 59;
+
+ // the renewal time is 54 minutes
+ private static final int RENEWAL_INTERVAL_IN_MINUTES = 54;
+
+ private static final String TOKEN_TYPE = "KEYPAIR_JWT";
+
+ // The public - private key pair we're using to connect to the service
+ private final transient KeyPair keyPair;
+
+ // the token itself
+ private final AtomicReference token;
+
+ /**
+ * Creates a JWTManager entity for a given account, user and KeyPair with a specified time to
+ * renew the token
+ *
+ * @param accountName - the snowflake account name of this user
+ * @param username - the snowflake username of the current user
+ * @param keyPair - the public/private key pair we're using to connect
+ * @param timeTillRenewal - the time measure until we renew the token
+ * @param unit the unit by which timeTillRenewal is measured
+ * @param telemetryService reference to the telemetry service
+ */
+ JWTManager(
+ String accountName,
+ String username,
+ KeyPair keyPair,
+ int timeTillRenewal,
+ TimeUnit unit,
+ TelemetryService telemetryService) {
+ super(accountName, username, telemetryService);
+ // if any of our arguments are null, throw an exception
+ if (keyPair == null) {
+ throw new IllegalArgumentException();
+ }
+ token = new AtomicReference<>();
+
+ // we have to keep around the keys
+ this.keyPair = keyPair;
+
+ // generate our first token
+ refreshToken();
+
+ // schedule all future renewals
+ tokenRefresher.scheduleAtFixedRate(this::refreshToken, timeTillRenewal, timeTillRenewal, unit);
+ }
+
+ /**
+ * Creates a JWTManager entity for a given account, user and KeyPair with the default time to
+ * renew (RENEWAL_INTERVAL_IN_MINUTES minutes)
+ *
+ * @param accountName - the snowflake account name of this user
+ * @param username - the snowflake username of the current user
+ * @param keyPair - the public/private key pair we're using to connect
+ * @param telemetryService reference to the telemetry service
+ */
+ public JWTManager(
+ String accountName, String username, KeyPair keyPair, TelemetryService telemetryService) {
+ this(
+ accountName,
+ username,
+ keyPair,
+ RENEWAL_INTERVAL_IN_MINUTES,
+ TimeUnit.MINUTES,
+ telemetryService);
+ }
+
+ @Override
+ public String getToken() {
+ // if we failed to regenerate the token at some point, throw
+ if (refreshFailed.get()) {
+ LOGGER.error("getToken request failed due to token regeneration failure");
+ throw new SecurityException();
+ }
+
+ return token.get();
+ }
+
+ @Override
+ String getTokenType() {
+ return TOKEN_TYPE;
+ }
+
+ /** regenerateToken - Regenerates our Token given our current user, account and keypair */
+ @Override
+ void refreshToken() {
+ // create our JWT claim builder object
+ JWTClaimsSet.Builder builder = new JWTClaimsSet.Builder();
+
+ // get the subject to the fully qualified username
+ String subject = String.format("%s.%s", account, user);
+
+ // get the issuer
+ String publicKeyFPInJwt = calculatePublicKeyFp(keyPair);
+ String issuer = String.format("%s.%s.%s", account, user, publicKeyFPInJwt);
+
+ // iat set to now
+ Date iat = new Date(System.currentTimeMillis());
+
+ // expiration in 59 minutes
+ Date exp = new Date(iat.getTime() + 59 * 60 * 1000);
+
+ // build claim set
+ JWTClaimsSet claimsSet =
+ builder.issuer(issuer).subject(subject).issueTime(iat).expirationTime(exp).build();
+ LOGGER.debug("Creating new JWT with subject {} and issuer {}...", subject, issuer);
+
+ SignedJWT signedJWT = new SignedJWT(new JWSHeader(JWSAlgorithm.RS256), claimsSet);
+
+ JWSSigner signer = new RSASSASigner(this.keyPair.getPrivate());
+
+ String newToken;
+ try {
+ signedJWT.sign(signer);
+ newToken = signedJWT.serialize();
+ } catch (JOSEException e) {
+ refreshFailed.set(true);
+ LOGGER.error("Failed to regenerate token! Exception is as follows : {}", e.getMessage());
+ throw new SecurityException();
+ }
+
+ // atomically update the string
+ LOGGER.info("Successfully created new JWT");
+ token.set(newToken);
+
+ // Refresh the token used in the telemetry service as well
+ if (telemetryService != null) {
+ telemetryService.refreshToken(newToken);
+ }
+ }
+
+ /**
+ * Given a keypair
+ *
+ * @return the fingerprint of public key
+ * The idea is to hash public key's raw bytes using SHA-256 and converts hash into a string
+ * using Base64 encoding.
+ */
+ private String calculatePublicKeyFp(KeyPair keyPair) {
+ // get the raw bytes of public key
+ byte[] publicKeyRawBytes = keyPair.getPublic().getEncoded();
+
+ // take sha256 on raw bytes and do base64 encode
+ publicKeyFingerPrint = String.format("SHA256:%s", Cryptor.sha256HashBase64(publicKeyRawBytes));
+ return publicKeyFingerPrint;
+ }
+
+ /** Currently, it only shuts down the instance of ExecutorService. */
+ @Override
+ public void close() {
+ tokenRefresher.shutdown();
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/connection/OAuthClient.java b/src/main/java/net/snowflake/ingest/connection/OAuthClient.java
new file mode 100644
index 000000000..bdff6c96e
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/connection/OAuthClient.java
@@ -0,0 +1,14 @@
+/*
+ * Copyright (c) 2023 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.connection;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+/** Interface to perform token refresh request from {@link OAuthManager} */
+public interface OAuthClient {
+ AtomicReference getoAuthCredentialRef();
+
+ void refreshToken();
+}
diff --git a/src/main/java/net/snowflake/ingest/connection/OAuthCredential.java b/src/main/java/net/snowflake/ingest/connection/OAuthCredential.java
new file mode 100644
index 000000000..e222478d1
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/connection/OAuthCredential.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2023 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.connection;
+
+import java.util.Base64;
+
+/** This class hold credentials for OAuth authentication */
+public class OAuthCredential {
+ private static final String BASIC_AUTH_HEADER_PREFIX = "Basic ";
+ private final String authHeader;
+ private final String clientId;
+ private final String clientSecret;
+ private String accessToken;
+ private String refreshToken;
+ private int expiresIn;
+
+ public OAuthCredential(String clientId, String clientSecret, String refreshToken) {
+ this.authHeader =
+ BASIC_AUTH_HEADER_PREFIX
+ + Base64.getEncoder().encodeToString((clientId + ":" + clientSecret).getBytes());
+ this.clientId = clientId;
+ this.clientSecret = clientSecret;
+ this.refreshToken = refreshToken;
+ }
+
+ public String getAuthHeader() {
+ return authHeader;
+ }
+
+ public String getAccessToken() {
+ return accessToken;
+ }
+
+ public void setAccessToken(String accessToken) {
+ this.accessToken = accessToken;
+ }
+
+ public String getRefreshToken() {
+ return refreshToken;
+ }
+
+ public void setRefreshToken(String refreshToken) {
+ this.refreshToken = refreshToken;
+ }
+
+ public void setExpiresIn(int expiresIn) {
+ this.expiresIn = expiresIn;
+ }
+
+ public int getExpiresIn() {
+ return expiresIn;
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/connection/OAuthManager.java b/src/main/java/net/snowflake/ingest/connection/OAuthManager.java
new file mode 100644
index 000000000..fefb94299
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/connection/OAuthManager.java
@@ -0,0 +1,176 @@
+/*
+ * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.connection;
+
+import java.util.concurrent.TimeUnit;
+import net.snowflake.client.jdbc.internal.apache.http.client.utils.URIBuilder;
+import net.snowflake.ingest.utils.Constants;
+import net.snowflake.ingest.utils.ErrorCode;
+import net.snowflake.ingest.utils.SFException;
+
+/** This class manages creating and automatically refresh the OAuth token */
+public final class OAuthManager extends SecurityManager {
+ private static final double DEFAULT_UPDATE_THRESHOLD_RATIO = 0.8;
+
+ // the endpoint for token request
+ private static final String TOKEN_REQUEST_ENDPOINT = "/oauth/token-request";
+
+ // Content type header to specify the encoding
+ private static final String OAUTH_CONTENT_TYPE_HEADER = "application/x-www-form-urlencoded";
+
+ // Properties for token refresh request
+ private static final String TOKEN_TYPE = "OAUTH";
+
+ // Update threshold, a floating-point value representing the ratio between the expiration time of
+ // an access token and the time needed to update it. It must be a value greater than 0 and less
+ // than 1. E.g. An access token with expires_in=600 and update_threshold_ratio=0.8 would be
+ // updated after 600*0.8 = 480.
+ private final double updateThresholdRatio;
+
+ private OAuthClient oAuthClient;
+
+ /**
+ * Creates a OAuthManager entity for a given account, user and OAuthCredential with default time
+ * to refresh the access token
+ *
+ * @param accountName - the snowflake account name of this user
+ * @param username - the snowflake username of the current user
+ * @param oAuthCredential - the OAuth credential we're using to connect
+ * @param baseURIBuilder - the uri builder with common scheme, host and port
+ * @param telemetryService reference to the telemetry service
+ */
+ OAuthManager(
+ String accountName,
+ String username,
+ OAuthCredential oAuthCredential,
+ URIBuilder baseURIBuilder,
+ TelemetryService telemetryService) {
+ this(
+ accountName,
+ username,
+ oAuthCredential,
+ baseURIBuilder,
+ DEFAULT_UPDATE_THRESHOLD_RATIO,
+ telemetryService);
+ }
+
+ /**
+ * Creates a OAuthManager entity for a given account, user and OAuthCredential with a specified
+ * time to refresh the token
+ *
+ * @param accountName - the snowflake account name of this user
+ * @param username - the snowflake username of the current user
+ * @param oAuthCredential - the OAuth credential we're using to connect
+ * @param baseURIBuilder - the uri builder with common scheme, host and port
+ * @param updateThresholdRatio - the ratio between the expiration time of a token and the time
+ * needed to refresh it.
+ * @param telemetryService reference to the telemetry service
+ */
+ OAuthManager(
+ String accountName,
+ String username,
+ OAuthCredential oAuthCredential,
+ URIBuilder baseURIBuilder,
+ double updateThresholdRatio,
+ TelemetryService telemetryService) {
+ // disable telemetry service until jdbc v3.13.34 is released
+ super(accountName, username, null);
+
+ // if any of our arguments are null, throw an exception
+ if (oAuthCredential == null || baseURIBuilder == null) {
+ throw new IllegalArgumentException();
+ }
+
+ // check if update threshold is within (0, 1)
+ if (updateThresholdRatio <= 0 || updateThresholdRatio >= 1) {
+ throw new IllegalArgumentException("updateThresholdRatio should fall in (0, 1)");
+ }
+ this.updateThresholdRatio = updateThresholdRatio;
+ this.oAuthClient = new SnowflakeOAuthClient(accountName, oAuthCredential, baseURIBuilder);
+
+ // generate our first token
+ refreshToken();
+ }
+
+ /**
+ * Creates a OAuthManager entity for a given account, user and OAuthClient with a specified time
+ * to refresh the token. Use for testing only.
+ *
+ * @param accountName - the snowflake account name of this user
+ * @param username - the snowflake username of the current user
+ * @param oAuthClient - the OAuth client to perform token refresh
+ * @param updateThresholdRatio - the ratio between the expiration time of a token and the time
+ * needed to refresh it.
+ */
+ public OAuthManager(
+ String accountName, String username, OAuthClient oAuthClient, double updateThresholdRatio) {
+ super(accountName, username, null);
+
+ this.updateThresholdRatio = updateThresholdRatio;
+ this.oAuthClient = oAuthClient;
+
+ refreshToken();
+ }
+
+ @Override
+ String getToken() {
+ if (refreshFailed.get()) {
+ throw new SecurityException("getToken request failed due to token refresh failure");
+ }
+ return oAuthClient.getoAuthCredentialRef().get().getAccessToken();
+ }
+
+ @Override
+ String getTokenType() {
+ return TOKEN_TYPE;
+ }
+
+ /**
+ * 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.
+ *
+ * @param refreshToken the new refresh token
+ */
+ void setRefreshToken(String refreshToken) {
+ oAuthClient.getoAuthCredentialRef().get().setRefreshToken(refreshToken);
+ }
+
+ /** refreshToken - Get new access token using refresh_token, client_id, client_secret */
+ @Override
+ void refreshToken() {
+ for (int retries = 0; retries < Constants.MAX_OAUTH_REFRESH_TOKEN_RETRY; retries++) {
+ try {
+ oAuthClient.refreshToken();
+
+ // Schedule next refresh
+ long nextRefreshDelay =
+ (long)
+ (oAuthClient.getoAuthCredentialRef().get().getExpiresIn()
+ * this.updateThresholdRatio);
+ tokenRefresher.schedule(this::refreshToken, nextRefreshDelay, TimeUnit.SECONDS);
+ LOGGER.debug(
+ "Refresh access token, next refresh is scheduled after {} seconds", nextRefreshDelay);
+
+ return;
+ } catch (SFException e1) {
+ // Exponential backoff retries
+ try {
+ Thread.sleep((1L << retries) * 1000L);
+ } catch (InterruptedException e2) {
+ throw new SFException(ErrorCode.OAUTH_REFRESH_TOKEN_ERROR, e2.getMessage());
+ }
+ }
+ }
+
+ refreshFailed.set(true);
+ throw new SecurityException("Fail to refresh access token");
+ }
+
+ /** Currently, it only shuts down the instance of ExecutorService. */
+ @Override
+ public void close() {
+ tokenRefresher.shutdown();
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java b/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java
index 44d846b94..3fdb8a660 100644
--- a/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java
+++ b/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java
@@ -1,41 +1,32 @@
/*
- * Copyright (c) 2012-2017 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
*/
package net.snowflake.ingest.connection;
+import static net.snowflake.ingest.utils.Constants.ENABLE_TELEMETRY_TO_SF;
import static net.snowflake.ingest.utils.Utils.isNullOrEmpty;
-import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
-import java.io.FileNotFoundException;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.UncheckedIOException;
import java.net.URI;
import java.net.URISyntaxException;
-import java.net.URL;
-import java.nio.file.Files;
-import java.nio.file.Paths;
import java.security.KeyPair;
import java.util.List;
import java.util.Map;
-import java.util.Optional;
-import java.util.Properties;
import java.util.UUID;
-import net.snowflake.ingest.SimpleIngestManager;
+import net.snowflake.client.jdbc.internal.apache.http.HttpHeaders;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpGet;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpPost;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpUriRequest;
+import net.snowflake.client.jdbc.internal.apache.http.client.utils.URIBuilder;
+import net.snowflake.client.jdbc.internal.apache.http.entity.ContentType;
+import net.snowflake.client.jdbc.internal.apache.http.entity.StringEntity;
+import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient;
import net.snowflake.ingest.utils.ErrorCode;
import net.snowflake.ingest.utils.SFException;
import net.snowflake.ingest.utils.SnowflakeURL;
import net.snowflake.ingest.utils.StagedFileWrapper;
-import org.apache.http.HttpHeaders;
-import org.apache.http.client.methods.HttpGet;
-import org.apache.http.client.methods.HttpPost;
-import org.apache.http.client.methods.HttpUriRequest;
-import org.apache.http.client.utils.URIBuilder;
-import org.apache.http.entity.ContentType;
-import org.apache.http.entity.StringEntity;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -52,7 +43,7 @@ public class RequestBuilder {
/* Member variables Begin */
// the security manager which will handle token generation
- private SecurityManager securityManager;
+ private final SecurityManager securityManager;
// whatever the actual scheme is
private String scheme;
@@ -65,6 +56,9 @@ public class RequestBuilder {
private final String userAgentSuffix;
+ // Reference to the telemetry service
+ private final TelemetryService telemetryService;
+
/* Member variables End */
/* Static constants Begin */
@@ -87,13 +81,6 @@ public class RequestBuilder {
// the endpoint for history time range queries
private static final String HISTORY_RANGE_ENDPOINT_FORMAT = "/v1/data/pipes/%s/loadHistoryScan";
- // the endpoint for configure snowpipe client
- private static final String CONFIGURE_CLIENT_ENDPOINT_FORMAT =
- "/v1/data/pipes/%s/client/configure";
-
- // the endpoint for get snowpipe client status
- private static final String CLIENT_STATUS_ENDPOINT_FORMAT = "/v1/data/pipes/%s/client/status";
-
// optional number of max seconds of items to fetch(eg. in the last hour)
private static final String RECENT_HISTORY_IN_SECONDS = "recentSeconds";
@@ -118,16 +105,12 @@ public class RequestBuilder {
// and object mapper for all marshalling and unmarshalling
private static final ObjectMapper objectMapper = new ObjectMapper();
- private static final String RESOURCES_FILE = "project.properties";
-
- private static final Properties PROPERTIES = loadProperties();
-
private static final String USER_AGENT = getDefaultUserAgent();
// Don't change!
public static final String CLIENT_NAME = "SnowpipeJavaSDK";
- public static final String DEFAULT_VERSION = "1.0.1-beta";
+ public static final String DEFAULT_VERSION = "2.0.0";
public static final String JAVA_USER_AGENT = "JAVA";
@@ -145,10 +128,11 @@ public class RequestBuilder {
*
* @param accountName - the name of the Snowflake account to which we're connecting
* @param userName - the username of the entity loading files
- * @param keyPair - the Public/Private key pair we'll use to authenticate
+ * @param credential - the credential we'll use to authenticate
*/
- public RequestBuilder(String accountName, String userName, KeyPair keyPair) {
- this(accountName, userName, keyPair, DEFAULT_SCHEME, DEFAULT_HOST_SUFFIX, DEFAULT_PORT, null);
+ public RequestBuilder(String accountName, String userName, Object credential) {
+ this(
+ accountName, userName, credential, DEFAULT_SCHEME, DEFAULT_HOST_SUFFIX, DEFAULT_PORT, null);
}
/**
@@ -156,15 +140,17 @@ public RequestBuilder(String accountName, String userName, KeyPair keyPair) {
*
* @param accountName - the name of the Snowflake account to which we're connecting
* @param userName - the username of the entity loading files
- * @param keyPair - the Public/Private key pair we'll use to authenticate
+ * @param credential - the credential we'll use to authenticate
+ * @param userAgentSuffix - The suffix part of HTTP Header User-Agent
*/
public RequestBuilder(
String accountName,
String userName,
String hostName,
- KeyPair keyPair,
+ Object credential,
String userAgentSuffix) {
- this(accountName, userName, keyPair, DEFAULT_SCHEME, hostName, DEFAULT_PORT, userAgentSuffix);
+ this(
+ accountName, userName, credential, DEFAULT_SCHEME, hostName, DEFAULT_PORT, userAgentSuffix);
}
/**
@@ -173,7 +159,7 @@ public RequestBuilder(
*
* @param accountName - the account name to which we're connecting
* @param userName - for whom are we connecting?
- * @param keyPair - our auth credentials
+ * @param credential - the credential we'll use to authenticate
* @param schemeName - are we HTTP or HTTPS?
* @param hostName - the host for this snowflake instance
* @param portNum - the port number
@@ -181,11 +167,11 @@ public RequestBuilder(
public RequestBuilder(
String accountName,
String userName,
- KeyPair keyPair,
+ Object credential,
String schemeName,
String hostName,
int portNum) {
- this(accountName, userName, keyPair, schemeName, hostName, portNum, null);
+ this(accountName, userName, credential, schemeName, hostName, portNum, null);
}
/**
@@ -193,7 +179,7 @@ public RequestBuilder(
*
* @param accountName - the account name to which we're connecting
* @param userName - for whom are we connecting?
- * @param keyPair - our auth credentials
+ * @param credential - the credential we'll use to authenticate
* @param schemeName - are we HTTP or HTTPS?
* @param hostName - the host for this snowflake instance
* @param portNum - the port number
@@ -202,20 +188,91 @@ public RequestBuilder(
public RequestBuilder(
String accountName,
String userName,
- KeyPair keyPair,
+ Object credential,
String schemeName,
String hostName,
int portNum,
String userAgentSuffix) {
+ this(
+ accountName,
+ userName,
+ credential,
+ schemeName,
+ hostName,
+ portNum,
+ userAgentSuffix,
+ null,
+ null,
+ null);
+ }
+
+ /**
+ * RequestBuilder - constructor used by streaming ingest
+ *
+ * @param url - the Snowflake account to which we're connecting
+ * @param userName - the username of the entity loading files
+ * @param credential - the credential we'll use to authenticate
+ * @param httpClient - reference to the http client
+ * @param clientName - name of the client, used to uniquely identify a client if used
+ */
+ public RequestBuilder(
+ SnowflakeURL url,
+ String userName,
+ Object credential,
+ CloseableHttpClient httpClient,
+ String clientName) {
+ this(
+ url.getAccount(),
+ userName,
+ credential,
+ url.getScheme(),
+ url.getUrlWithoutPort(),
+ url.getPort(),
+ null,
+ null,
+ httpClient,
+ clientName);
+ }
+
+ /**
+ * RequestBuilder - this constructor is for testing purposes only
+ *
+ * @param accountName - the account name to which we're connecting
+ * @param userName - for whom are we connecting?
+ * @param credential - our auth credentials, either JWT key pair or OAuth credential
+ * @param schemeName - are we HTTP or HTTPS?
+ * @param hostName - the host for this snowflake instance
+ * @param portNum - the port number
+ * @param userAgentSuffix - The suffix part of HTTP Header User-Agent
+ * @param securityManager - The security manager for authentication
+ * @param httpClient - reference to the http client
+ * @param clientName - name of the client, used to uniquely identify a client if used
+ */
+ public RequestBuilder(
+ String accountName,
+ String userName,
+ Object credential,
+ String schemeName,
+ String hostName,
+ int portNum,
+ String userAgentSuffix,
+ SecurityManager securityManager,
+ CloseableHttpClient httpClient,
+ String clientName) {
// none of these arguments should be null
- if (accountName == null || userName == null || keyPair == null) {
+ if (accountName == null || userName == null || credential == null) {
throw new IllegalArgumentException();
}
- // create our security/token manager
- securityManager = new SecurityManager(accountName, userName, keyPair);
+ // Set up the telemetry service if needed
+ // TODO: SNOW-854272 Support telemetry service when using OAuth authentication
+ this.telemetryService =
+ ENABLE_TELEMETRY_TO_SF && credential instanceof KeyPair
+ ? new TelemetryService(
+ httpClient, clientName, schemeName + "://" + hostName + ":" + portNum)
+ : null;
- // stash references to the account and user name as well
+ // stash references to the account and username as well
String account = accountName.toUpperCase();
String user = userName.toUpperCase();
@@ -225,6 +282,27 @@ public RequestBuilder(
this.host = hostName;
this.userAgentSuffix = userAgentSuffix;
+ // create our security/token manager
+ if (securityManager == null) {
+ if (credential instanceof KeyPair) {
+ this.securityManager =
+ new JWTManager(accountName, userName, (KeyPair) credential, telemetryService);
+ } else if (credential instanceof OAuthCredential) {
+ this.securityManager =
+ new OAuthManager(
+ accountName,
+ userName,
+ (OAuthCredential) credential,
+ makeBaseURI(),
+ telemetryService);
+ } else {
+ throw new IllegalArgumentException(
+ "Credential should be instance of either KeyPair or OAuthCredential");
+ }
+ } else {
+ this.securityManager = securityManager;
+ }
+
LOGGER.info(
"Creating a RequestBuilder with arguments : "
+ "Account : {}, User : {}, Scheme : {}, Host : {}, Port : {}, userAgentSuffix: {}",
@@ -236,52 +314,6 @@ public RequestBuilder(
this.userAgentSuffix);
}
- /**
- * RequestBuilder - constructor used by streaming ingest
- *
- * @param url - the Snowflake account to which we're connecting
- * @param userName - the username of the entity loading files
- * @param keyPair - the Public/Private key pair we'll use to authenticate
- */
- public RequestBuilder(SnowflakeURL url, String userName, KeyPair keyPair) {
- this(
- url.getAccount(),
- userName,
- keyPair,
- url.getScheme(),
- url.getUrlWithoutPort(),
- url.getPort());
- }
-
- private static Properties loadProperties() {
- Properties properties = new Properties();
- properties.put("version", DEFAULT_VERSION);
-
- try {
- URL res = SimpleIngestManager.class.getClassLoader().getResource(RESOURCES_FILE);
- if (res == null) {
- throw new UncheckedIOException(new FileNotFoundException(RESOURCES_FILE));
- }
-
- URI uri;
- try {
- uri = res.toURI();
- } catch (URISyntaxException ex) {
- throw new IllegalArgumentException(ex);
- }
-
- try (InputStream is = Files.newInputStream(Paths.get(uri))) {
- properties.load(is);
- } catch (IOException ex) {
- throw new UncheckedIOException("Failed to load resource", ex);
- }
- } catch (Exception e) {
- LOGGER.warn("Could not read version info: " + e.toString());
- }
-
- return properties;
- }
-
/**
* Creates a string for user agent which should always be present in all requests to Snowflake
* (Snowpipe APIs)
@@ -293,8 +325,7 @@ private static Properties loadProperties() {
* @return the default agent string
*/
private static String getDefaultUserAgent() {
- final String clientVersion = PROPERTIES.getProperty("version");
- StringBuilder defaultUserAgent = new StringBuilder(CLIENT_NAME + "/" + clientVersion);
+ StringBuilder defaultUserAgent = new StringBuilder(CLIENT_NAME + "/" + DEFAULT_VERSION);
final String osInformation =
String.format(
@@ -310,48 +341,35 @@ private static String getDefaultUserAgent() {
// Add Java Version
final String javaVersion = System.getProperty("java.version");
- defaultUserAgent.append(JAVA_USER_AGENT + "/" + javaVersion);
+ defaultUserAgent.append(JAVA_USER_AGENT + "/").append(javaVersion);
+ String userAgent = defaultUserAgent.toString();
- return defaultUserAgent.toString();
+ LOGGER.info("Default user agent " + userAgent);
+ return userAgent;
}
private static String buildCustomUserAgent(String additionalUserAgentInfo) {
return USER_AGENT.trim() + " " + additionalUserAgentInfo;
}
+
/** A simple POJO for generating our POST body to the insert endpoint */
private static class IngestRequest {
// the list of files we're loading
private final List files;
- // additional info passed along with files which includes clientSequencer and offsetToken
- // can be null
- @JsonInclude(JsonInclude.Include.NON_NULL)
- private final InsertFilesClientInfo clientInfo;
-
- /** Constructor used when both files and clientInfo are passed in request */
- public IngestRequest(List files, InsertFilesClientInfo clientInfo) {
- this.files = files;
- this.clientInfo = clientInfo;
- }
-
/**
* Ctor used when only files is used in request body.
*
* clientInfo will be defaulted to null
*/
public IngestRequest(List files) {
- this(files, null);
+ this.files = files;
}
/* Gets the list of files which were added in request */
public List getFiles() {
return files;
}
-
- /* Gets the clientInfo associated with files which was added in request */
- public InsertFilesClientInfo getClientInfo() {
- return clientInfo;
- }
}
/**
@@ -367,6 +385,17 @@ private URIBuilder makeBaseURI(UUID requestId) {
throw new IllegalArgumentException();
}
+ // construct the builder object without param
+ URIBuilder builder = makeBaseURI();
+
+ // set the request id
+ builder.setParameter(REQUEST_ID, requestId.toString());
+
+ return builder;
+ }
+
+ /** Construct a URIBuilder for common parts of request without any parameter */
+ private URIBuilder makeBaseURI() {
// construct the builder object
URIBuilder builder = new URIBuilder();
@@ -379,9 +408,6 @@ private URIBuilder makeBaseURI(UUID requestId) {
// set the port name
builder.setPort(port);
- // set the request id
- builder.setParameter(REQUEST_ID, requestId.toString());
-
return builder;
}
@@ -492,79 +518,14 @@ private URI makeHistoryRangeURI(
}
/**
- * Given a request UUID, and a fully qualified pipe name make a URI for configure snowpipe client
- * http://snowflakeURL{:PORT}/v1/data/pipes/{pipeName}/client/configure
- *
- * Where snowflake URL can be an old url or new regionless URL.
- *
- *
Request Body is Empty
- *
- *
And our response looks like:
- *
- *
- * {
- * 'clientSequencer': LONG
- * }
- *
- *
- * @param requestId the label for this request
- * @param pipe the pipe name
- * @return configure snowpipe client URI
- * @throws URISyntaxException
- */
- private URI makeConfigureClientURI(UUID requestId, String pipe) throws URISyntaxException {
- if (pipe == null) {
- throw new IllegalArgumentException();
- }
- URIBuilder builder = makeBaseURI(requestId);
- builder.setPath(String.format(CONFIGURE_CLIENT_ENDPOINT_FORMAT, pipe));
- LOGGER.info("Final Configure Client URIBuilder - {}", builder);
- return builder.build();
- }
-
- /**
- * makeGetClientURI - Given a request UUID, and a fully qualified pipe name make a URI for getting
- * snowpipe client http://snowflakeURL{:PORT}/v1/data/pipes/{pipeName}/client/status
- *
- * Where snowflake URL can be an old url or new regionless URL.
- *
- *
Request Body is Empty
- *
- *
And our response looks like:
- *
- *
- * {
- * 'offsetToken': STRING
- * 'clientSequencer': LONG
- * }
- *
- *
- * @param requestId the label for this request
- * @param pipe the pipe name
- * @return get client URI
- * @throws URISyntaxException
- */
- private URI makeGetClientURI(UUID requestId, String pipe) throws URISyntaxException {
- if (pipe == null) {
- throw new IllegalArgumentException();
- }
- URIBuilder builder = makeBaseURI(requestId);
- builder.setPath(String.format(CLIENT_STATUS_ENDPOINT_FORMAT, pipe));
- LOGGER.info("Final Get Client URIBuilder - {}", builder);
- return builder.build();
- }
-
- /**
- * Given a list of files, and an optional clientInfo generate json string which later can be
- * passed in request body of insertFiles API
+ * Given a list of files, generate a json string which later can be passed in request body of
+ * insertFiles API
*
* @param files the list of files we want to send
- * @param clientInfo optional clientInfo which can be empty.
* @return the string json blob
* @throws IllegalArgumentException if files passed in is null
*/
- private String serializeInsertFilesRequest(
- List files, Optional clientInfo) {
+ private String serializeInsertFilesRequest(List files) {
// if the files argument is null, throw
if (files == null) {
LOGGER.info("Null files argument in RequestBuilder");
@@ -572,10 +533,7 @@ private String serializeInsertFilesRequest(
}
// create pojo
- IngestRequest ingestRequest =
- clientInfo
- .map(insertFilesClientInfo -> new IngestRequest(files, insertFilesClientInfo))
- .orElseGet(() -> new IngestRequest(files));
+ IngestRequest ingestRequest = new IngestRequest(files);
// serialize to a string
try {
@@ -605,21 +563,26 @@ private static void addUserAgent(HttpUriRequest request, String userAgentSuffix)
}
/**
- * addToken - adds a the JWT token to a request
+ * addToken - adds a token to a request
*
* @param request the URI request
- * @param token the token to add
*/
- private static void addToken(HttpUriRequest request, String token) {
- request.setHeader(HttpHeaders.AUTHORIZATION, BEARER_PARAMETER + token);
- request.setHeader(SF_HEADER_AUTHORIZATION_TOKEN_TYPE, JWT_TOKEN_TYPE);
+ public void addToken(HttpUriRequest request) {
+ request.setHeader(HttpHeaders.AUTHORIZATION, BEARER_PARAMETER + securityManager.getToken());
+ request.setHeader(SF_HEADER_AUTHORIZATION_TOKEN_TYPE, this.securityManager.getTokenType());
}
- private static void addHeaders(HttpUriRequest request, String token, String userAgentSuffix) {
+ /**
+ * addHeader - adds necessary header to a request
+ *
+ * @param request the URI request
+ * @param userAgentSuffix user agent suffix
+ */
+ private void addHeaders(HttpUriRequest request, String userAgentSuffix) {
addUserAgent(request, userAgentSuffix);
// Add the auth token
- addToken(request, token);
+ addToken(request);
// Add Accept header
request.setHeader(HttpHeaders.ACCEPT, HTTP_HEADER_CONTENT_TYPE_JSON);
@@ -639,28 +602,6 @@ private static void addHeaders(HttpUriRequest request, String token, String user
public HttpPost generateInsertRequest(
UUID requestId, String pipe, List files, boolean showSkippedFiles)
throws URISyntaxException {
- return generateInsertRequest(requestId, pipe, files, showSkippedFiles, Optional.empty());
- }
-
- /**
- * generateInsertRequest - given a pipe, list of files and clientInfo, make a request for the
- * insert endpoint
- *
- * @param requestId a UUID we will use to label this request
- * @param pipe a fully qualified pipe name
- * @param files a list of files
- * @param showSkippedFiles a boolean which returns skipped files when set to true
- * @param clientInfo
- * @return a post request with all the data we need
- * @throws URISyntaxException if the URI components provided are improper
- */
- public HttpPost generateInsertRequest(
- UUID requestId,
- String pipe,
- List files,
- boolean showSkippedFiles,
- Optional clientInfo)
- throws URISyntaxException {
// make the insert URI
URI insertURI = makeInsertURI(requestId, pipe, showSkippedFiles);
LOGGER.info("Created Insert Request : {} ", insertURI);
@@ -668,12 +609,11 @@ public HttpPost generateInsertRequest(
// Make the post request
HttpPost post = new HttpPost(insertURI);
- addHeaders(post, securityManager.getToken(), this.userAgentSuffix);
+ addHeaders(post, this.userAgentSuffix);
// the entity for the containing the json
final StringEntity entity =
- new StringEntity(
- serializeInsertFilesRequest(files, clientInfo), ContentType.APPLICATION_JSON);
+ new StringEntity(serializeInsertFilesRequest(files), ContentType.APPLICATION_JSON);
post.setEntity(entity);
return post;
@@ -697,7 +637,7 @@ public HttpGet generateHistoryRequest(
// make the get request
HttpGet get = new HttpGet(historyURI);
- addHeaders(get, securityManager.getToken(), this.userAgentSuffix);
+ addHeaders(get, this.userAgentSuffix);
return get;
}
@@ -723,7 +663,7 @@ public HttpGet generateHistoryRangeRequest(
HttpGet get = new HttpGet(historyRangeURI);
- addHeaders(get, securityManager.getToken(), this.userAgentSuffix /*User agent information*/);
+ addHeaders(get, this.userAgentSuffix /*User agent information*/);
return get;
}
@@ -751,7 +691,7 @@ public HttpPost generateStreamingIngestPostRequest(
// Make the post request
HttpPost post = new HttpPost(uri);
- addHeaders(post, securityManager.getToken(), this.userAgentSuffix /*User agent information*/);
+ addHeaders(post, this.userAgentSuffix /*User agent information*/);
// The entity for the containing the json
final StringEntity entity = new StringEntity(payload, ContentType.APPLICATION_JSON);
@@ -760,22 +700,6 @@ public HttpPost generateStreamingIngestPostRequest(
return post;
}
- /**
- * Given a requestId and a pipe, make a configure client request
- *
- * @param requestID a UUID we will use to label this request
- * @param pipe a fully qualified pipe name
- * @return configure client request
- * @throws URISyntaxException
- */
- public HttpPost generateConfigureClientRequest(UUID requestID, String pipe)
- throws URISyntaxException {
- URI configureClientURI = makeConfigureClientURI(requestID, pipe);
- HttpPost post = new HttpPost(configureClientURI);
- addHeaders(post, securityManager.getToken(), this.userAgentSuffix);
- return post;
- }
-
/**
* Generate post request for streaming ingest related APIs
*
@@ -798,19 +722,25 @@ public HttpPost generateStreamingIngestPostRequest(
}
/**
- * Given a requestId and a pipe, make a get client status request
+ * 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
*
- * @param requestID UUID
- * @param pipe a fully qualified pipe name
- * @return get client status request
- * @throws URISyntaxException
+ * @param refreshToken the new refresh token
*/
- public HttpGet generateGetClientStatusRequest(UUID requestID, String pipe)
- throws URISyntaxException {
- URI getClientStatusURI = makeGetClientURI(requestID, pipe);
- HttpGet get = new HttpGet(getClientStatusURI);
- addHeaders(get, securityManager.getToken(), this.userAgentSuffix);
- return get;
+ public void setRefreshToken(String refreshToken) {
+ if (securityManager instanceof OAuthManager) {
+ ((OAuthManager) securityManager).setRefreshToken(refreshToken);
+ }
+ }
+
+ /** Get authorization type */
+ public String getAuthType() {
+ return securityManager.getTokenType();
+ }
+
+ /** Refresh token manually */
+ public void refreshToken() {
+ securityManager.refreshToken();
}
/**
@@ -822,5 +752,13 @@ public HttpGet generateGetClientStatusRequest(UUID requestID, String pipe)
*/
public void closeResources() {
securityManager.close();
+ if (telemetryService != null) {
+ telemetryService.close();
+ }
+ }
+
+ /** Get the telemetry service */
+ public TelemetryService getTelemetryService() {
+ return telemetryService;
}
}
diff --git a/src/main/java/net/snowflake/ingest/connection/SecurityManager.java b/src/main/java/net/snowflake/ingest/connection/SecurityManager.java
index 7832ab88e..3fe89aeff 100644
--- a/src/main/java/net/snowflake/ingest/connection/SecurityManager.java
+++ b/src/main/java/net/snowflake/ingest/connection/SecurityManager.java
@@ -1,128 +1,71 @@
/*
- * Copyright (c) 2012-2017 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
*/
package net.snowflake.ingest.connection;
-import com.nimbusds.jose.JOSEException;
-import com.nimbusds.jose.JWSAlgorithm;
-import com.nimbusds.jose.JWSHeader;
-import com.nimbusds.jose.JWSSigner;
-import com.nimbusds.jose.crypto.RSASSASigner;
-import com.nimbusds.jwt.JWTClaimsSet;
-import com.nimbusds.jwt.SignedJWT;
-import java.security.KeyPair;
-import java.util.Date;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
-import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicReference;
-import net.snowflake.ingest.utils.Cryptor;
import net.snowflake.ingest.utils.ThreadFactoryUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
- * This class manages creating and automatically renewing the JWT token
+ * This class manages creating and automatically renewing the token
*
* @author obabarinsa
- * @since 1.8
*/
-final class SecurityManager implements AutoCloseable {
+abstract class SecurityManager implements AutoCloseable {
// the logger for SecurityManager
- private static final Logger LOGGER = LoggerFactory.getLogger(SecurityManager.class);
+ protected static final Logger LOGGER = LoggerFactory.getLogger(SecurityManager.class);
- // the token lifetime is 59 minutes
- private static final float LIFETIME = 59;
-
- // the renewal time is 54 minutes
- private static final int RENEWAL_INTERVAL = 54;
-
- // The public - private key pair we're using to connect to the service
- private transient KeyPair keyPair;
-
- // the name of the account on behalf of which we're connecting
- private String account;
+ protected final String account;
// Fingerprint of public key sent from client in jwt payload
- private String publicKeyFingerPrint;
+ protected String publicKeyFingerPrint;
// the name of the user who will be loading the files
- private String user;
+ protected final String user;
- // the token itself
- private AtomicReference token;
-
- // Did we fail to regenerate our token at some point?
- private AtomicBoolean regenFailed;
+ // Did we fail to refresh our token at some point?
+ protected final AtomicBoolean refreshFailed;
// Thread factory for daemon threads so that application can shutdown
final ThreadFactory tf = ThreadFactoryUtil.poolThreadFactory(getClass().getSimpleName(), true);
- // the thread we use for renewing all tokens
- private final ScheduledExecutorService keyRenewer = Executors.newScheduledThreadPool(1, tf);
+ // the thread we use for refresh tokens
+ protected ScheduledExecutorService tokenRefresher = Executors.newScheduledThreadPool(1, tf);
+
+ // Reference to the Telemetry Service in the client
+ protected final TelemetryService telemetryService;
/**
- * Creates a SecurityManager entity for a given account, user and KeyPair with a specified time to
- * renew the token
+ * Creates a SecurityManager entity for a given account
*
- * @param accountname - the snowflake account name of this user
+ * @param accountName - the snowflake account name of this user
* @param username - the snowflake username of the current user
- * @param keyPair - the public/private key pair we're using to connect
- * @param timeTillRenewal - the time measure until we renew the token
- * @param unit the unit by which timeTillRenewal is measured
*/
- SecurityManager(
- String accountname, String username, KeyPair keyPair, int timeTillRenewal, TimeUnit unit) {
+ SecurityManager(String accountName, String username, TelemetryService telemetryService) {
// if any of our arguments are null, throw an exception
- if (accountname == null || username == null || keyPair == null) {
+ if (accountName == null || username == null) {
throw new IllegalArgumentException();
}
- account = parseAccount(accountname);
+ account = parseAccount(accountName);
user = username.toUpperCase();
- // create our automatic reference to a string (our token)
- token = new AtomicReference<>();
-
// we haven't yet failed to regenerate our token
- regenFailed = new AtomicBoolean();
-
- // we have to keep around the keys
- this.keyPair = keyPair;
-
- // generate our first token
- regenerateToken();
+ this.refreshFailed = new AtomicBoolean();
- // schedule all future renewals
- keyRenewer.scheduleAtFixedRate(this::regenerateToken, timeTillRenewal, timeTillRenewal, unit);
+ this.telemetryService = telemetryService;
}
- /**
- * Creates a SecurityManager entity for a given account, user and KeyPair with the default time to
- * renew (54 minutes)
- *
- * @param accountname - the snowflake account name of this user
- * @param username - the snowflake username of the current user
- * @param keyPair - the public/private key pair we're using to connect
- */
- SecurityManager(String accountname, String username, KeyPair keyPair) {
- this(accountname, username, keyPair, RENEWAL_INTERVAL, TimeUnit.MINUTES);
+ /** Test only */
+ void setRefreshFailed(Boolean refreshFailed) {
+ this.refreshFailed.set(refreshFailed);
}
- /**
- * Trims an account name if it contains a "."
- *
- * Snowflake's python connector trims an accountname if it contains a "."
- *
- * @param accountName given accountName in SimpleIngestManager's constructor
- * @return initial part of account name if originally it contained a ".", otherwise return the
- * same accountName.
- * @see Python
- * Connector
- */
private String parseAccount(final String accountName) {
String parseAccount = null;
if (accountName.contains(".")) {
@@ -136,94 +79,19 @@ private String parseAccount(final String accountName) {
return parseAccount.toUpperCase();
}
- /** regenerateToken - Regenerates our Token given our current user, account and keypair */
- private void regenerateToken() {
- // create our JWT claim builder object
- JWTClaimsSet.Builder builder = new JWTClaimsSet.Builder();
-
- // set the subject to the fully qualified username
- String subject = String.format("%s.%s", account, user);
- LOGGER.info("Creating Token with subject {}", subject);
-
- // set the issuer
- String publicKeyFPInJwt = calculatePublicKeyFp(keyPair);
- String issuer = String.format("%s.%s.%s", account, user, publicKeyFPInJwt);
- LOGGER.info("Creating Token with issuer {}", issuer);
-
- // iat set to now
- Date iat = new Date(System.currentTimeMillis());
-
- // expiration in 59 minutes
- Date exp = new Date(iat.getTime() + 59 * 60 * 1000);
-
- // build claim set
- JWTClaimsSet claimsSet =
- builder.issuer(issuer).subject(subject).issueTime(iat).expirationTime(exp).build();
-
- SignedJWT signedJWT = new SignedJWT(new JWSHeader(JWSAlgorithm.RS256), claimsSet);
-
- JWSSigner signer = new RSASSASigner(this.keyPair.getPrivate());
-
- String newToken;
- try {
- signedJWT.sign(signer);
- newToken = signedJWT.serialize();
- } catch (JOSEException e) {
- regenFailed.set(true);
- LOGGER.error("Failed to regenerate token! Exception is as follows : {}", e.getMessage());
- throw new SecurityException();
- }
-
- // atomically update the string
- LOGGER.info("Created new JWT");
- token.set(newToken);
- }
-
/**
* getToken - returns we've most recently generated
*
- * @return the string version of a valid JWT token
- * @throws SecurityException if we failed to regenerate a token since the last call
- */
- String getToken() {
- // if we failed to regenerate the token at some point, throw
- if (regenFailed.get()) {
- LOGGER.error("getToken request failed due to token regeneration failure");
- throw new SecurityException();
- }
-
- return token.get();
- }
-
- /* Only used in testing at the moment */
- String getAccount() {
- return this.account;
- }
-
- /**
- * Given a keypair
- *
- * @return the fingerprint of public key
- *
The idea is to hash public key's raw bytes using SHA-256 and converts hash into a string
- * using Base64 encoding.
+ * @return the string version of a valid token
+ * @throws SecurityException if we failed to regenerate or refresh a token since the last call
*/
- private String calculatePublicKeyFp(KeyPair keyPair) {
- // get the raw bytes of public key
- byte[] publicKeyRawBytes = keyPair.getPublic().getEncoded();
+ abstract String getToken();
- // take sha256 on raw bytes and do base64 encode
- publicKeyFingerPrint = String.format("SHA256:%s", Cryptor.sha256HashBase64(publicKeyRawBytes));
- return publicKeyFingerPrint;
- }
+ /** getTokenType - returns the token type, either "KEYPAIR_JWT" or "OAUTH" */
+ abstract String getTokenType();
- /** Only called by SecurityManagerTest */
- String getPublicKeyFingerPrint() {
- return publicKeyFingerPrint;
- }
+ /** refreshToken - regenerate or fetch a new token */
+ abstract void refreshToken();
- /** Currently, it only shuts down the instance of ExecutorService. */
- @Override
- public void close() {
- keyRenewer.shutdown();
- }
+ public abstract void close();
}
diff --git a/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java b/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java
index 8c9590a22..bf20afda8 100644
--- a/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java
+++ b/src/main/java/net/snowflake/ingest/connection/ServiceResponseHandler.java
@@ -8,11 +8,16 @@
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.util.UUID;
+import net.snowflake.client.jdbc.internal.apache.http.HttpEntity;
+import net.snowflake.client.jdbc.internal.apache.http.HttpResponse;
+import net.snowflake.client.jdbc.internal.apache.http.HttpStatus;
+import net.snowflake.client.jdbc.internal.apache.http.StatusLine;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpGet;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpPost;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpUriRequest;
+import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient;
+import net.snowflake.client.jdbc.internal.apache.http.util.EntityUtils;
import net.snowflake.ingest.utils.BackOffException;
-import org.apache.http.HttpResponse;
-import org.apache.http.HttpStatus;
-import org.apache.http.StatusLine;
-import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -33,9 +38,8 @@ public enum ApiName {
INSERT_FILES("POST"),
INSERT_REPORT("GET"),
LOAD_HISTORY_SCAN("GET"),
- CLIENT_CONFIGURE("POST"),
- CLIENT_STATUS("GET"),
STREAMING_OPEN_CHANNEL("POST"),
+ STREAMING_DROP_CHANNEL("POST"),
STREAMING_CHANNEL_STATUS("POST"),
STREAMING_REGISTER_BLOB("POST"),
STREAMING_CLIENT_CONFIGURE("POST");
@@ -75,12 +79,20 @@ private static boolean isStatusOK(StatusLine statusLine) {
*
* @param response the HTTPResponse we want to distill into an IngestResponse
* @param requestId
+ * @param httpClient HttpClient for retries
+ * @param httpPostForIngestFile HttpRequest for retries
+ * @param builder RequestBuilder for retries
* @return An IngestResponse with all of the parsed out information
* @throws IOException if our entity is somehow corrupt or we can't get it
* @throws IngestResponseException - if we have an uncategorized network issue
* @throws BackOffException - if we have a 503 issue
*/
- public static IngestResponse unmarshallIngestResponse(HttpResponse response, UUID requestId)
+ public static IngestResponse unmarshallIngestResponse(
+ HttpResponse response,
+ UUID requestId,
+ CloseableHttpClient httpClient,
+ HttpPost httpPostForIngestFile,
+ RequestBuilder builder)
throws IOException, IngestResponseException, BackOffException {
// we can't unmarshall a null response
if (response == null) {
@@ -89,10 +101,11 @@ public static IngestResponse unmarshallIngestResponse(HttpResponse response, UUI
}
// handle the exceptional status code
- handleExceptionalStatus(response, requestId, ApiName.INSERT_FILES);
+ response =
+ handleExceptionalStatus(
+ response, requestId, ApiName.INSERT_FILES, httpClient, httpPostForIngestFile, builder);
- // grab the response entity
- String blob = EntityUtils.toString(response.getEntity());
+ String blob = consumeAndReturnResponseEntityAsString(response.getEntity());
// Read out the blob entity into a class
return mapper.readValue(blob, IngestResponse.class);
@@ -104,12 +117,20 @@ public static IngestResponse unmarshallIngestResponse(HttpResponse response, UUI
*
* @param response the HttpResponse object we are trying to deserialize
* @param requestId
+ * @param httpClient HttpClient for retries
+ * @param httpGetHistory HttpRequest for retries
+ * @param builder RequestBuilder for retries
* @return a HistoryResponse with all the parsed out information
* @throws IOException if our entity is somehow corrupt or we can't get it
* @throws IngestResponseException - if we have an uncategorized network issue
* @throws BackOffException - if we have a 503 issue
*/
- public static HistoryResponse unmarshallHistoryResponse(HttpResponse response, UUID requestId)
+ public static HistoryResponse unmarshallHistoryResponse(
+ HttpResponse response,
+ UUID requestId,
+ CloseableHttpClient httpClient,
+ HttpGet httpGetHistory,
+ RequestBuilder builder)
throws IOException, IngestResponseException, BackOffException {
// we can't unmarshall a null response
if (response == null) {
@@ -118,10 +139,11 @@ public static HistoryResponse unmarshallHistoryResponse(HttpResponse response, U
}
// handle the exceptional status code
- handleExceptionalStatus(response, requestId, ApiName.INSERT_REPORT);
+ response =
+ handleExceptionalStatus(
+ response, requestId, ApiName.INSERT_REPORT, httpClient, httpGetHistory, builder);
- // grab the string version of the response entity
- String blob = EntityUtils.toString(response.getEntity());
+ String blob = consumeAndReturnResponseEntityAsString(response.getEntity());
// read out our blob into a pojo
return mapper.readValue(blob, HistoryResponse.class);
@@ -132,13 +154,20 @@ public static HistoryResponse unmarshallHistoryResponse(HttpResponse response, U
*
* @param response the HttpResponse object we are trying to deserialize
* @param requestId
+ * @param httpClient HttpClient for retries
+ * @param request HttpRequest for retries
+ * @param builder HttpBuilder for retries
* @return HistoryRangeResponse
* @throws IOException if our entity is somehow corrupt or we can't get it
* @throws IngestResponseException - if we have an uncategorized network issue
* @throws BackOffException - if we have a 503 issue
*/
public static HistoryRangeResponse unmarshallHistoryRangeResponse(
- HttpResponse response, UUID requestId)
+ HttpResponse response,
+ UUID requestId,
+ CloseableHttpClient httpClient,
+ HttpGet request,
+ RequestBuilder builder)
throws IOException, IngestResponseException, BackOffException {
// we can't unmarshall a null response
@@ -148,73 +177,15 @@ public static HistoryRangeResponse unmarshallHistoryRangeResponse(
}
// handle the exceptional status code
- handleExceptionalStatus(response, requestId, ApiName.LOAD_HISTORY_SCAN);
-
- // grab the string version of the response entity
- String blob = EntityUtils.toString(response.getEntity());
+ response =
+ handleExceptionalStatus(
+ response, requestId, ApiName.LOAD_HISTORY_SCAN, httpClient, request, builder);
+ String blob = consumeAndReturnResponseEntityAsString(response.getEntity());
// read out our blob into a pojo
return mapper.readValue(blob, HistoryRangeResponse.class);
}
- /**
- * unmarshallConfigureClientResponse - Given an HttpResponse object, attempts to deserialize it
- * into a ConfigureClientResponse
- *
- * @param response HttpResponse
- * @param requestId
- * @return ConfigureClientResponse
- * @throws IOException if our entity is somehow corrupt or we can't get it
- * @throws IngestResponseException - if we have an uncategorized network issue
- * @throws BackOffException - if we have a 503 issue
- */
- public static ConfigureClientResponse unmarshallConfigureClientResponse(
- HttpResponse response, UUID requestId)
- throws IOException, IngestResponseException, BackOffException {
- if (response == null) {
- LOGGER.warn("Null response passed to unmarshallConfigureClientResponse");
- throw new IllegalArgumentException();
- }
-
- // handle the exceptional status code
- handleExceptionalStatus(response, requestId, ApiName.CLIENT_CONFIGURE);
-
- // grab the string version of the response entity
- String blob = EntityUtils.toString(response.getEntity());
-
- // read out our blob into a pojo
- return mapper.readValue(blob, ConfigureClientResponse.class);
- }
-
- /**
- * unmarshallGetClientStatus - Given an HttpResponse object, attempts to deserialize it into a
- * ClientStatusResponse
- *
- * @param response HttpResponse
- * @param requestId
- * @return ClientStatusResponse
- * @throws IOException if our entity is somehow corrupt or we can't get it
- * @throws IngestResponseException - if we have an uncategorized network issue
- * @throws BackOffException - if we have a 503 issue
- */
- public static ClientStatusResponse unmarshallGetClientStatus(
- HttpResponse response, UUID requestId)
- throws IOException, IngestResponseException, BackOffException {
- if (response == null) {
- LOGGER.warn("Null response passed to unmarshallClientStatusResponse");
- throw new IllegalArgumentException();
- }
-
- // handle the exceptional status code
- handleExceptionalStatus(response, requestId, ApiName.CLIENT_STATUS);
-
- // grab the string version of the response entity
- String blob = EntityUtils.toString(response.getEntity());
-
- // read out our blob into a pojo
- return mapper.readValue(blob, ClientStatusResponse.class);
- }
-
/**
* unmarshallStreamingIngestResponse Given an HttpResponse object - attempts to deserialize it
* into an Object based on input type
@@ -227,7 +198,12 @@ public static ClientStatusResponse unmarshallGetClientStatus(
* @throws IngestResponseException if received an exceptional status code
*/
public static T unmarshallStreamingIngestResponse(
- HttpResponse response, Class valueType, ApiName apiName)
+ HttpResponse response,
+ Class valueType,
+ ApiName apiName,
+ CloseableHttpClient httpClient,
+ HttpUriRequest request,
+ RequestBuilder requestBuilder)
throws IOException, IngestResponseException {
// We can't unmarshall a null response
if (response == null) {
@@ -236,10 +212,11 @@ public static T unmarshallStreamingIngestResponse(
}
// Handle the exceptional status code
- handleExceptionalStatus(response, null, apiName);
+ response =
+ handleExceptionalStatus(response, null, apiName, httpClient, request, requestBuilder);
// Grab the string version of the response entity
- String blob = EntityUtils.toString(response.getEntity());
+ String blob = consumeAndReturnResponseEntityAsString(response.getEntity());
// Read out our blob into a pojo
return mapper.readValue(blob, valueType);
@@ -250,35 +227,73 @@ public static T unmarshallStreamingIngestResponse(
*
* @param response HttpResponse
* @param requestId
+ * @param apiName enum to represent the corresponding api name
+ * @param httpClient HttpClient for retries
+ * @param request HttpRequest for retries
+ * @param requestBuilder RequestBuilder for retries
+ * @return modified HttpResponse
* @throws IOException if our entity is somehow corrupt or we can't get it
* @throws IngestResponseException - for all other non OK status
* @throws BackOffException - if we have a 503 issue
*/
- private static void handleExceptionalStatus(
- HttpResponse response, UUID requestId, ApiName apiName)
+ private static HttpResponse handleExceptionalStatus(
+ HttpResponse response,
+ UUID requestId,
+ ApiName apiName,
+ CloseableHttpClient httpClient,
+ HttpUriRequest request,
+ RequestBuilder requestBuilder)
throws IOException, IngestResponseException, BackOffException {
- StatusLine statusLine = response.getStatusLine();
- if (!isStatusOK(statusLine)) {
+ if (!isStatusOK(response.getStatusLine())) {
+ StatusLine statusLine = response.getStatusLine();
+ LOGGER.warn(
+ "{} Status hit from {}, requestId:{}",
+ statusLine.getStatusCode(),
+ apiName,
+ requestId == null ? "" : requestId.toString());
+
// if we have a 503 exception throw a backoff
switch (statusLine.getStatusCode()) {
// If we have a 503, BACKOFF
case HttpStatus.SC_SERVICE_UNAVAILABLE:
- LOGGER.warn(
- "503 Status hit from {}, backoff, requestId:{}",
- apiName,
- requestId == null ? "" : requestId.toString());
throw new BackOffException();
+ case HttpStatus.SC_UNAUTHORIZED:
+ LOGGER.warn("Authorization failed, refreshing Token succeeded, retry");
+ requestBuilder.refreshToken();
+ requestBuilder.addToken(request);
+ response = httpClient.execute(request);
+ if (!isStatusOK(response.getStatusLine())) {
+ throw new SecurityException("Authorization failed after retry");
+ }
+ break;
default:
- LOGGER.error(
- "Exceptional Status Code from {}: {}, requestId:{}",
- apiName,
- statusLine.getStatusCode(),
- requestId == null ? "" : requestId.toString());
- String blob = EntityUtils.toString(response.getEntity());
+ String blob = consumeAndReturnResponseEntityAsString(response.getEntity());
throw new IngestResponseException(
statusLine.getStatusCode(),
IngestResponseException.IngestExceptionBody.parseBody(blob));
}
}
+ return response;
+ }
+
+ /**
+ * Consumes the HttpEntity as mentioned in HttpClient Docs
+ *
+ * Also returns the string version of this entity.
+ *
+ * @param httpResponseEntity the response entity obtained after successfully calling associated
+ * Rest APIs
+ * @return String version of this http response which will be later used to deserialize into
+ * respective Response Object
+ * @throws IOException if parsing error
+ */
+ private static String consumeAndReturnResponseEntityAsString(HttpEntity httpResponseEntity)
+ throws IOException {
+ // grab the string version of the response entity
+ String responseEntityAsString = EntityUtils.toString(httpResponseEntity);
+
+ EntityUtils.consumeQuietly(httpResponseEntity);
+ return responseEntityAsString;
}
}
diff --git a/src/main/java/net/snowflake/ingest/connection/SnowflakeOAuthClient.java b/src/main/java/net/snowflake/ingest/connection/SnowflakeOAuthClient.java
new file mode 100644
index 000000000..a90540659
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/connection/SnowflakeOAuthClient.java
@@ -0,0 +1,138 @@
+/*
+ * Copyright (c) 2023 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.connection;
+
+import java.io.UnsupportedEncodingException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.net.URLEncoder;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
+import net.snowflake.client.jdbc.internal.apache.http.HttpHeaders;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.CloseableHttpResponse;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpPost;
+import net.snowflake.client.jdbc.internal.apache.http.client.methods.HttpUriRequest;
+import net.snowflake.client.jdbc.internal.apache.http.client.utils.URIBuilder;
+import net.snowflake.client.jdbc.internal.apache.http.entity.ContentType;
+import net.snowflake.client.jdbc.internal.apache.http.entity.StringEntity;
+import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient;
+import net.snowflake.client.jdbc.internal.apache.http.util.EntityUtils;
+import net.snowflake.client.jdbc.internal.google.api.client.http.HttpStatusCodes;
+import net.snowflake.client.jdbc.internal.google.gson.JsonObject;
+import net.snowflake.client.jdbc.internal.google.gson.JsonParser;
+import net.snowflake.ingest.utils.ErrorCode;
+import net.snowflake.ingest.utils.HttpUtil;
+import net.snowflake.ingest.utils.SFException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/*
+ * Implementation of Snowflake OAuth Client, used for refreshing an OAuth access token.
+ */
+public class SnowflakeOAuthClient implements OAuthClient {
+
+ static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeOAuthClient.class);
+ private static final String TOKEN_REQUEST_ENDPOINT = "/oauth/token-request";
+
+ // Content type header to specify the encoding
+ private static final String OAUTH_CONTENT_TYPE_HEADER = "application/x-www-form-urlencoded";
+ private static final String GRANT_TYPE_PARAM = "grant_type";
+ private static final String ACCESS_TOKEN = "access_token";
+ private static final String REFRESH_TOKEN = "refresh_token";
+ private static final String EXPIRES_IN = "expires_in";
+
+ // OAuth credential
+ private final AtomicReference oAuthCredential;
+
+ // exact uri for token request
+ private final URI tokenRequestURI;
+
+ // Http client for submitting token refresh request
+ private final CloseableHttpClient httpClient;
+
+ /**
+ * Creates an SnowflakeOAuthClient given account, credential and base uri
+ *
+ * @param accountName - the snowflake account name of this user
+ * @param oAuthCredential - the OAuth credential we're using to connect
+ * @param baseURIBuilder - the uri builder with common scheme, host and port
+ */
+ SnowflakeOAuthClient(
+ String accountName, OAuthCredential oAuthCredential, URIBuilder baseURIBuilder) {
+ this.oAuthCredential = new AtomicReference<>(oAuthCredential);
+
+ // build token request uri
+ baseURIBuilder.setPath(TOKEN_REQUEST_ENDPOINT);
+ try {
+ this.tokenRequestURI = baseURIBuilder.build();
+ } catch (URISyntaxException e) {
+ throw new SFException(e, ErrorCode.MAKE_URI_FAILURE, e.getMessage());
+ }
+
+ this.httpClient = HttpUtil.getHttpClient(accountName);
+ }
+
+ /** Get access token */
+ @Override
+ public AtomicReference getoAuthCredentialRef() {
+ return oAuthCredential;
+ }
+
+ /** Refresh access token using a valid refresh token */
+ @Override
+ public void refreshToken() {
+ String respBodyString = null;
+ try (CloseableHttpResponse httpResponse = httpClient.execute(makeRefreshTokenRequest())) {
+ respBodyString = EntityUtils.toString(httpResponse.getEntity());
+
+ if (httpResponse.getStatusLine().getStatusCode() == HttpStatusCodes.STATUS_CODE_OK) {
+ JsonObject respBody = JsonParser.parseString(respBodyString).getAsJsonObject();
+
+ if (respBody.has(ACCESS_TOKEN) && respBody.has(EXPIRES_IN)) {
+ // Trim surrounding quotation marks
+ String newAccessToken = respBody.get(ACCESS_TOKEN).toString().replaceAll("^\"|\"$", "");
+
+ oAuthCredential.get().setAccessToken(newAccessToken);
+ oAuthCredential.get().setExpiresIn(respBody.get(EXPIRES_IN).getAsInt());
+ return;
+ }
+ }
+ throw new SFException(
+ ErrorCode.OAUTH_REFRESH_TOKEN_ERROR,
+ "Refresh access token fail with response: " + respBodyString);
+ } catch (Exception e) {
+ throw new SFException(ErrorCode.OAUTH_REFRESH_TOKEN_ERROR, e.getMessage());
+ }
+ }
+
+ /** Helper method for making refresh request */
+ private HttpUriRequest makeRefreshTokenRequest() {
+ HttpPost post = new HttpPost(tokenRequestURI);
+ post.addHeader(HttpHeaders.CONTENT_TYPE, OAUTH_CONTENT_TYPE_HEADER);
+ post.addHeader(HttpHeaders.AUTHORIZATION, oAuthCredential.get().getAuthHeader());
+
+ Map payload = new HashMap<>();
+ try {
+ payload.put(GRANT_TYPE_PARAM, URLEncoder.encode(REFRESH_TOKEN, "UTF-8"));
+ payload.put(
+ REFRESH_TOKEN, URLEncoder.encode(oAuthCredential.get().getRefreshToken(), "UTF-8"));
+ } catch (UnsupportedEncodingException e) {
+ throw new SFException(e, ErrorCode.OAUTH_REFRESH_TOKEN_ERROR, e.getMessage());
+ }
+
+ String payloadString =
+ payload.entrySet().stream()
+ .map(e -> e.getKey() + "=" + e.getValue())
+ .collect(Collectors.joining("&"));
+
+ final StringEntity entity =
+ new StringEntity(payloadString, ContentType.APPLICATION_FORM_URLENCODED);
+ post.setEntity(entity);
+
+ return post;
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/connection/TelemetryService.java b/src/main/java/net/snowflake/ingest/connection/TelemetryService.java
new file mode 100644
index 000000000..15e2a0c5a
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/connection/TelemetryService.java
@@ -0,0 +1,184 @@
+/*
+ * Copyright (c) 2022 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.connection;
+
+import static net.snowflake.ingest.connection.RequestBuilder.DEFAULT_VERSION;
+
+import com.codahale.metrics.Histogram;
+import com.codahale.metrics.Meter;
+import com.codahale.metrics.Snapshot;
+import com.codahale.metrics.Timer;
+import com.google.common.util.concurrent.RateLimiter;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient;
+import net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.ObjectMapper;
+import net.snowflake.client.jdbc.internal.fasterxml.jackson.databind.node.ObjectNode;
+import net.snowflake.client.jdbc.telemetry.TelemetryClient;
+import net.snowflake.client.jdbc.telemetry.TelemetryUtil;
+import net.snowflake.ingest.utils.Logging;
+
+/**
+ * Telemetry service to collect logs in the SDK and send them to Snowflake through the JDBC client
+ * telemetry API
+ */
+public class TelemetryService {
+ // Enum for different client telemetries
+ enum TelemetryType {
+ STREAMING_INGEST_LATENCY_IN_SEC("streaming_ingest_latency_in_ms"),
+ STREAMING_INGEST_CLIENT_FAILURE("streaming_ingest_client_failure"),
+ STREAMING_INGEST_THROUGHPUT_BYTES_PER_SEC("streaming_ingest_throughput_bytes_per_sec"),
+ STREAMING_INGEST_CPU_MEMORY_USAGE("streaming_ingest_cpu_memory_usage"),
+ STREAMING_INGEST_BATCH_OFFSET_MISMATCH("streaming_ingest_batch_offset_mismatch");
+
+ private final String name;
+
+ TelemetryType(String name) {
+ this.name = name;
+ }
+
+ @Override
+ public String toString() {
+ return this.name;
+ }
+ }
+
+ private static final Logging logger = new Logging(TelemetryService.class);
+ private static final ObjectMapper MAPPER = new ObjectMapper();
+
+ private static final String TYPE = "type";
+ private static final String CLIENT_NAME = "client_name";
+ private static final String COUNT = "count";
+ private static final String MAX = "max";
+ private static final String MIN = "min";
+ private static final String MEDIAN = "median";
+ private static final String MEAN = "mean";
+ private static final String PERCENTILE99TH = "99thPercentile";
+ private final TelemetryClient telemetry;
+ private final String clientName;
+ private final Map rateLimitersMap;
+
+ /**
+ * Default constructor
+ *
+ * @param httpClient http client
+ * @param clientName name of the client
+ * @param url account url
+ */
+ TelemetryService(CloseableHttpClient httpClient, String clientName, String url) {
+ this.clientName = clientName;
+ this.telemetry = (TelemetryClient) TelemetryClient.createSessionlessTelemetry(httpClient, url);
+ this.rateLimitersMap = new HashMap<>();
+ }
+
+ /** Flush the telemetry buffer and close the telemetry service */
+ public void close() {
+ this.telemetry.close();
+ }
+
+ /** Report the Streaming Ingest latency metrics */
+ public void reportLatencyInSec(
+ Timer buildLatency, Timer uploadLatency, Timer registerLatency, Timer flushLatency) {
+ if (flushLatency.getCount() > 0) {
+ ObjectNode msg = MAPPER.createObjectNode();
+ msg.set("build_latency_ms", buildMsgFromTimer(buildLatency));
+ msg.set("upload_latency_ms", buildMsgFromTimer(uploadLatency));
+ msg.set("register_latency_ms", buildMsgFromTimer(registerLatency));
+ msg.set("flush_latency_ms", buildMsgFromTimer(flushLatency));
+ send(TelemetryType.STREAMING_INGEST_LATENCY_IN_SEC, msg);
+ }
+ }
+
+ /** Report the Streaming Ingest failure metrics */
+ public void reportClientFailure(String summary, String exception) {
+ ObjectNode msg = MAPPER.createObjectNode();
+ msg.put("summary", summary);
+ msg.put("client_version", DEFAULT_VERSION);
+ msg.put("exception", exception);
+ send(TelemetryType.STREAMING_INGEST_CLIENT_FAILURE, msg);
+ }
+
+ /** Report the Streaming Ingest throughput metrics */
+ public void reportThroughputBytesPerSecond(Meter inputThroughput, Meter uploadThroughput) {
+ if (inputThroughput.getCount() > 0) {
+ ObjectNode msg = MAPPER.createObjectNode();
+ msg.put(COUNT, inputThroughput.getCount());
+ msg.put("input_mean_rate_bytes_per_sec", inputThroughput.getMeanRate());
+ msg.put("upload_mean_rate_bytes_per_sec", uploadThroughput.getMeanRate());
+ send(TelemetryType.STREAMING_INGEST_THROUGHPUT_BYTES_PER_SEC, msg);
+ }
+ }
+
+ /** Report the Streaming Ingest CUP/memory usage metrics */
+ public void reportCpuMemoryUsage(Histogram cpuUsage) {
+ if (cpuUsage.getCount() > 0) {
+ ObjectNode msg = MAPPER.createObjectNode();
+ Snapshot cpuSnapshot = cpuUsage.getSnapshot();
+ Runtime runTime = Runtime.getRuntime();
+ msg.put(COUNT, cpuUsage.getCount());
+ msg.put("cpu_max", cpuSnapshot.getMax());
+ msg.put("cpu_mean", cpuSnapshot.getMean());
+ msg.put("max_memory", runTime.maxMemory());
+ msg.put("total_memory", runTime.totalMemory());
+ msg.put("free_memory", runTime.freeMemory());
+ send(TelemetryType.STREAMING_INGEST_CPU_MEMORY_USAGE, msg);
+ }
+ }
+
+ /** Report the offset token mismatch in a batch */
+ public void reportBatchOffsetMismatch(
+ String channelName,
+ String prevBatchEndOffset,
+ String startOffset,
+ String endOffset,
+ long rowCount) {
+ // Add a rate limiter to report the mismatch at most once every second per channel
+ RateLimiter rateLimiter =
+ rateLimitersMap.computeIfAbsent(channelName, v -> RateLimiter.create(1.0));
+ if (rateLimiter.tryAcquire()) {
+ ObjectNode msg = MAPPER.createObjectNode();
+ msg.put("channel_name", channelName);
+ msg.put("prev_batch_end_offset", prevBatchEndOffset);
+ msg.put("start_offset", startOffset);
+ msg.put("end_offset", endOffset);
+ msg.put("row_count", rowCount);
+ send(TelemetryType.STREAMING_INGEST_BATCH_OFFSET_MISMATCH, msg);
+ } else {
+ logger.logDebug(
+ "Rate limit exceeded on reportBatchOffsetMismatch, skipping report it to SF.");
+ }
+ }
+
+ /** Send log to Snowflake asynchronously through JDBC client telemetry */
+ void send(TelemetryType type, ObjectNode msg) {
+ try {
+ msg.put(TYPE, type.toString());
+ msg.put(CLIENT_NAME, clientName);
+ telemetry.addLogToBatch(TelemetryUtil.buildJobData(msg));
+ } catch (Exception e) {
+ logger.logWarn("Failed to send telemetry data, error: {}", e.getMessage());
+ }
+ }
+
+ /** Build message from a Timer metric */
+ private ObjectNode buildMsgFromTimer(Timer timer) {
+ ObjectNode msg = MAPPER.createObjectNode();
+ Snapshot buildSnapshot = timer.getSnapshot();
+ msg.put(COUNT, timer.getCount());
+ msg.put(MAX, TimeUnit.NANOSECONDS.toMillis(buildSnapshot.getMax()));
+ msg.put(MIN, TimeUnit.NANOSECONDS.toMillis(buildSnapshot.getMin()));
+ msg.put(MEAN, TimeUnit.NANOSECONDS.toMillis((long) buildSnapshot.getMean()));
+ msg.put(MEDIAN, TimeUnit.NANOSECONDS.toMillis((long) buildSnapshot.getMedian()));
+ msg.put(
+ PERCENTILE99TH, TimeUnit.NANOSECONDS.toMillis((long) buildSnapshot.get99thPercentile()));
+ return msg;
+ }
+
+ /** Refresh JWT token stored in the telemetry client */
+ public void refreshToken(String token) {
+ telemetry.refreshToken(token);
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/DropChannelRequest.java b/src/main/java/net/snowflake/ingest/streaming/DropChannelRequest.java
new file mode 100644
index 000000000..3bc096c16
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/DropChannelRequest.java
@@ -0,0 +1,89 @@
+/*
+ * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming;
+
+import net.snowflake.ingest.utils.Utils;
+
+/** A class that is used to drop a {@link SnowflakeStreamingIngestChannel} */
+public class DropChannelRequest {
+ // Name of the channel
+ private final String channelName;
+
+ // Name of the database that the channel belongs to
+ private final String dbName;
+
+ // Name of the schema that the channel belongs to
+ private final String schemaName;
+
+ // Name of the table that the channel belongs to
+ private final String tableName;
+
+ public static DropChannelRequestBuilder builder(String channelName) {
+ return new DropChannelRequestBuilder(channelName);
+ }
+
+ /** A builder class to build a DropChannelRequest */
+ public static class DropChannelRequestBuilder {
+ private final String channelName;
+ private String dbName;
+ private String schemaName;
+ private String tableName;
+
+ public DropChannelRequestBuilder(String channelName) {
+ this.channelName = channelName;
+ }
+
+ public DropChannelRequestBuilder setDBName(String dbName) {
+ this.dbName = dbName;
+ return this;
+ }
+
+ public DropChannelRequestBuilder setSchemaName(String schemaName) {
+ this.schemaName = schemaName;
+ return this;
+ }
+
+ public DropChannelRequestBuilder setTableName(String tableName) {
+ this.tableName = tableName;
+ return this;
+ }
+
+ public DropChannelRequest build() {
+ return new DropChannelRequest(this);
+ }
+ }
+
+ public DropChannelRequest(DropChannelRequestBuilder builder) {
+ Utils.assertStringNotNullOrEmpty("channel name", builder.channelName);
+ Utils.assertStringNotNullOrEmpty("database name", builder.dbName);
+ Utils.assertStringNotNullOrEmpty("schema name", builder.schemaName);
+ Utils.assertStringNotNullOrEmpty("table name", builder.tableName);
+
+ this.channelName = builder.channelName;
+ this.dbName = builder.dbName;
+ this.schemaName = builder.schemaName;
+ this.tableName = builder.tableName;
+ }
+
+ public String getDBName() {
+ return this.dbName;
+ }
+
+ public String getSchemaName() {
+ return this.schemaName;
+ }
+
+ public String getTableName() {
+ return this.tableName;
+ }
+
+ public String getChannelName() {
+ return this.channelName;
+ }
+
+ public String getFullyQualifiedTableName() {
+ return String.format("%s.%s.%s", this.dbName, this.schemaName, this.tableName);
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/InsertValidationResponse.java b/src/main/java/net/snowflake/ingest/streaming/InsertValidationResponse.java
index f64260864..f57d384c9 100644
--- a/src/main/java/net/snowflake/ingest/streaming/InsertValidationResponse.java
+++ b/src/main/java/net/snowflake/ingest/streaming/InsertValidationResponse.java
@@ -47,16 +47,21 @@ public void addError(InsertError error) {
/** Wraps the row content and exception when there is a failure */
public static class InsertError {
private final Object rowContent;
- private final SFException exception;
+ private SFException exception;
// Used to map this error row with original row in the insertRows Iterable.
// i.e the rowIndex can be index 9 in the list of 10 rows.
// index is 0 based so as to match with incoming Iterable
- private final long rowIndex;
+ private long rowIndex;
- public InsertError(Object row, SFException exception, long rowIndex) {
+ // List of extra column names in the input row compared with the table schema
+ private List extraColNames;
+
+ // List of missing non-nullable column names in the input row compared with the table schema
+ private List missingNotNullColNames;
+
+ public InsertError(Object row, long rowIndex) {
this.rowContent = row;
- this.exception = exception;
this.rowIndex = rowIndex;
}
@@ -70,11 +75,29 @@ public String getMessage() {
return this.exception.getMessage();
}
+ /**
+ * Set the insert exception
+ *
+ * @param exception exception encountered during the insert
+ */
+ public void setException(SFException exception) {
+ this.exception = exception;
+ }
+
/** Get the exception */
public SFException getException() {
return this.exception;
}
+ /**
+ * Set the row index
+ *
+ * @param rowIndex the corresponding row index in the original input row list
+ */
+ public void setRowIndex(long rowIndex) {
+ this.rowIndex = rowIndex;
+ }
+
/**
* Get the rowIndex. Please note, this index is 0 based so it can be used in fetching nth row
* from the input. ({@link SnowflakeStreamingIngestChannel#insertRows(Iterable, String)})
@@ -82,5 +105,28 @@ public SFException getException() {
public long getRowIndex() {
return rowIndex;
}
+
+ /** Set the extra column names in the input row compared with the table schema */
+ public void setExtraColNames(List extraColNames) {
+ this.extraColNames = extraColNames;
+ }
+
+ /** Get the list of extra column names in the input row compared with the table schema */
+ public List getExtraColNames() {
+ return extraColNames;
+ }
+
+ /** Set the missing non-nullable column names in the input row compared with the table schema */
+ public void setMissingNotNullColNames(List missingNotNullColNames) {
+ this.missingNotNullColNames = missingNotNullColNames;
+ }
+
+ /**
+ * Get the list of missing non-nullable column names in the input row compared with the table
+ * schema
+ */
+ public List getMissingNotNullColNames() {
+ return missingNotNullColNames;
+ }
}
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/OffsetTokenVerificationFunction.java b/src/main/java/net/snowflake/ingest/streaming/OffsetTokenVerificationFunction.java
new file mode 100644
index 000000000..847e56501
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/OffsetTokenVerificationFunction.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming;
+
+/**
+ * Interface to provide a custom offset verification logic. If specified, verification failures will
+ * be logged as warnings and reported to Snowflake. This interface could be used when there are
+ * certain assumption about the offset token behavior. Please reach out to Snowflake if you notice
+ * any unexpected behaviors.
+ *
+ * Below is an example that verifies that all offset tokens are monotonically increasing numbers:
+ *
+ *
+ *
+ * private static final OffsetTokenVerificationFunction offsetTokenVerificationFunction =
+ * (prevBatchEndOffset, curBatchStartOffset, curBatchEndOffset, rowCount) -> {
+ * boolean isMatch = true;
+ *
+ * if (curBatchStartOffset != null) {
+ * try {
+ * long curStart = Long.parseLong(curBatchStartOffset);
+ * long curEnd = Long.parseLong(curBatchEndOffset);
+ *
+ * // We verify that the end_offset - start_offset + 1 = row_count
+ * if (curEnd - curStart + 1 != rowCount) {
+ * isMatch = false;
+ * }
+ *
+ * // We verify that start_offset_of_current_batch = end_offset_of_previous_batch+1
+ * if (prevBatchEndOffset != null) {
+ * long prevEnd = Long.parseLong(prevBatchEndOffset);
+ * if (curStart != prevEnd + 1) {
+ * isMatch = false;
+ * }
+ * }
+ * } catch (NumberFormatException ignored) {
+ * // Do nothing since we can't compare the offset, or report a mismatch if number is expected
+ * }
+ * }
+ *
+ * return isMatch;
+ * };
+ *
+ *
+ */
+public interface OffsetTokenVerificationFunction {
+ /**
+ * @param prevBatchEndOffset end offset token of the previous batch
+ * @param curBatchStartOffset start offset token of the current batch
+ * @param curBatchEndOffset end offset token of the current batch
+ * @param rowCount number of rows in the current batch
+ * @return a boolean indicates whether the verification passed or not, if not, we will log a
+ * warning and report it to SF
+ */
+ boolean verify(
+ String prevBatchEndOffset,
+ String curBatchStartOffset,
+ String curBatchEndOffset,
+ long rowCount);
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java b/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java
index 0370771e1..cc8782dbd 100644
--- a/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java
+++ b/src/main/java/net/snowflake/ingest/streaming/OpenChannelRequest.java
@@ -4,6 +4,7 @@
package net.snowflake.ingest.streaming;
+import java.time.ZoneId;
import net.snowflake.ingest.utils.Utils;
/** A class that is used to open/create a {@link SnowflakeStreamingIngestChannel} */
@@ -11,8 +12,16 @@ public class OpenChannelRequest {
public enum OnErrorOption {
CONTINUE, // CONTINUE loading the rows, and return all the errors in the response
ABORT, // ABORT the entire batch, and throw an exception when we hit the first error
+ SKIP_BATCH, // If an error in the batch is detected return a response containing all error row
+ // indexes. No data is ingested
}
+ /**
+ * Default value of the timezone, which will be used for TIMESTAMP_LTZ and TIMESTAMP_TZ column
+ * types when the user input does not have any timezone information.
+ */
+ private static final ZoneId DEFAULT_DEFAULT_TIMEZONE = ZoneId.of("America/Los_Angeles");
+
// Name of the channel
private final String channelName;
@@ -28,6 +37,14 @@ public enum OnErrorOption {
// On_error option on this channel
private final OnErrorOption onErrorOption;
+ // Default timezone for TIMESTAMP_LTZ and TIMESTAMP_TZ columns
+ private final ZoneId defaultTimezone;
+
+ private final String offsetToken;
+ private final boolean isOffsetTokenProvided;
+
+ private final OffsetTokenVerificationFunction offsetTokenVerificationFunction;
+
public static OpenChannelRequestBuilder builder(String channelName) {
return new OpenChannelRequestBuilder(channelName);
}
@@ -39,9 +56,16 @@ public static class OpenChannelRequestBuilder {
private String schemaName;
private String tableName;
private OnErrorOption onErrorOption;
+ private ZoneId defaultTimezone;
+
+ private String offsetToken;
+ private boolean isOffsetTokenProvided = false;
+
+ private OffsetTokenVerificationFunction offsetTokenVerificationFunction;
public OpenChannelRequestBuilder(String channelName) {
this.channelName = channelName;
+ this.defaultTimezone = DEFAULT_DEFAULT_TIMEZONE;
}
public OpenChannelRequestBuilder setDBName(String dbName) {
@@ -64,6 +88,23 @@ public OpenChannelRequestBuilder setOnErrorOption(OnErrorOption onErrorOption) {
return this;
}
+ public OpenChannelRequestBuilder setDefaultTimezone(ZoneId defaultTimezone) {
+ this.defaultTimezone = defaultTimezone;
+ return this;
+ }
+
+ public OpenChannelRequestBuilder setOffsetToken(String offsetToken) {
+ this.offsetToken = offsetToken;
+ this.isOffsetTokenProvided = true;
+ return this;
+ }
+
+ public OpenChannelRequestBuilder setOffsetTokenVerificationFunction(
+ OffsetTokenVerificationFunction function) {
+ this.offsetTokenVerificationFunction = function;
+ return this;
+ }
+
public OpenChannelRequest build() {
return new OpenChannelRequest(this);
}
@@ -75,12 +116,17 @@ private OpenChannelRequest(OpenChannelRequestBuilder builder) {
Utils.assertStringNotNullOrEmpty("schema name", builder.schemaName);
Utils.assertStringNotNullOrEmpty("table name", builder.tableName);
Utils.assertNotNull("on_error option", builder.onErrorOption);
+ Utils.assertNotNull("default_timezone", builder.defaultTimezone);
this.channelName = builder.channelName;
this.dbName = builder.dbName;
this.schemaName = builder.schemaName;
this.tableName = builder.tableName;
this.onErrorOption = builder.onErrorOption;
+ this.defaultTimezone = builder.defaultTimezone;
+ this.offsetToken = builder.offsetToken;
+ this.isOffsetTokenProvided = builder.isOffsetTokenProvided;
+ this.offsetTokenVerificationFunction = builder.offsetTokenVerificationFunction;
}
public String getDBName() {
@@ -99,6 +145,10 @@ public String getChannelName() {
return this.channelName;
}
+ public ZoneId getDefaultTimezone() {
+ return this.defaultTimezone;
+ }
+
public String getFullyQualifiedTableName() {
return String.format("%s.%s.%s", this.dbName, this.schemaName, this.tableName);
}
@@ -106,4 +156,16 @@ public String getFullyQualifiedTableName() {
public OnErrorOption getOnErrorOption() {
return this.onErrorOption;
}
+
+ 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/SnowflakeStreamingIngestChannel.java b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestChannel.java
index e1f0fdbc1..7ffbf2549 100644
--- a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestChannel.java
+++ b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestChannel.java
@@ -4,17 +4,23 @@
package net.snowflake.ingest.streaming;
+import java.time.ZoneId;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import javax.annotation.Nullable;
+import net.snowflake.ingest.streaming.internal.ColumnProperties;
/**
* A logical partition that represents a connection to a single Snowflake table, data will be
- * ingested into the channel, and then flush to Snowflake table periodically in the background. Note
- * that only one client (or thread) could write to a channel at a given time, so if there are
- * multiple clients (or multiple threads in the same client) try to ingest using the same channel,
- * the latest client (or thread) that opens the channel will win and all the other opened channels
- * will be invalid
+ * ingested into the channel, and then flushed to Snowflake table periodically in the background.
+ *
+ * Channels are identified by their name and only one channel with the same name may ingest data
+ * at the same time. When a new channel is opened, all previously opened channels with the same name
+ * are invalidated (this applies for the table globally. not just in a single JVM). In order to
+ * ingest data from multiple threads/clients/applications, we recommend opening multiple channels,
+ * each with a different name.
+ *
+ *
Thread safety note: Implementations of this interface are required to be thread safe.
*/
public interface SnowflakeStreamingIngestChannel {
/**
@@ -67,17 +73,174 @@ public interface SnowflakeStreamingIngestChannel {
boolean isClosed();
/**
- * Close the channel, this will make sure all the data in this channel is committed before closing
+ * Close the channel, this function will make sure all the data in this channel is committed
*
* @return a completable future which will be completed when the channel is closed
*/
CompletableFuture close();
+ /**
+ * Close the channel, this function will make sure all the data in this channel is committed
+ *
+ * Note that this call with drop=true will delete Offset
+ * Token and other state from Snowflake servers unless the channel has already been opened in
+ * another client. So only use it if you are completely done ingesting data for this channel. If
+ * you open a channel with the same name in the future, it will behave like a new channel.
+ *
+ * @param drop if true, the channel will be dropped after all data is successfully committed.
+ * @return a completable future which will be completed when the channel is closed
+ */
+ CompletableFuture close(boolean drop);
+
/**
* Insert one row into the channel, the row is represented using Map where the key is column name
- * and the value is a row of data
+ * and the value is a row of data. The following table summarizes supported value types and their
+ * formats:
+ *
+ *
+ *
+ * | Snowflake Column Type |
+ * Allowed Java Data Type |
+ *
+ *
+ * | CHAR, VARCHAR |
+ *
+ *
+ * - String
+ * - primitive data types (int, boolean, char, …)
+ *
+ * |
+ *
+ *
+ * | BINARY |
+ *
+ *
+ * - byte[]
+ * - String (hex-encoded)
+ *
+ * |
+ *
+ *
+ * | NUMBER, FLOAT |
+ *
+ *
+ * - numeric types (BigInteger, BigDecimal, byte, int, double, …)
+ * - String
+ *
+ * |
+ *
+ *
+ * | BOOLEAN |
+ *
+ *
+ * - boolean
+ * - numeric types (BigInteger, BigDecimal, byte, int, double, …)
+ * - String
+ *
+ * See boolean conversion details.
+ * |
+ *
+ *
+ * | TIME |
+ *
+ *
+ * - {@link java.time.LocalTime}
+ * - {@link java.time.OffsetTime}
+ * -
+ * String (in one of the following formats):
+ *
+ * - {@link java.time.format.DateTimeFormatter#ISO_LOCAL_TIME}
+ * - {@link java.time.format.DateTimeFormatter#ISO_OFFSET_TIME}
+ * - Integer-stored time (see Snowflake Docs for more details)
+ *
+ *
+ *
+ *
+ * |
+ *
+ *
+ * | DATE |
+ *
+ *
+ * - {@link java.time.LocalDate}
+ * - {@link java.time.LocalDateTime}
+ * - {@link java.time.OffsetDateTime}
+ * - {@link java.time.ZonedDateTime}
+ * - {@link java.time.Instant}
+ * -
+ * String (in one of the following formats):
+ *
+ * - {@link java.time.format.DateTimeFormatter#ISO_LOCAL_DATE}
+ * - {@link java.time.format.DateTimeFormatter#ISO_LOCAL_DATE_TIME}
+ * - {@link java.time.format.DateTimeFormatter#ISO_OFFSET_DATE_TIME}
+ * - {@link java.time.format.DateTimeFormatter#ISO_ZONED_DATE_TIME}
+ * - Integer-stored date (see Snowflake Docs for more details)
+ *
+ *
+ *
+ * |
+ *
+ *
+ * | TIMESTAMP_NTZ, TIMESTAMP_LTZ, TIMESTAMP_TZ |
+ *
+ *
+ * - {@link java.time.LocalDate}
+ * - {@link java.time.LocalDateTime}
+ * - {@link java.time.OffsetDateTime}
+ * - {@link java.time.ZonedDateTime}
+ * - {@link java.time.Instant}
+ * -
+ * String (in one of the following formats):
+ *
+ * - {@link java.time.format.DateTimeFormatter#ISO_LOCAL_DATE}
+ * - {@link java.time.format.DateTimeFormatter#ISO_LOCAL_DATE_TIME}
+ * - {@link java.time.format.DateTimeFormatter#ISO_OFFSET_DATE_TIME}
+ * - {@link java.time.format.DateTimeFormatter#ISO_ZONED_DATE_TIME}
+ * - Integer-stored timestamp (see Snowflake Docs for more details)
+ *
+ *
+ *
+ *
+ *
+ * For TIMESTAMP_LTZ and TIMESTAMP_TZ, all input without timezone will be by default interpreted in the timezone "America/Los_Angeles". This can be changed by calling {@link net.snowflake.ingest.streaming.OpenChannelRequest.OpenChannelRequestBuilder#setDefaultTimezone(ZoneId)}.
+ * |
+ *
+ * | VARIANT, ARRAY |
+ *
+ *
+ * - String (must be a valid JSON value)
+ * - primitive data types and their arrays
+ * - BigInteger, BigDecimal
+ * - {@link java.time.LocalDate}
+ * - {@link java.time.LocalDateTime}
+ * - {@link java.time.OffsetDateTime}
+ * - {@link java.time.ZonedDateTime}
+ * - Map<String, T> where T is a valid VARIANT type
+ * - T[] where T is a valid VARIANT type
+ * - List<T> where T is a valid VARIANT type
+ *
+ * |
+ *
+ *
+ * | OBJECT |
+ *
+ *
+ * - String (must be a valid JSON object)
+ * - Map<String, T> where T is a valid variant type
+ *
+ * |
+ *
+ *
+ * | GEOGRAPHY, GEOMETRY |
+ * Not supported |
+ *
+ *
*
- * @param row object data to write
+ *
+ *
+ * @param row object data to write. For predictable results, we recommend not to concurrently
+ * modify the input row data.
* @param offsetToken offset of given row, used for replay in case of failures. It could be null
* if you don't plan on replaying or can't replay
* @return insert response that possibly contains errors because of insertion failures
@@ -86,13 +249,25 @@ public interface SnowflakeStreamingIngestChannel {
/**
* Insert a batch of rows into the channel, each row is represented using Map where the key is
- * column name and the value is a row of data
+ * column name and the value is a row of data. See {@link
+ * SnowflakeStreamingIngestChannel#insertRow(Map, String)} for more information about accepted
+ * values.
*
* @param rows object data to write
- * @param offsetToken offset of last row in the row-set, used for replay in case of failures, It
- * could be null if you don't plan on replaying or can't replay
+ * @param startOffsetToken start offset of the batch/row-set
+ * @param endOffsetToken end offset of the batch/row-set, used for replay in case of failures, *
+ * It could be null if you don't plan on replaying or can't replay
* @return insert response that possibly contains errors because of insertion failures
*/
+ InsertValidationResponse insertRows(
+ Iterable> rows,
+ @Nullable String startOffsetToken,
+ @Nullable String endOffsetToken);
+
+ /**
+ * Insert a batch of rows into the channel with the end offset token only, please see {@link
+ * SnowflakeStreamingIngestChannel#insertRows(Iterable, String, String)} for more information.
+ */
InsertValidationResponse insertRows(
Iterable> rows, @Nullable String offsetToken);
@@ -101,5 +276,15 @@ InsertValidationResponse insertRows(
*
* @return the latest committed offset token
*/
+ @Nullable
String getLatestCommittedOffsetToken();
+
+ /**
+ * Gets the table schema associated with this channel. Note that this is the table schema at the
+ * time of a channel open event. The schema may be changed on the Snowflake side in which case
+ * this will continue to show an old schema version until the channel is re-opened.
+ *
+ * @return map representing Column Name to Column Properties
+ */
+ Map getTableSchema();
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClient.java b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClient.java
index 9dcc3b716..2de8a45c2 100644
--- a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClient.java
+++ b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClient.java
@@ -4,11 +4,16 @@
package net.snowflake.ingest.streaming;
+import java.util.List;
+import java.util.Map;
+
/**
* A class that is the starting point for using the Streaming Ingest client APIs, a single client
* maps to exactly one account in Snowflake; however, multiple clients can point to the same
* account. Each client will contain information for Snowflake authentication and authorization, and
* it will be used to create one or more {@link SnowflakeStreamingIngestChannel}
+ *
+ * Thread safety note: Implementations of this interface are required to be thread safe.
*/
public interface SnowflakeStreamingIngestClient extends AutoCloseable {
@@ -20,6 +25,24 @@ public interface SnowflakeStreamingIngestClient extends AutoCloseable {
*/
SnowflakeStreamingIngestChannel openChannel(OpenChannelRequest request);
+ /**
+ * Drop the specified channel on the server using a {@link DropChannelRequest}
+ *
+ *
Note that this call will blindly drop the latest version of the channel and any
+ * pending data will be lost. It will also delete Offset
+ * Token and other state from Snowflake servers. So only use it if you are completely done
+ * ingesting data for this channel. If you open a channel with the same name in the future, it
+ * will behave like a new channel.
+ *
+ *
Also see {@link SnowflakeStreamingIngestChannel#close(boolean)} to drop channels on close.
+ * That approach will drop the local version of the channel and if the channel has been
+ * concurrently reopened by another client, that version of the channel won't be affected.
+ *
+ * @param request the drop channel request
+ */
+ void dropChannel(DropChannelRequest request);
+
/**
* Get the client name
*
@@ -27,6 +50,28 @@ public interface SnowflakeStreamingIngestClient extends AutoCloseable {
*/
String getName();
- /** @return a boolean to indicate whether the client is closed */
+ /**
+ * 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.
+ *
+ * @param refreshToken the new refresh token
+ */
+ void setRefreshToken(String refreshToken);
+
+ /**
+ * Check whether the client is closed or not, if you want to make sure all data are committed
+ * before closing, please call {@link SnowflakeStreamingIngestClient#close()} before closing the
+ * entire client
+ *
+ * @return a boolean to indicate whether the client is closed
+ */
boolean isClosed();
+
+ /**
+ * Return the latest committed offset token for a list of channels
+ *
+ * @return a map of channel fully qualified name to latest committed offset token
+ */
+ Map getLatestCommittedOffsetTokens(
+ List channels);
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java
index d99f54e56..cd6d78787 100644
--- a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java
+++ b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java
@@ -8,8 +8,6 @@
import java.util.Properties;
import net.snowflake.ingest.streaming.internal.SnowflakeStreamingIngestClientInternal;
import net.snowflake.ingest.utils.Constants;
-import net.snowflake.ingest.utils.ErrorCode;
-import net.snowflake.ingest.utils.SFException;
import net.snowflake.ingest.utils.SnowflakeURL;
import net.snowflake.ingest.utils.Utils;
@@ -30,6 +28,9 @@ public static class Builder {
// Allows client to override some default parameter values
private Map parameterOverrides;
+ // Indicates whether it's under test mode
+ private boolean isTestMode;
+
private Builder(String name) {
this.name = name;
}
@@ -44,37 +45,20 @@ public Builder setParameterOverrides(Map parameterOverrides) {
return this;
}
+ public Builder setIsTestMode(boolean isTestMode) {
+ this.isTestMode = isTestMode;
+ return this;
+ }
+
public SnowflakeStreamingIngestClient build() {
Utils.assertStringNotNullOrEmpty("client name", this.name);
Utils.assertNotNull("connection properties", this.prop);
- if (!this.prop.containsKey(Constants.ACCOUNT_URL)) {
- if (!this.prop.containsKey(Constants.HOST)) {
- throw new SFException(ErrorCode.MISSING_CONFIG, "host");
- }
- if (!this.prop.containsKey(Constants.SCHEME)) {
- throw new SFException(ErrorCode.MISSING_CONFIG, "scheme");
- }
- if (!this.prop.containsKey(Constants.PORT)) {
- throw new SFException(ErrorCode.MISSING_CONFIG, "port");
- }
-
- this.prop.put(
- Constants.ACCOUNT_URL,
- Utils.constructAccountUrl(
- this.prop.get(Constants.SCHEME).toString(),
- this.prop.get(Constants.HOST).toString(),
- Integer.parseInt(prop.get(Constants.PORT).toString())));
- }
-
- if (!this.prop.containsKey(Constants.ROLE)) {
- throw new SFException(ErrorCode.MISSING_CONFIG, "role");
- }
+ Properties prop = Utils.createProperties(this.prop);
+ SnowflakeURL accountURL = new SnowflakeURL(prop.getProperty(Constants.ACCOUNT_URL));
- SnowflakeURL accountURL = new SnowflakeURL(this.prop.getProperty(Constants.ACCOUNT_URL));
- Properties prop = Utils.createProperties(this.prop, accountURL.sslEnabled());
- return new SnowflakeStreamingIngestClientInternal(
- this.name, accountURL, prop, parameterOverrides);
+ return new SnowflakeStreamingIngestClientInternal<>(
+ this.name, accountURL, prop, this.parameterOverrides, this.isTestMode);
}
}
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/example/SnowflakeStreamingIngestExample.java b/src/main/java/net/snowflake/ingest/streaming/example/SnowflakeStreamingIngestExample.java
index 85409ce20..ceb8aaa1d 100644
--- a/src/main/java/net/snowflake/ingest/streaming/example/SnowflakeStreamingIngestExample.java
+++ b/src/main/java/net/snowflake/ingest/streaming/example/SnowflakeStreamingIngestExample.java
@@ -4,11 +4,12 @@
package net.snowflake.ingest.streaming.example;
+import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.node.ObjectNode;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.HashMap;
+import java.util.Iterator;
import java.util.Map;
import java.util.Properties;
import net.snowflake.ingest.streaming.InsertValidationResponse;
@@ -16,43 +17,83 @@
import net.snowflake.ingest.streaming.SnowflakeStreamingIngestChannel;
import net.snowflake.ingest.streaming.SnowflakeStreamingIngestClient;
import net.snowflake.ingest.streaming.SnowflakeStreamingIngestClientFactory;
-import net.snowflake.ingest.utils.Utils;
-/** Examples on how to use the Streaming Ingest client APIs */
+/**
+ * Example on how to use the Streaming Ingest client APIs.
+ *
+ * Please read the README.md file for detailed steps
+ */
public class SnowflakeStreamingIngestExample {
+ // Please follow the example in profile_streaming.json.example to see the required properties, or
+ // if you have already set up profile.json with Snowpipe before, all you need is to add the "role"
+ // property. If the "role" is not specified, the default user role will be applied.
private static String PROFILE_PATH = "profile.json";
private static final ObjectMapper mapper = new ObjectMapper();
public static void main(String[] args) throws Exception {
- ObjectNode profile =
- (ObjectNode) mapper.readTree(new String(Files.readAllBytes(Paths.get(PROFILE_PATH))));
- Properties prop = Utils.getPropertiesFromJson(profile);
+ Properties props = new Properties();
+ Iterator> propIt =
+ mapper.readTree(new String(Files.readAllBytes(Paths.get(PROFILE_PATH)))).fields();
+ while (propIt.hasNext()) {
+ Map.Entry prop = propIt.next();
+ props.put(prop.getKey(), prop.getValue().asText());
+ }
// Create a streaming ingest client
try (SnowflakeStreamingIngestClient client =
- SnowflakeStreamingIngestClientFactory.builder("CLIENT").setProperties(prop).build()) {
- // Create an open channel request on table T_STREAMINGINGEST
+ SnowflakeStreamingIngestClientFactory.builder("MY_CLIENT").setProperties(props).build()) {
+
+ // Create an open channel request on table MY_TABLE, note that the corresponding
+ // db/schema/table needs to be present
+ // Example: create or replace table MY_TABLE(c1 number);
OpenChannelRequest request1 =
- OpenChannelRequest.builder("CHANNEL")
- .setDBName("DB_STREAMINGINGEST")
- .setSchemaName("SCHEMA_STREAMINGINGEST")
- .setTableName("T_STREAMINGINGEST")
- .setOnErrorOption(OpenChannelRequest.OnErrorOption.CONTINUE)
+ OpenChannelRequest.builder("MY_CHANNEL")
+ .setDBName("MY_DATABASE")
+ .setSchemaName("MY_SCHEMA")
+ .setTableName("MY_TABLE")
+ .setOnErrorOption(
+ OpenChannelRequest.OnErrorOption.CONTINUE) // Another ON_ERROR option is ABORT
.build();
// Open a streaming ingest channel from the given client
SnowflakeStreamingIngestChannel channel1 = client.openChannel(request1);
- // Insert a few rows into the channel,
- for (int val = 0; val < 1000; val++) {
+ // Insert rows into the channel (Using insertRows API)
+ final int totalRowsInTable = 1000;
+ for (int val = 0; val < totalRowsInTable; val++) {
Map row = new HashMap<>();
+
+ // c1 corresponds to the column name in table
row.put("c1", val);
- InsertValidationResponse response = channel1.insertRow(row, null);
+
+ // Insert the row with the current offset_token
+ InsertValidationResponse response = channel1.insertRow(row, String.valueOf(val));
if (response.hasErrors()) {
- // Simply throw if there is an exception
+ // Simply throw if there is an exception, or you can do whatever you want with the
+ // erroneous row
throw response.getInsertErrors().get(0).getException();
}
}
+
+ // If needed, you can check the offset_token registered in Snowflake to make sure everything
+ // is committed
+ final int expectedOffsetTokenInSnowflake = totalRowsInTable - 1; // 0 based offset_token
+ final int maxRetries = 10;
+ int retryCount = 0;
+
+ do {
+ String offsetTokenFromSnowflake = channel1.getLatestCommittedOffsetToken();
+ if (offsetTokenFromSnowflake != null
+ && offsetTokenFromSnowflake.equals(String.valueOf(expectedOffsetTokenInSnowflake))) {
+ System.out.println("SUCCESSFULLY inserted " + totalRowsInTable + " rows");
+ break;
+ }
+ retryCount++;
+ } while (retryCount < maxRetries);
+
+ // Close the channel, the function internally will make sure everything is committed (or throw
+ // an exception if there is any issue)
+ channel1.close().get();
}
}
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/AbstractRowBuffer.java b/src/main/java/net/snowflake/ingest/streaming/internal/AbstractRowBuffer.java
new file mode 100644
index 000000000..4a13f5adb
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/AbstractRowBuffer.java
@@ -0,0 +1,735 @@
+/*
+ * Copyright (c) 2022 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming.internal;
+
+import com.google.common.annotations.VisibleForTesting;
+import java.time.ZoneId;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import net.snowflake.ingest.connection.TelemetryService;
+import net.snowflake.ingest.streaming.InsertValidationResponse;
+import net.snowflake.ingest.streaming.OffsetTokenVerificationFunction;
+import net.snowflake.ingest.streaming.OpenChannelRequest;
+import net.snowflake.ingest.utils.Constants;
+import net.snowflake.ingest.utils.ErrorCode;
+import net.snowflake.ingest.utils.Logging;
+import net.snowflake.ingest.utils.Pair;
+import net.snowflake.ingest.utils.SFException;
+
+/**
+ * The abstract implementation of the buffer in the Streaming Ingest channel that holds the
+ * un-flushed rows, these rows will be converted to the underlying format implementation for faster
+ * processing
+ *
+ * @param type of column data ({@link ParquetChunkData} for Parquet)
+ */
+abstract class AbstractRowBuffer implements RowBuffer {
+ private static final Logging logger = new Logging(AbstractRowBuffer.class);
+ // these types cannot be packed into the data chunk because they are not readable by the server
+ // side scanner
+ private static final int INVALID_SERVER_SIDE_DATA_TYPE_ORDINAL = -1;
+
+ // The maximum recommended batch size of data passed into a single insertRows() call
+ private static final int INSERT_ROWS_RECOMMENDED_MAX_BATCH_SIZE_IN_BYTES = 16 * 1024 * 1024;
+
+ // Snowflake table column logical type
+ enum ColumnLogicalType {
+ ANY,
+ BOOLEAN(1),
+ ROWINDEX,
+ NULL(15),
+ REAL(8),
+ FIXED(2),
+ TEXT(9),
+ CHAR,
+ BINARY(10),
+ DATE(7),
+ TIME(6),
+ TIMESTAMP_LTZ(3),
+ TIMESTAMP_NTZ(4),
+ TIMESTAMP_TZ(5),
+ INTERVAL,
+ RAW,
+ ARRAY(13, true),
+ OBJECT(12, true),
+ VARIANT(11, true),
+ ROW,
+ SEQUENCE,
+ FUNCTION,
+ USER_DEFINED_TYPE,
+ ;
+
+ // ordinal should be in sync with the server side scanner
+ private final int ordinal;
+ // whether it is a composite data type: array, object or variant
+ private final boolean object;
+
+ ColumnLogicalType() {
+ // no valid server side ordinal by default
+ this(INVALID_SERVER_SIDE_DATA_TYPE_ORDINAL);
+ }
+
+ ColumnLogicalType(int ordinal) {
+ this(ordinal, false);
+ }
+
+ ColumnLogicalType(int ordinal, boolean object) {
+ this.ordinal = ordinal;
+ this.object = object;
+ }
+
+ /**
+ * Ordinal to encode the data type for the server side scanner
+ *
+ * currently used for Parquet format
+ */
+ public int getOrdinal() {
+ return ordinal;
+ }
+
+ /** Whether the data type is a composite type: OBJECT, VARIANT, ARRAY. */
+ public boolean isObject() {
+ return object;
+ }
+ }
+
+ // Snowflake table column physical type
+ enum ColumnPhysicalType {
+ ROWINDEX(9),
+ DOUBLE(7),
+ SB1(1),
+ SB2(2),
+ SB4(3),
+ SB8(4),
+ SB16(5),
+ LOB(8),
+ BINARY,
+ ROW(10),
+ ;
+
+ // ordinal should be in sync with the server side scanner
+ private final int ordinal;
+
+ ColumnPhysicalType() {
+ // no valid server side ordinal by default
+ this(INVALID_SERVER_SIDE_DATA_TYPE_ORDINAL);
+ }
+
+ ColumnPhysicalType(int ordinal) {
+ this.ordinal = ordinal;
+ }
+
+ /**
+ * Ordinal to encode the data type for the server side scanner
+ *
+ *
currently used for Parquet format
+ */
+ public int getOrdinal() {
+ return ordinal;
+ }
+ }
+
+ /** Insert rows function strategy for ON_ERROR=CONTINUE */
+ public class ContinueIngestionStrategy implements IngestionStrategy {
+ @Override
+ public InsertValidationResponse insertRows(
+ AbstractRowBuffer rowBuffer,
+ Iterable> rows,
+ String startOffsetToken,
+ String endOffsetToken) {
+ InsertValidationResponse response = new InsertValidationResponse();
+ float rowsSizeInBytes = 0F;
+ int rowIndex = 0;
+ int prevRowCount = rowBuffer.bufferedRowCount;
+ for (Map row : rows) {
+ InsertValidationResponse.InsertError error =
+ new InsertValidationResponse.InsertError(row, rowIndex);
+ try {
+ if (rowBuffer.bufferedRowCount == Integer.MAX_VALUE) {
+ throw new SFException(ErrorCode.INTERNAL_ERROR, "Row count reaches MAX value");
+ }
+ Set inputColumnNames = verifyInputColumns(row, error, rowIndex);
+ rowsSizeInBytes +=
+ addRow(
+ row, rowBuffer.bufferedRowCount, rowBuffer.statsMap, inputColumnNames, rowIndex);
+ rowBuffer.bufferedRowCount++;
+ } catch (SFException e) {
+ error.setException(e);
+ response.addError(error);
+ } catch (Throwable e) {
+ logger.logWarn("Unexpected error happens during insertRows: {}", e);
+ error.setException(new SFException(e, ErrorCode.INTERNAL_ERROR, e.getMessage()));
+ response.addError(error);
+ }
+ rowIndex++;
+ }
+ checkBatchSizeRecommendedMaximum(rowsSizeInBytes);
+ checkOffsetMismatch(
+ rowBuffer.channelState.getEndOffsetToken(), startOffsetToken, endOffsetToken, rowIndex);
+ rowBuffer.channelState.updateOffsetToken(startOffsetToken, endOffsetToken, prevRowCount);
+ rowBuffer.bufferSize += rowsSizeInBytes;
+ rowBuffer.rowSizeMetric.accept(rowsSizeInBytes);
+ return response;
+ }
+ }
+
+ /** Insert rows function strategy for ON_ERROR=ABORT */
+ public class AbortIngestionStrategy implements IngestionStrategy {
+ @Override
+ public InsertValidationResponse insertRows(
+ AbstractRowBuffer rowBuffer,
+ Iterable> rows,
+ String startOffsetToken,
+ String endOffsetToken) {
+ // If the on_error option is ABORT, simply throw the first exception
+ InsertValidationResponse response = new InsertValidationResponse();
+ float rowsSizeInBytes = 0F;
+ float tempRowsSizeInBytes = 0F;
+ int tempRowCount = 0;
+ for (Map row : rows) {
+ Set inputColumnNames = verifyInputColumns(row, null, tempRowCount);
+ tempRowsSizeInBytes +=
+ addTempRow(row, tempRowCount, rowBuffer.tempStatsMap, inputColumnNames, tempRowCount);
+ tempRowCount++;
+ checkBatchSizeEnforcedMaximum(tempRowsSizeInBytes);
+ if ((long) rowBuffer.bufferedRowCount + tempRowCount >= Integer.MAX_VALUE) {
+ throw new SFException(ErrorCode.INTERNAL_ERROR, "Row count reaches MAX value");
+ }
+ }
+ checkBatchSizeRecommendedMaximum(tempRowsSizeInBytes);
+
+ moveTempRowsToActualBuffer(tempRowCount);
+
+ rowsSizeInBytes = tempRowsSizeInBytes;
+ rowBuffer.statsMap.forEach(
+ (colName, stats) ->
+ rowBuffer.statsMap.put(
+ colName,
+ RowBufferStats.getCombinedStats(stats, rowBuffer.tempStatsMap.get(colName))));
+ checkOffsetMismatch(
+ rowBuffer.channelState.getEndOffsetToken(),
+ startOffsetToken,
+ endOffsetToken,
+ tempRowCount);
+ rowBuffer.channelState.updateOffsetToken(
+ startOffsetToken, endOffsetToken, rowBuffer.bufferedRowCount);
+ rowBuffer.bufferedRowCount += tempRowCount;
+ rowBuffer.bufferSize += rowsSizeInBytes;
+ rowBuffer.rowSizeMetric.accept(rowsSizeInBytes);
+ return response;
+ }
+ }
+
+ /** Insert rows function strategy for ON_ERROR=SKIP_BATCH */
+ public class SkipBatchIngestionStrategy implements IngestionStrategy {
+ @Override
+ public InsertValidationResponse insertRows(
+ AbstractRowBuffer rowBuffer,
+ Iterable> rows,
+ String startOffsetToken,
+ String endOffsetToken) {
+ InsertValidationResponse response = new InsertValidationResponse();
+ float rowsSizeInBytes = 0F;
+ float tempRowsSizeInBytes = 0F;
+ int tempRowCount = 0;
+ int rowIndex = 0;
+ for (Map row : rows) {
+ InsertValidationResponse.InsertError error =
+ new InsertValidationResponse.InsertError(row, rowIndex);
+ try {
+ Set inputColumnNames = verifyInputColumns(row, error, rowIndex);
+ tempRowsSizeInBytes +=
+ addTempRow(row, tempRowCount, rowBuffer.tempStatsMap, inputColumnNames, rowIndex);
+ tempRowCount++;
+ } catch (SFException e) {
+ error.setException(e);
+ response.addError(error);
+ } catch (Throwable e) {
+ logger.logWarn("Unexpected error happens during insertRows: {}", e);
+ error.setException(new SFException(e, ErrorCode.INTERNAL_ERROR, e.getMessage()));
+ response.addError(error);
+ }
+ rowIndex++;
+ checkBatchSizeEnforcedMaximum(tempRowsSizeInBytes);
+ if ((long) rowBuffer.bufferedRowCount + rowIndex >= Integer.MAX_VALUE) {
+ throw new SFException(ErrorCode.INTERNAL_ERROR, "Row count reaches MAX value");
+ }
+ }
+
+ if (!response.hasErrors()) {
+ checkBatchSizeRecommendedMaximum(tempRowsSizeInBytes);
+ moveTempRowsToActualBuffer(tempRowCount);
+ rowsSizeInBytes = tempRowsSizeInBytes;
+ rowBuffer.statsMap.forEach(
+ (colName, stats) ->
+ rowBuffer.statsMap.put(
+ colName,
+ RowBufferStats.getCombinedStats(stats, rowBuffer.tempStatsMap.get(colName))));
+ checkOffsetMismatch(
+ rowBuffer.channelState.getEndOffsetToken(), startOffsetToken, endOffsetToken, rowIndex);
+ rowBuffer.channelState.updateOffsetToken(
+ startOffsetToken, endOffsetToken, rowBuffer.bufferedRowCount);
+ rowBuffer.bufferedRowCount += tempRowCount;
+ rowBuffer.bufferSize += rowsSizeInBytes;
+ rowBuffer.rowSizeMetric.accept(rowsSizeInBytes);
+ }
+ return response;
+ }
+ }
+
+ // Map the column name to the stats
+ @VisibleForTesting Map statsMap;
+
+ // Temp stats map to use until all the rows are validated
+ @VisibleForTesting Map tempStatsMap;
+
+ // Lock used to protect the buffers from concurrent read/write
+ private final Lock flushLock;
+
+ // Current row count buffered
+ @VisibleForTesting volatile int bufferedRowCount;
+
+ // Current buffer size
+ private volatile float bufferSize;
+
+ // Names of non-nullable columns
+ private final Set nonNullableFieldNames;
+
+ // Buffer's channel fully qualified name with database, schema and table
+ final String channelFullyQualifiedName;
+
+ // Metric callback to report size of inserted rows
+ private final Consumer rowSizeMetric;
+
+ // State of the owning channel
+ final ChannelRuntimeState channelState;
+
+ // ON_ERROR option for this channel
+ final OpenChannelRequest.OnErrorOption onErrorOption;
+
+ final ZoneId defaultTimezone;
+
+ // Buffer parameters that are set at the owning client level
+ final ClientBufferParameters clientBufferParameters;
+
+ // Function used to verify offset token logic
+ final OffsetTokenVerificationFunction offsetTokenVerificationFunction;
+
+ // Telemetry service use to report telemetry to SF
+ final TelemetryService telemetryService;
+
+ AbstractRowBuffer(
+ OpenChannelRequest.OnErrorOption onErrorOption,
+ ZoneId defaultTimezone,
+ String fullyQualifiedChannelName,
+ Consumer rowSizeMetric,
+ ChannelRuntimeState channelRuntimeState,
+ ClientBufferParameters clientBufferParameters,
+ OffsetTokenVerificationFunction offsetTokenVerificationFunction,
+ TelemetryService telemetryService) {
+ this.onErrorOption = onErrorOption;
+ this.defaultTimezone = defaultTimezone;
+ this.rowSizeMetric = rowSizeMetric;
+ this.channelState = channelRuntimeState;
+ this.channelFullyQualifiedName = fullyQualifiedChannelName;
+ this.nonNullableFieldNames = new HashSet<>();
+ this.flushLock = new ReentrantLock();
+ this.bufferedRowCount = 0;
+ this.bufferSize = 0F;
+ this.clientBufferParameters = clientBufferParameters;
+ this.offsetTokenVerificationFunction = offsetTokenVerificationFunction;
+ this.telemetryService = telemetryService;
+
+ // Initialize empty stats
+ this.statsMap = new HashMap<>();
+ this.tempStatsMap = new HashMap<>();
+ }
+
+ /**
+ * Adds non-nullable field name. It is used to check if all non-nullable fields have been
+ * provided.
+ *
+ * @param nonNullableFieldName non-nullable filed name
+ */
+ void addNonNullableFieldName(String nonNullableFieldName) {
+ nonNullableFieldNames.add(nonNullableFieldName);
+ }
+
+ /** Throws an exception if the column has a collation defined. */
+ void validateColumnCollation(ColumnMetadata column) {
+ if (column.getCollation() != null) {
+ throw new SFException(
+ ErrorCode.UNSUPPORTED_DATA_TYPE,
+ String.format(
+ "Column %s with collation %s detected. Ingestion into collated columns is not"
+ + " supported",
+ column.getName(), column.getCollation()));
+ }
+ }
+
+ /**
+ * Get the current buffer size
+ *
+ * @return the current buffer size
+ */
+ @Override
+ public float getSize() {
+ return bufferSize;
+ }
+
+ /**
+ * Verify that the input row columns are all valid.
+ *
+ * Checks that the columns, specified in the row, are present in the table and values for all
+ * non-nullable columns are specified.
+ *
+ * @param row the input row
+ * @param error the insert error that we return to the customer
+ * @param rowIndex the index of the current row in the input batch
+ * @return the set of input column names
+ */
+ Set verifyInputColumns(
+ Map row, InsertValidationResponse.InsertError error, int rowIndex) {
+ // Map of unquoted column name -> original column name
+ Map inputColNamesMap =
+ row.keySet().stream()
+ .collect(Collectors.toMap(LiteralQuoteUtils::unquoteColumnName, value -> value));
+
+ // Check for extra columns in the row
+ List extraCols = new ArrayList<>();
+ for (String columnName : inputColNamesMap.keySet()) {
+ if (!hasColumn(columnName)) {
+ extraCols.add(inputColNamesMap.get(columnName));
+ }
+ }
+
+ if (!extraCols.isEmpty()) {
+ if (error != null) {
+ error.setExtraColNames(extraCols);
+ }
+ throw new SFException(
+ ErrorCode.INVALID_FORMAT_ROW,
+ "Extra columns: " + extraCols,
+ String.format(
+ "Columns not present in the table shouldn't be specified, rowIndex:%d", rowIndex));
+ }
+
+ // Check for missing columns in the row
+ List missingCols = new ArrayList<>();
+ for (String columnName : this.nonNullableFieldNames) {
+ if (!inputColNamesMap.containsKey(columnName)) {
+ missingCols.add(statsMap.get(columnName).getColumnDisplayName());
+ }
+ }
+
+ if (!missingCols.isEmpty()) {
+ if (error != null) {
+ error.setMissingNotNullColNames(missingCols);
+ }
+ throw new SFException(
+ ErrorCode.INVALID_FORMAT_ROW,
+ "Missing columns: " + missingCols,
+ String.format(
+ "Values for all non-nullable columns must be specified, rowIndex:%d", rowIndex));
+ }
+
+ return inputColNamesMap.keySet();
+ }
+
+ /**
+ * Insert a batch of rows into the row buffer
+ *
+ * @param rows input row
+ * @param startOffsetToken start offset token of the batch
+ * @param endOffsetToken offset token of the latest row in the batch
+ * @return insert response that possibly contains errors because of insertion failures
+ */
+ @Override
+ public InsertValidationResponse insertRows(
+ Iterable> rows, String startOffsetToken, String endOffsetToken) {
+ if (!hasColumns()) {
+ throw new SFException(ErrorCode.INTERNAL_ERROR, "Empty column fields");
+ }
+ InsertValidationResponse response = null;
+ this.flushLock.lock();
+ try {
+ this.channelState.updateInsertStats(System.currentTimeMillis(), this.bufferedRowCount);
+ IngestionStrategy ingestionStrategy = createIngestionStrategy(onErrorOption);
+ response = ingestionStrategy.insertRows(this, rows, startOffsetToken, endOffsetToken);
+ } finally {
+ this.tempStatsMap.values().forEach(RowBufferStats::reset);
+ clearTempRows();
+ this.flushLock.unlock();
+ }
+ return response;
+ }
+
+ /**
+ * Flush the data in the row buffer by taking the ownership of the old vectors and pass all the
+ * required info back to the flush service to build the blob
+ *
+ * @param filePath the name of the file the data will be written in
+ * @return A ChannelData object that contains the info needed by the flush service to build a blob
+ */
+ @Override
+ public ChannelData flush(final String filePath) {
+ logger.logDebug("Start get data for channel={}", channelFullyQualifiedName);
+ if (this.bufferedRowCount > 0) {
+ Optional oldData = Optional.empty();
+ int oldRowCount = 0;
+ float oldBufferSize = 0F;
+ long oldRowSequencer = 0;
+ String oldEndOffsetToken = null;
+ String oldStartOffsetToken = null;
+ Map oldColumnEps = null;
+ Pair oldMinMaxInsertTimeInMs = null;
+
+ logger.logDebug("Buffer flush about to take lock on channel={}", channelFullyQualifiedName);
+
+ this.flushLock.lock();
+ try {
+ if (this.bufferedRowCount > 0) {
+ // Transfer the ownership of the vectors
+ oldData = getSnapshot(filePath);
+ oldRowCount = this.bufferedRowCount;
+ oldBufferSize = this.bufferSize;
+ oldRowSequencer = this.channelState.incrementAndGetRowSequencer();
+ oldEndOffsetToken = this.channelState.getEndOffsetToken();
+ oldStartOffsetToken = this.channelState.getStartOffsetToken();
+ oldColumnEps = new HashMap<>(this.statsMap);
+ oldMinMaxInsertTimeInMs =
+ new Pair<>(
+ this.channelState.getFirstInsertInMs(), this.channelState.getLastInsertInMs());
+ // Reset everything in the buffer once we save all the info
+ reset();
+ }
+ } finally {
+ this.flushLock.unlock();
+ }
+
+ logger.logDebug(
+ "Buffer flush released lock on channel={}, rowCount={}, bufferSize={}",
+ channelFullyQualifiedName,
+ oldRowCount,
+ oldBufferSize);
+
+ if (oldData.isPresent()) {
+ ChannelData data = new ChannelData<>();
+ data.setVectors(oldData.get());
+ data.setRowCount(oldRowCount);
+ data.setBufferSize(oldBufferSize);
+ data.setRowSequencer(oldRowSequencer);
+ data.setEndOffsetToken(oldEndOffsetToken);
+ data.setStartOffsetToken(oldStartOffsetToken);
+ data.setColumnEps(oldColumnEps);
+ data.setMinMaxInsertTimeInMs(oldMinMaxInsertTimeInMs);
+ data.setFlusherFactory(this::createFlusher);
+ return data;
+ }
+ }
+ return null;
+ }
+
+ /** Whether row has a column with a given name. */
+ abstract boolean hasColumn(String name);
+
+ /**
+ * Add an input row to the buffer.
+ *
+ * @param row input row
+ * @param bufferedRowIndex Row number of buffered rows which will eventually by flushed.
+ * @param statsMap column stats map
+ * @param formattedInputColumnNames list of input column names after formatting
+ * @param insertRowIndex Index of the rows given in insertRows API. Not the same as
+ * bufferedRowIndex
+ * @return row size
+ */
+ abstract float addRow(
+ Map row,
+ int bufferedRowIndex,
+ Map statsMap,
+ Set formattedInputColumnNames,
+ final long insertRowIndex);
+
+ /**
+ * Add an input row to the temporary row buffer.
+ *
+ * The temporary row buffer is used to add rows before validating all of them. Once all the
+ * rows to add have been validated, the temporary row buffer is moved to the actual one.
+ *
+ * @param row input row
+ * @param curRowIndex current row index to use
+ * @param statsMap column stats map
+ * @param formattedInputColumnNames list of input column names after formatting
+ * @param insertRowIndex index of the row being inserteed from User Input List
+ * @return row size
+ */
+ abstract float addTempRow(
+ Map row,
+ int curRowIndex,
+ Map statsMap,
+ Set formattedInputColumnNames,
+ long insertRowIndex);
+
+ /** Move rows from the temporary buffer to the current row buffer. */
+ abstract void moveTempRowsToActualBuffer(int tempRowCount);
+
+ /** Clear the temporary row buffer. */
+ abstract void clearTempRows();
+
+ /** If row has any column. */
+ abstract boolean hasColumns();
+
+ /** Reset the variables after each flush. Note that the caller needs to handle synchronization. */
+ void reset() {
+ this.bufferedRowCount = 0;
+ this.bufferSize = 0F;
+ this.statsMap.replaceAll((key, value) -> value.forkEmpty());
+ }
+
+ /**
+ * Get buffered data snapshot for later flushing.
+ *
+ * @param filePath the name of the file the data will be written in
+ */
+ abstract Optional getSnapshot(final String filePath);
+
+ @VisibleForTesting
+ abstract Object getVectorValueAt(String column, int index);
+
+ @VisibleForTesting
+ abstract int getTempRowCount();
+
+ /**
+ * Close row buffer by releasing its internal resources only. Note that the allocated memory for
+ * the channel is released elsewhere (in close).
+ */
+ abstract void closeInternal();
+
+ /** Close the row buffer and release allocated memory for the channel. */
+ @Override
+ public synchronized void close(String name) {
+
+ closeInternal();
+ }
+
+ /**
+ * Given a set of col names to stats, build the right ep info
+ *
+ * @param rowCount: count of rows in the given buffer
+ * @param colStats: map of column name to RowBufferStats
+ * @return the EPs built from column stats
+ */
+ static EpInfo buildEpInfoFromStats(long rowCount, Map colStats) {
+ EpInfo epInfo = new EpInfo(rowCount, new HashMap<>());
+ for (Map.Entry colStat : colStats.entrySet()) {
+ RowBufferStats stat = colStat.getValue();
+ FileColumnProperties dto = new FileColumnProperties(stat);
+ String colName = colStat.getValue().getColumnDisplayName();
+ epInfo.getColumnEps().put(colName, dto);
+ }
+ epInfo.verifyEpInfo();
+ return epInfo;
+ }
+
+ /** Row buffer factory. */
+ static AbstractRowBuffer createRowBuffer(
+ OpenChannelRequest.OnErrorOption onErrorOption,
+ ZoneId defaultTimezone,
+ Constants.BdecVersion bdecVersion,
+ String fullyQualifiedChannelName,
+ Consumer rowSizeMetric,
+ ChannelRuntimeState channelRuntimeState,
+ ClientBufferParameters clientBufferParameters,
+ OffsetTokenVerificationFunction offsetTokenVerificationFunction,
+ TelemetryService telemetryService) {
+ switch (bdecVersion) {
+ case THREE:
+ //noinspection unchecked
+ return (AbstractRowBuffer)
+ new ParquetRowBuffer(
+ onErrorOption,
+ defaultTimezone,
+ fullyQualifiedChannelName,
+ rowSizeMetric,
+ channelRuntimeState,
+ clientBufferParameters,
+ offsetTokenVerificationFunction,
+ telemetryService);
+ default:
+ throw new SFException(
+ ErrorCode.INTERNAL_ERROR, "Unsupported BDEC format version: " + bdecVersion);
+ }
+ }
+
+ private void checkBatchSizeEnforcedMaximum(float batchSizeInBytes) {
+ if (batchSizeInBytes > clientBufferParameters.getMaxChunkSizeInBytes()) {
+ throw new SFException(
+ ErrorCode.MAX_BATCH_SIZE_EXCEEDED,
+ clientBufferParameters.getMaxChunkSizeInBytes(),
+ INSERT_ROWS_RECOMMENDED_MAX_BATCH_SIZE_IN_BYTES);
+ }
+ }
+
+ private void checkBatchSizeRecommendedMaximum(float batchSizeInBytes) {
+ if (batchSizeInBytes > INSERT_ROWS_RECOMMENDED_MAX_BATCH_SIZE_IN_BYTES) {
+ logger.logWarn(
+ "The batch of rows passed to 'insertRows' is over the recommended max batch size. Given"
+ + " {} bytes, recommended max batch is {} bytes. For optimal performance and memory"
+ + " utilization, we recommend splitting large batches into multiple smaller ones and"
+ + " call insertRows for each smaller batch separately.",
+ batchSizeInBytes,
+ INSERT_ROWS_RECOMMENDED_MAX_BATCH_SIZE_IN_BYTES);
+ }
+ }
+
+ /**
+ * We verify the offset token based on the provided verification logic and report to SF if there
+ * is a mismatch.
+ */
+ private void checkOffsetMismatch(
+ String prevEndOffset, String curStartOffset, String curEndOffset, int rowCount) {
+ if (offsetTokenVerificationFunction != null
+ && !offsetTokenVerificationFunction.verify(
+ prevEndOffset, curStartOffset, curEndOffset, rowCount)) {
+ logger.logWarn(
+ "The channel {} might have an offset token mismatch based on the provided offset token"
+ + " verification logic, preEndOffset={}, curStartOffset={}, curEndOffset={},"
+ + " rowCount={}.",
+ channelFullyQualifiedName,
+ prevEndOffset,
+ curStartOffset,
+ curEndOffset,
+ rowCount);
+ if (telemetryService != null) {
+ telemetryService.reportBatchOffsetMismatch(
+ channelFullyQualifiedName, prevEndOffset, curStartOffset, curEndOffset, rowCount);
+ }
+ }
+ }
+
+ /** Create the ingestion strategy based on the channel on_error option */
+ IngestionStrategy createIngestionStrategy(OpenChannelRequest.OnErrorOption onErrorOption) {
+ switch (onErrorOption) {
+ case CONTINUE:
+ return new ContinueIngestionStrategy<>();
+ case ABORT:
+ return new AbortIngestionStrategy<>();
+ case SKIP_BATCH:
+ return new SkipBatchIngestionStrategy<>();
+ default:
+ throw new IllegalArgumentException("Unknown on error option: " + onErrorOption);
+ }
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ArrowRowBuffer.java b/src/main/java/net/snowflake/ingest/streaming/internal/ArrowRowBuffer.java
deleted file mode 100644
index ea70c50be..000000000
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ArrowRowBuffer.java
+++ /dev/null
@@ -1,962 +0,0 @@
-/*
- * Copyright (c) 2021 Snowflake Computing Inc. All rights reserved.
- */
-
-package net.snowflake.ingest.streaming.internal;
-
-import java.math.BigDecimal;
-import java.math.BigInteger;
-import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
-import java.util.stream.Collectors;
-import net.snowflake.client.jdbc.internal.google.common.collect.Sets;
-import net.snowflake.client.jdbc.internal.snowflake.common.util.Power10;
-import net.snowflake.ingest.streaming.InsertValidationResponse;
-import net.snowflake.ingest.streaming.OpenChannelRequest;
-import net.snowflake.ingest.utils.ErrorCode;
-import net.snowflake.ingest.utils.Logging;
-import net.snowflake.ingest.utils.SFException;
-import net.snowflake.ingest.utils.Utils;
-import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.util.VisibleForTesting;
-import org.apache.arrow.vector.BaseFixedWidthVector;
-import org.apache.arrow.vector.BaseVariableWidthVector;
-import org.apache.arrow.vector.BigIntVector;
-import org.apache.arrow.vector.BitVector;
-import org.apache.arrow.vector.DateDayVector;
-import org.apache.arrow.vector.DecimalVector;
-import org.apache.arrow.vector.FieldVector;
-import org.apache.arrow.vector.Float8Vector;
-import org.apache.arrow.vector.IntVector;
-import org.apache.arrow.vector.SmallIntVector;
-import org.apache.arrow.vector.TinyIntVector;
-import org.apache.arrow.vector.VarBinaryVector;
-import org.apache.arrow.vector.VarCharVector;
-import org.apache.arrow.vector.VectorSchemaRoot;
-import org.apache.arrow.vector.complex.StructVector;
-import org.apache.arrow.vector.types.Types;
-import org.apache.arrow.vector.types.pojo.ArrowType;
-import org.apache.arrow.vector.types.pojo.Field;
-import org.apache.arrow.vector.types.pojo.FieldType;
-import org.apache.arrow.vector.util.Text;
-import org.apache.arrow.vector.util.TransferPair;
-
-/**
- * The buffer in the Streaming Ingest channel that holds the un-flushed rows, these rows will be
- * converted to Arrow format for faster processing
- */
-class ArrowRowBuffer {
-
- // Snowflake table column logical type
- private static enum ColumnLogicalType {
- ANY,
- BOOLEAN,
- ROWINDEX,
- NULL,
- REAL,
- FIXED,
- TEXT,
- CHAR,
- BINARY,
- DATE,
- TIME,
- TIMESTAMP_LTZ,
- TIMESTAMP_NTZ,
- TIMESTAMP_TZ,
- INTERVAL,
- RAW,
- ARRAY,
- OBJECT,
- VARIANT,
- ROW,
- SEQUENCE,
- FUNCTION,
- USER_DEFINED_TYPE,
- }
-
- // Snowflake table column physical type
- private static enum ColumnPhysicalType {
- ROWINDEX,
- DOUBLE,
- SB1,
- SB2,
- SB4,
- SB8,
- SB16,
- LOB,
- BINARY,
- ROW,
- }
-
- private static final Logging logger = new Logging(ArrowRowBuffer.class);
-
- // Constants for column fields
- private static final String FIELD_EPOCH_IN_SECONDS = "epoch"; // seconds since epoch
- private static final String FIELD_TIME_ZONE = "timezone"; // time zone index
- private static final String FIELD_FRACTION_IN_NANOSECONDS = "fraction"; // fraction in nanoseconds
-
- // Column metadata that will send back to server as part of the blob, and will be used by the
- // Arrow reader
- private static final String COLUMN_PHYSICAL_TYPE = "physicalType";
- private static final String COLUMN_LOGICAL_TYPE = "logicalType";
- private static final String COLUMN_NULLABLE = "nullable";
- static final String COLUMN_SCALE = "scale";
- private static final String COLUMN_PRECISION = "precision";
- private static final String COLUMN_CHAR_LENGTH = "charLength";
- private static final String COLUMN_BYTE_LENGTH = "byteLength";
- @VisibleForTesting static final int DECIMAL_BIT_WIDTH = 128;
-
- // Holder for a set of the Arrow vectors (buffers)
- @VisibleForTesting VectorSchemaRoot vectorsRoot;
-
- // For ABORT on_error option, temp vectors are needed to temporarily holding the rows until
- // they're all validated, then the rows will be transferred to the final VectorSchemaRoot
- @VisibleForTesting VectorSchemaRoot tempVectorsRoot;
-
- // Map the column name to Arrow column field
- private final Map fields;
-
- // Map the column name to the stats
- @VisibleForTesting Map statsMap;
-
- // Temp stats map to use until all the rows are validated
- @VisibleForTesting Map tempStatsMap;
-
- // Lock used to protect the buffers from concurrent read/write
- private final Lock flushLock;
-
- // Current row count
- @VisibleForTesting volatile int rowCount;
-
- // Allocator used to allocate the buffers
- private final BufferAllocator allocator;
-
- // Current buffer size
- private volatile float bufferSize;
-
- // Reference to the Streaming Ingest channel that owns this buffer
- @VisibleForTesting final SnowflakeStreamingIngestChannelInternal owningChannel;
-
- // Names of non-nullable columns
- private final Set nonNullableFieldNames;
-
- /**
- * Given a set of col names to stats, build the right ep info
- *
- * @param rowCount: count of rows in the given arrow buffer
- * @param colStats: map of column name to RowBufferStats
- * @return the EPs built from column stats
- */
- static EpInfo buildEpInfoFromStats(long rowCount, Map colStats) {
- EpInfo epInfo = new EpInfo(rowCount, new HashMap<>());
- for (Map.Entry colStat : colStats.entrySet()) {
- RowBufferStats stat = colStat.getValue();
- FileColumnProperties dto = new FileColumnProperties(stat);
-
- String colName = colStat.getKey();
- epInfo.getColumnEps().put(colName, dto);
- }
- return epInfo;
- }
-
- /**
- * Construct a ArrowRowBuffer object
- *
- * @param channel
- */
- ArrowRowBuffer(SnowflakeStreamingIngestChannelInternal channel) {
- this.owningChannel = channel;
- this.allocator = channel.getAllocator();
- this.fields = new HashMap<>();
- this.nonNullableFieldNames = new HashSet<>();
- this.flushLock = new ReentrantLock();
- this.rowCount = 0;
- this.bufferSize = 0F;
-
- // Initialize empty stats
- this.statsMap = new HashMap<>();
- this.tempStatsMap = new HashMap<>();
- }
-
- /**
- * Setup the column fields and vectors using the column metadata from the server
- *
- * @param columns list of column metadata
- */
- void setupSchema(List columns) {
- List vectors = new ArrayList<>();
- List tempVectors = new ArrayList<>();
-
- for (ColumnMetadata column : columns) {
- Field field = buildField(column);
- FieldVector vector = field.createVector(this.allocator);
- if (!field.isNullable()) {
- this.nonNullableFieldNames.add(field.getName());
- }
- this.fields.put(column.getName(), field);
- vectors.add(vector);
- this.statsMap.put(column.getName(), new RowBufferStats(column.getCollation()));
-
- if (this.owningChannel.getOnErrorOption() == OpenChannelRequest.OnErrorOption.ABORT) {
- FieldVector tempVector = field.createVector(this.allocator);
- tempVectors.add(tempVector);
- this.tempStatsMap.put(column.getName(), new RowBufferStats(column.getCollation()));
- }
- }
-
- this.vectorsRoot = new VectorSchemaRoot(vectors);
- this.tempVectorsRoot = new VectorSchemaRoot(tempVectors);
- }
-
- /**
- * Close the row buffer and release resources. Note that the caller needs to handle
- * synchronization
- */
- void close() {
- if (this.vectorsRoot != null) {
- this.vectorsRoot.close();
- this.tempVectorsRoot.close();
- }
- this.fields.clear();
- Utils.closeAllocator(this.allocator);
- }
-
- /** Reset the variables after each flush. Note that the caller needs to handle synchronization */
- void reset() {
- this.vectorsRoot.clear();
- this.rowCount = 0;
- this.bufferSize = 0F;
- this.statsMap.replaceAll(
- (key, value) -> new RowBufferStats(value.getCollationDefinitionString()));
- }
-
- /**
- * Get the current buffer size
- *
- * @return the current buffer size
- */
- float getSize() {
- return this.bufferSize;
- }
-
- /**
- * Insert a batch of rows into the row buffer
- *
- * @param rows input row
- * @param offsetToken offset token of the latest row in the batch
- * @return insert response that possibly contains errors because of insertion failures
- */
- InsertValidationResponse insertRows(Iterable> rows, String offsetToken) {
- float rowSize = 0F;
- if (this.fields.isEmpty()) {
- throw new SFException(ErrorCode.INTERNAL_ERROR, "Empty column fields");
- }
-
- InsertValidationResponse response = new InsertValidationResponse();
- this.flushLock.lock();
- try {
- if (this.owningChannel.getOnErrorOption() == OpenChannelRequest.OnErrorOption.CONTINUE) {
- // Used to map incoming row(nth row) to InsertError(for nth row) in response
- long rowIndex = 0;
- for (Map row : rows) {
- try {
- rowSize += convertRowToArrow(row, this.vectorsRoot, this.rowCount, this.statsMap);
- this.rowCount++;
- this.bufferSize += rowSize;
- } catch (SFException e) {
- response.addError(new InsertValidationResponse.InsertError(row, e, rowIndex));
- } catch (Throwable e) {
- logger.logWarn("Unexpected error happens during insertRows: {}", e.getMessage());
- response.addError(
- new InsertValidationResponse.InsertError(
- row, new SFException(e, ErrorCode.INTERNAL_ERROR, e.getMessage()), rowIndex));
- }
- rowIndex++;
- }
- } else {
- // If the on_error option is ABORT, simply throw the first exception
- float tempRowSize = 0F;
- int tempRowCount = 0;
- for (Map row : rows) {
- tempRowSize +=
- convertRowToArrow(row, this.tempVectorsRoot, tempRowCount, this.tempStatsMap);
- tempRowCount++;
- }
-
- // If all the rows are inserted successfully, transfer the rows from temp vectors to
- // the final vectors and update the row size and row count
- // TODO: switch to VectorSchemaRootAppender once it works for all vector types
- for (Field field : fields.values()) {
- FieldVector from = this.tempVectorsRoot.getVector(field);
- FieldVector to = this.vectorsRoot.getVector(field);
- for (int rowIdx = 0; rowIdx < tempRowCount; rowIdx++) {
- to.copyFromSafe(rowIdx, this.rowCount + rowIdx, from);
- }
- }
- rowSize = tempRowSize;
- this.rowCount += tempRowCount;
- this.bufferSize += rowSize;
- this.statsMap.forEach(
- (colName, stats) -> {
- this.statsMap.put(
- colName, RowBufferStats.getCombinedStats(stats, this.tempStatsMap.get(colName)));
- });
- }
-
- this.owningChannel.setOffsetToken(offsetToken);
- this.owningChannel.collectRowSize(rowSize);
- } finally {
- this.tempStatsMap.values().forEach(RowBufferStats::reset);
- this.tempVectorsRoot.clear();
- this.flushLock.unlock();
- }
-
- return response;
- }
-
- /**
- * Flush the data in the row buffer by taking the ownership of the old vectors and pass all the
- * required info back to the flush service to build the blob
- *
- * @return A ChannelData object that contains the info needed by the flush service to build a blob
- */
- ChannelData flush() {
- logger.logDebug("Start get data for channel={}", this.owningChannel.getFullyQualifiedName());
- if (this.rowCount > 0) {
- List oldVectors = new ArrayList<>();
- int oldRowCount = 0;
- float oldBufferSize = 0F;
- long oldRowSequencer = 0;
- String oldOffsetToken = null;
- Map oldColumnEps = null;
-
- logger.logDebug(
- "Arrow buffer flush about to take lock on channel={}",
- this.owningChannel.getFullyQualifiedName());
-
- this.flushLock.lock();
- try {
- if (this.rowCount > 0) {
- // Transfer the ownership of the vectors
- for (FieldVector vector : this.vectorsRoot.getFieldVectors()) {
- vector.setValueCount(this.rowCount);
- if (vector instanceof DecimalVector) {
- // DecimalVectors do not transfer FieldType metadata when using
- // vector.getTransferPair. We need to explicitly create the new vector to transfer to
- // in order to keep the metadata.
- ArrowType arrowType =
- new ArrowType.Decimal(
- ((DecimalVector) vector).getPrecision(),
- ((DecimalVector) vector).getScale(),
- DECIMAL_BIT_WIDTH);
- FieldType fieldType =
- new FieldType(
- vector.getField().isNullable(),
- arrowType,
- null,
- vector.getField().getMetadata());
- Field f = new Field(vector.getName(), fieldType, null);
- DecimalVector newVector = new DecimalVector(f, this.allocator);
- TransferPair t = vector.makeTransferPair(newVector);
- t.transfer();
- oldVectors.add((FieldVector) t.getTo());
- } else {
- TransferPair t = vector.getTransferPair(this.allocator);
- t.transfer();
- oldVectors.add((FieldVector) t.getTo());
- }
- }
-
- oldRowCount = this.rowCount;
- oldBufferSize = this.bufferSize;
- oldRowSequencer = this.owningChannel.incrementAndGetRowSequencer();
- oldOffsetToken = this.owningChannel.getOffsetToken();
- oldColumnEps = new HashMap<>(this.statsMap);
- // Reset everything in the buffer once we save all the info
- reset();
- }
- } finally {
- this.flushLock.unlock();
- }
-
- logger.logDebug(
- "Arrow buffer flush released lock on channel={}, rowCount={}, bufferSize={}",
- this.owningChannel.getFullyQualifiedName(),
- rowCount,
- bufferSize);
-
- if (!oldVectors.isEmpty()) {
- ChannelData data = new ChannelData();
- VectorSchemaRoot vectors = new VectorSchemaRoot(oldVectors);
- vectors.setRowCount(oldRowCount);
- data.setVectors(vectors);
- data.setBufferSize(oldBufferSize);
- data.setChannel(this.owningChannel);
- data.setRowSequencer(oldRowSequencer);
- data.setOffsetToken(oldOffsetToken);
- data.setColumnEps(oldColumnEps);
-
- return data;
- }
- }
- return null;
- }
-
- /**
- * Build the column field from the column metadata
- *
- * @param column column metadata
- * @return Column field object
- */
- Field buildField(ColumnMetadata column) {
- ArrowType arrowType;
- FieldType fieldType;
- List children = null;
-
- // Put info into the metadata, which will be used by the Arrow reader later
- Map metadata = new HashMap<>();
- metadata.put(COLUMN_LOGICAL_TYPE, column.getLogicalType());
- metadata.put(COLUMN_PHYSICAL_TYPE, column.getPhysicalType());
- metadata.put(COLUMN_NULLABLE, String.valueOf(column.getNullable()));
-
- ColumnPhysicalType physicalType;
- ColumnLogicalType logicalType;
- try {
- physicalType = ColumnPhysicalType.valueOf(column.getPhysicalType());
- logicalType = ColumnLogicalType.valueOf(column.getLogicalType());
- } catch (IllegalArgumentException e) {
- throw new SFException(
- ErrorCode.UNKNOWN_DATA_TYPE, column.getLogicalType(), column.getPhysicalType());
- }
-
- if (column.getPrecision() != null) {
- metadata.put(COLUMN_PRECISION, column.getPrecision().toString());
- }
- if (column.getScale() != null) {
- metadata.put(COLUMN_SCALE, column.getScale().toString());
- }
- if (column.getByteLength() != null) {
- metadata.put(COLUMN_BYTE_LENGTH, column.getByteLength().toString());
- }
- if (column.getLength() != null) {
- metadata.put(COLUMN_CHAR_LENGTH, column.getLength().toString());
- }
-
- // Handle differently depends on the column logical and physical types
- switch (logicalType) {
- case FIXED:
- if ((column.getScale() != null && column.getScale() != 0)
- || physicalType == ColumnPhysicalType.SB16) {
- arrowType =
- new ArrowType.Decimal(column.getPrecision(), column.getScale(), DECIMAL_BIT_WIDTH);
- } else {
- switch (physicalType) {
- case SB1:
- arrowType = Types.MinorType.TINYINT.getType();
- break;
- case SB2:
- arrowType = Types.MinorType.SMALLINT.getType();
- break;
- case SB4:
- arrowType = Types.MinorType.INT.getType();
- break;
- case SB8:
- arrowType = Types.MinorType.BIGINT.getType();
- break;
- default:
- throw new SFException(
- ErrorCode.UNKNOWN_DATA_TYPE, column.getLogicalType(), column.getPhysicalType());
- }
- }
- break;
- case ANY:
- case ARRAY:
- case CHAR:
- case TEXT:
- case OBJECT:
- case VARIANT:
- arrowType = Types.MinorType.VARCHAR.getType();
- break;
- case TIMESTAMP_LTZ:
- case TIMESTAMP_NTZ:
- switch (physicalType) {
- case SB8:
- arrowType = Types.MinorType.BIGINT.getType();
- break;
- case SB16:
- {
- arrowType = Types.MinorType.STRUCT.getType();
- FieldType fieldTypeEpoch =
- new FieldType(true, Types.MinorType.BIGINT.getType(), null, metadata);
- FieldType fieldTypeFraction =
- new FieldType(true, Types.MinorType.INT.getType(), null, metadata);
- Field fieldEpoch = new Field(FIELD_EPOCH_IN_SECONDS, fieldTypeEpoch, null);
- Field fieldFraction =
- new Field(FIELD_FRACTION_IN_NANOSECONDS, fieldTypeFraction, null);
- children = new ArrayList<>();
- children.add(fieldEpoch);
- children.add(fieldFraction);
- break;
- }
- default:
- throw new SFException(
- ErrorCode.UNKNOWN_DATA_TYPE, column.getLogicalType(), column.getPhysicalType());
- }
- break;
- case TIMESTAMP_TZ:
- switch (physicalType) {
- case SB8:
- {
- arrowType = Types.MinorType.STRUCT.getType();
- FieldType fieldTypeEpoch =
- new FieldType(true, Types.MinorType.BIGINT.getType(), null, metadata);
- FieldType fieldTypeTimezone =
- new FieldType(true, Types.MinorType.INT.getType(), null, metadata);
- Field fieldEpoch = new Field(FIELD_EPOCH_IN_SECONDS, fieldTypeEpoch, null);
- Field fieldTimezone = new Field(FIELD_TIME_ZONE, fieldTypeTimezone, null);
-
- children = new ArrayList<>();
- children.add(fieldEpoch);
- children.add(fieldTimezone);
- break;
- }
- case SB16:
- {
- arrowType = Types.MinorType.STRUCT.getType();
- FieldType fieldTypeEpoch =
- new FieldType(true, Types.MinorType.BIGINT.getType(), null, metadata);
- FieldType fieldTypeFraction =
- new FieldType(true, Types.MinorType.INT.getType(), null, metadata);
- FieldType fieldTypeTimezone =
- new FieldType(true, Types.MinorType.INT.getType(), null, metadata);
- Field fieldEpoch = new Field(FIELD_EPOCH_IN_SECONDS, fieldTypeEpoch, null);
- Field fieldFraction =
- new Field(FIELD_FRACTION_IN_NANOSECONDS, fieldTypeFraction, null);
- Field fieldTimezone = new Field(FIELD_TIME_ZONE, fieldTypeTimezone, null);
-
- children = new ArrayList<>();
- children.add(fieldEpoch);
- children.add(fieldFraction);
- children.add(fieldTimezone);
- break;
- }
- default:
- throw new SFException(
- ErrorCode.UNKNOWN_DATA_TYPE,
- "Unknown physical type for TIMESTAMP_TZ: " + physicalType);
- }
- break;
- case DATE:
- arrowType = Types.MinorType.DATEDAY.getType();
- break;
- case TIME:
- switch (physicalType) {
- case SB4:
- arrowType = Types.MinorType.INT.getType();
- break;
- case SB8:
- arrowType = Types.MinorType.BIGINT.getType();
- break;
- default:
- throw new SFException(
- ErrorCode.UNKNOWN_DATA_TYPE, column.getLogicalType(), column.getPhysicalType());
- }
- break;
- case BOOLEAN:
- arrowType = Types.MinorType.BIT.getType();
- break;
- case BINARY:
- arrowType = Types.MinorType.VARBINARY.getType();
- break;
- case REAL:
- arrowType = Types.MinorType.FLOAT8.getType();
- break;
- default:
- throw new SFException(
- ErrorCode.UNKNOWN_DATA_TYPE, column.getLogicalType(), column.getPhysicalType());
- }
-
- // Create the corresponding column field base on the column data type
- fieldType = new FieldType(column.getNullable(), arrowType, null, metadata);
- return new Field(column.getName(), fieldType, children);
- }
-
- private String formatColumnName(String columnName) {
- Utils.assertStringNotNullOrEmpty("invalid column name", columnName);
- return (columnName.charAt(0) == '"' && columnName.charAt(columnName.length() - 1) == '"')
- ? columnName.substring(1, columnName.length() - 1)
- : columnName.toUpperCase();
- }
-
- /**
- * Verify that the input row columns are all valid
- *
- * @param row the input row
- * @return the set of input column names
- */
- private Set verifyInputColumns(Map row) {
- Set inputColumns =
- row.keySet().stream().map(this::formatColumnName).collect(Collectors.toSet());
-
- for (String columnName : this.nonNullableFieldNames) {
- if (!inputColumns.contains(columnName)) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- "Missing column: " + columnName,
- "Values for all non-nullable columns must be specified.");
- }
- }
-
- for (String columnName : inputColumns) {
- Field field = this.fields.get(columnName);
- if (field == null) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- "Extra column: " + columnName,
- "Columns not present in the table shouldn't be specified.");
- }
- }
-
- return inputColumns;
- }
-
- /**
- * Convert the input row to the correct Arrow format
- *
- * @param row input row
- * @param sourceVectors vectors (buffers) that hold the row
- * @param curRowIndex current row index to use
- * @param statsMap column stats map
- * @return row size
- */
- private float convertRowToArrow(
- Map row,
- VectorSchemaRoot sourceVectors,
- int curRowIndex,
- Map statsMap) {
- // Verify all the input columns are valid
- Set inputColumnNames = verifyInputColumns(row);
-
- float rowBufferSize = 0F;
- for (Map.Entry entry : row.entrySet()) {
- rowBufferSize += 0.125; // 1/8 for null value bitmap
- String columnName = this.formatColumnName(entry.getKey());
- Object value = entry.getValue();
- Field field = this.fields.get(columnName);
- Utils.assertNotNull("Arrow column field", field);
- FieldVector vector = sourceVectors.getVector(field);
- Utils.assertNotNull("Arrow column vector", vector);
- RowBufferStats stats = statsMap.get(columnName);
- Utils.assertNotNull("Arrow column stats", stats);
- ColumnLogicalType logicalType =
- ColumnLogicalType.valueOf(field.getMetadata().get(COLUMN_LOGICAL_TYPE));
- ColumnPhysicalType physicalType =
- ColumnPhysicalType.valueOf(field.getMetadata().get(COLUMN_PHYSICAL_TYPE));
-
- if (value == null) {
- if (!field.getFieldType().isNullable()) {
- throw new SFException(
- ErrorCode.INVALID_ROW, columnName, "Passed null to non nullable field");
- }
- insertNull(vector, stats, curRowIndex);
- } else {
- switch (logicalType) {
- case FIXED:
- if (!field.getMetadata().get(COLUMN_SCALE).equals("0")
- || physicalType == ColumnPhysicalType.SB16) {
- int scale =
- DataValidationUtil.validateAndParseInteger(field.getMetadata().get(COLUMN_SCALE));
- BigDecimal bigDecimalValue = DataValidationUtil.validateAndParseBigDecimal(value);
-
- // vector.setSafe requires the BigDecimal input scale explicitly match its scale
- bigDecimalValue = bigDecimalValue.setScale(scale);
- ((DecimalVector) vector).setSafe(curRowIndex, bigDecimalValue);
- BigInteger intRep =
- bigDecimalValue
- .multiply(BigDecimal.valueOf(Power10.intTable[scale]))
- .toBigInteger();
- stats.addIntValue(intRep);
- rowBufferSize += 16;
- } else {
- switch (physicalType) {
- case SB1:
- byte byteValue = DataValidationUtil.validateAndParseByte(value);
- ((TinyIntVector) vector).setSafe(curRowIndex, byteValue);
- stats.addIntValue(BigInteger.valueOf(byteValue));
- rowBufferSize += 1;
- break;
- case SB2:
- short shortValue = DataValidationUtil.validateAndParseShort(value);
- ((SmallIntVector) vector).setSafe(curRowIndex, shortValue);
- stats.addIntValue(BigInteger.valueOf(shortValue));
- rowBufferSize += 2;
- break;
- case SB4:
- int intVal = DataValidationUtil.validateAndParseInteger(value);
- ((IntVector) vector).setSafe(curRowIndex, intVal);
- stats.addIntValue(BigInteger.valueOf(intVal));
- rowBufferSize += 4;
- break;
- case SB8:
- long longValue = DataValidationUtil.validateAndParseLong(value);
- ((BigIntVector) vector).setSafe(curRowIndex, longValue);
- stats.addIntValue(BigInteger.valueOf(longValue));
- rowBufferSize += 8;
- break;
- default:
- throw new SFException(ErrorCode.UNKNOWN_DATA_TYPE, logicalType, physicalType);
- }
- }
- break;
- case ANY:
- case CHAR:
- case TEXT:
- {
- String maxLengthString = field.getMetadata().get(COLUMN_CHAR_LENGTH);
- String str =
- DataValidationUtil.validateAndParseString(
- value,
- Optional.ofNullable(maxLengthString)
- .map(s -> DataValidationUtil.validateAndParseInteger(maxLengthString)));
- Text text = new Text(str);
- ((VarCharVector) vector).setSafe(curRowIndex, text);
- stats.addStrValue(str);
- rowBufferSize += text.getBytes().length;
- break;
- }
- case OBJECT:
- {
- String str = DataValidationUtil.validateAndParseObject(value);
- Text text = new Text(str);
- ((VarCharVector) vector).setSafe(curRowIndex, text);
- rowBufferSize += text.getBytes().length;
- break;
- }
- case ARRAY:
- case VARIANT:
- {
- String str = DataValidationUtil.validateAndParseVariant(value);
- Text text = new Text(str);
- ((VarCharVector) vector).setSafe(curRowIndex, text);
- rowBufferSize += text.getBytes().length;
- break;
- }
- case TIMESTAMP_LTZ:
- case TIMESTAMP_NTZ:
- switch (physicalType) {
- case SB8:
- {
- BigIntVector bigIntVector = (BigIntVector) vector;
- BigInteger timeInScale =
- DataValidationUtil.validateAndParseTime(value, field.getMetadata());
- bigIntVector.setSafe(curRowIndex, timeInScale.longValue());
- stats.addIntValue(timeInScale);
- rowBufferSize += 8;
- break;
- }
- case SB16:
- {
- StructVector structVector = (StructVector) vector;
- BigIntVector epochVector =
- (BigIntVector) structVector.getChild(FIELD_EPOCH_IN_SECONDS);
- IntVector fractionVector =
- (IntVector) structVector.getChild(FIELD_FRACTION_IN_NANOSECONDS);
- rowBufferSize += 0.25; // for children vector's null value
- structVector.setIndexDefined(curRowIndex);
-
- TimestampWrapper timestampWrapper =
- DataValidationUtil.validateAndParseTimestampNtzSb16(
- value, field.getMetadata());
- epochVector.setSafe(curRowIndex, timestampWrapper.getEpoch());
- fractionVector.setSafe(curRowIndex, timestampWrapper.getFraction());
- rowBufferSize += 12;
- stats.addIntValue(timestampWrapper.getTimeInScale());
- break;
- }
- default:
- throw new SFException(ErrorCode.UNKNOWN_DATA_TYPE, logicalType, physicalType);
- }
- break;
- case TIMESTAMP_TZ:
- switch (physicalType) {
- case SB8:
- {
- StructVector structVector = (StructVector) vector;
- BigIntVector epochVector =
- (BigIntVector) structVector.getChild(FIELD_EPOCH_IN_SECONDS);
- IntVector timezoneVector = (IntVector) structVector.getChild(FIELD_TIME_ZONE);
-
- rowBufferSize += 0.25; // for children vector's null value
- structVector.setIndexDefined(curRowIndex);
-
- TimestampWrapper timestampWrapper =
- DataValidationUtil.validateAndParseTimestampTz(value, field.getMetadata());
- epochVector.setSafe(curRowIndex, timestampWrapper.getTimeInScale().longValue());
- timezoneVector.setSafe(
- curRowIndex,
- timestampWrapper
- .getTimeZoneIndex()
- .orElseThrow(
- () ->
- new SFException(
- ErrorCode.INVALID_ROW,
- value,
- "Unable to parse timezone for TIMESTAMP_TZ column")));
- rowBufferSize += 12;
- BigInteger timeInBinary =
- timestampWrapper
- .getSfTimestamp()
- .orElseThrow(
- () ->
- new SFException(
- ErrorCode.INVALID_ROW,
- value,
- "Unable to parse timestamp for TIMESTAMP_TZ column"))
- .toBinary(Integer.parseInt(field.getMetadata().get(COLUMN_SCALE)), true);
- stats.addIntValue(timeInBinary);
- break;
- }
- case SB16:
- {
- StructVector structVector = (StructVector) vector;
- BigIntVector epochVector =
- (BigIntVector) structVector.getChild(FIELD_EPOCH_IN_SECONDS);
- IntVector fractionVector =
- (IntVector) structVector.getChild(FIELD_FRACTION_IN_NANOSECONDS);
- IntVector timezoneVector = (IntVector) structVector.getChild(FIELD_TIME_ZONE);
-
- rowBufferSize += 0.375; // for children vector's null value
- structVector.setIndexDefined(curRowIndex);
-
- TimestampWrapper timestampWrapper =
- DataValidationUtil.validateAndParseTimestampTz(value, field.getMetadata());
- epochVector.setSafe(curRowIndex, timestampWrapper.getEpoch());
- fractionVector.setSafe(curRowIndex, timestampWrapper.getFraction());
- timezoneVector.setSafe(
- curRowIndex,
- timestampWrapper
- .getTimeZoneIndex()
- .orElseThrow(
- () ->
- new SFException(
- ErrorCode.INVALID_ROW,
- value,
- "Unable to parse timezone for TIMESTAMP_TZ column")));
- rowBufferSize += 16;
- BigInteger timeInBinary =
- timestampWrapper
- .getSfTimestamp()
- .orElseThrow(
- () ->
- new SFException(
- ErrorCode.INVALID_ROW,
- value,
- "Unable to parse timestamp for TIMESTAMP_TZ column"))
- .toBinary(Integer.parseInt(field.getMetadata().get(COLUMN_SCALE)), true);
- stats.addIntValue(timeInBinary);
- break;
- }
- default:
- throw new SFException(ErrorCode.UNKNOWN_DATA_TYPE, logicalType, physicalType);
- }
- break;
- case DATE:
- {
- DateDayVector dateDayVector = (DateDayVector) vector;
- // Expect days past the epoch
- int intValue = DataValidationUtil.validateAndParseDate(value);
- dateDayVector.setSafe(curRowIndex, intValue);
- stats.addIntValue(BigInteger.valueOf(intValue));
- rowBufferSize += 4;
- break;
- }
- case TIME:
- switch (physicalType) {
- case SB4:
- {
- BigInteger timeInScale =
- DataValidationUtil.validateAndParseTime(value, field.getMetadata());
- stats.addIntValue(timeInScale);
- ((IntVector) vector).setSafe(curRowIndex, timeInScale.intValue());
- stats.addIntValue(timeInScale);
- rowBufferSize += 4;
- break;
- }
- case SB8:
- {
- BigInteger timeInScale =
- DataValidationUtil.validateAndParseTime(value, field.getMetadata());
- ((BigIntVector) vector).setSafe(curRowIndex, timeInScale.longValue());
- stats.addIntValue(timeInScale);
- rowBufferSize += 8;
- break;
- }
- default:
- throw new SFException(ErrorCode.UNKNOWN_DATA_TYPE, logicalType, physicalType);
- }
- break;
- case BOOLEAN:
- {
- int intValue = DataValidationUtil.validateAndParseBoolean(value);
- ((BitVector) vector).setSafe(curRowIndex, intValue);
- rowBufferSize += 0.125;
- stats.addIntValue(BigInteger.valueOf(intValue));
- break;
- }
- case BINARY:
- byte[] bytes = DataValidationUtil.validateAndParseBinary(value);
- ((VarBinaryVector) vector).setSafe(curRowIndex, bytes);
- stats.addStrValue(new String(bytes, StandardCharsets.UTF_8));
- rowBufferSize += bytes.length;
- break;
- case REAL:
- double doubleValue = DataValidationUtil.validateAndParseReal(value);
- ((Float8Vector) vector).setSafe(curRowIndex, doubleValue);
- stats.addRealValue(doubleValue);
- rowBufferSize += 8;
- break;
- default:
- throw new SFException(ErrorCode.UNKNOWN_DATA_TYPE, logicalType, physicalType);
- }
- }
- }
-
- // Insert nulls to the columns that doesn't show up in the input
- for (String columnName : Sets.difference(this.fields.keySet(), inputColumnNames)) {
- insertNull(
- sourceVectors.getVector(this.fields.get(columnName)),
- statsMap.get(columnName),
- curRowIndex);
- }
-
- return rowBufferSize;
- }
-
- /** Helper function to insert null value to a field vector */
- private void insertNull(FieldVector vector, RowBufferStats stats, int curRowIndex) {
- if (BaseFixedWidthVector.class.isAssignableFrom(vector.getClass())) {
- ((BaseFixedWidthVector) vector).setNull(curRowIndex);
- } else if (BaseVariableWidthVector.class.isAssignableFrom(vector.getClass())) {
- ((BaseVariableWidthVector) vector).setNull(curRowIndex);
- } else if (vector instanceof StructVector) {
- ((StructVector) vector).setNull(curRowIndex);
- ((StructVector) vector)
- .getChildrenFromFields()
- .forEach(
- child -> {
- ((BaseFixedWidthVector) child).setNull(curRowIndex);
- });
- } else {
- throw new SFException(ErrorCode.INTERNAL_ERROR, "Unexpected FieldType");
- }
- stats.incCurrentNullCount();
- }
-}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/BinaryStringUtils.java b/src/main/java/net/snowflake/ingest/streaming/internal/BinaryStringUtils.java
new file mode 100644
index 000000000..21abed8f5
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/BinaryStringUtils.java
@@ -0,0 +1,44 @@
+package net.snowflake.ingest.streaming.internal;
+
+import java.nio.ByteBuffer;
+import org.apache.commons.codec.binary.Hex;
+
+public class BinaryStringUtils {
+ private static final int MAX_LOB_LEN = 32;
+
+ /** Returns the number of unicode code points in a string */
+ static int unicodeCharactersCount(String s) {
+ return s.codePointCount(0, s.length());
+ }
+
+ /**
+ * Truncate an array of bytes to 32 bytes and optionally increment the last byte(s). More the one
+ * byte can be incremented in case it overflows.
+ */
+ static String truncateBytesAsHex(byte[] bytes, boolean truncateUp) {
+ if (bytes.length <= MAX_LOB_LEN) {
+ return Hex.encodeHexString(bytes);
+ }
+
+ // In order not to mutate the original byte array, let's make a copy
+ byte[] result = new byte[MAX_LOB_LEN];
+ System.arraycopy(bytes, 0, result, 0, MAX_LOB_LEN);
+
+ // Round the least significant byte(s) up
+ if (truncateUp) {
+ int idx;
+ for (idx = MAX_LOB_LEN - 1; idx >= 0; idx--) {
+ // increment the current byte, if there was no overflow, we can stop
+ if (++result[idx] != 0) {
+ break;
+ }
+ }
+ // Whole prefix has overflown, return infinity
+ if (idx == -1) {
+ return "Z";
+ }
+ }
+
+ return Hex.encodeHexString(ByteBuffer.wrap(result, 0, MAX_LOB_LEN));
+ }
+}
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 9cb366a2c..b88090e01 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java
@@ -8,23 +8,30 @@
import static net.snowflake.ingest.utils.Constants.BLOB_CHUNK_METADATA_LENGTH_SIZE_IN_BYTES;
import static net.snowflake.ingest.utils.Constants.BLOB_EXTENSION_TYPE;
import static net.snowflake.ingest.utils.Constants.BLOB_FILE_SIZE_SIZE_IN_BYTES;
-import static net.snowflake.ingest.utils.Constants.BLOB_FORMAT_VERSION;
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.Constants.COMPRESS_BLOB_TWICE;
import static net.snowflake.ingest.utils.Utils.toByteArray;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
+import java.util.ArrayList;
import java.util.List;
-import java.util.zip.GZIPOutputStream;
-import javax.xml.bind.DatatypeConverter;
+import java.util.zip.CRC32;
+import javax.crypto.BadPaddingException;
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.NoSuchPaddingException;
+import net.snowflake.ingest.utils.Constants;
+import net.snowflake.ingest.utils.Cryptor;
import net.snowflake.ingest.utils.Logging;
+import net.snowflake.ingest.utils.Pair;
+import org.apache.commons.codec.binary.Hex;
/**
* Build a single blob file that contains file header plus data. The header will be a
@@ -39,8 +46,8 @@
* variable size chunks metadata in json format
*
*
- * After the metadata, it will be one or more chunks of variable size Arrow data, and each chunk
- * will be encrypted and compressed separately.
+ *
After the metadata, it will be one or more chunks of variable size Parquet data, and each
+ * chunk will be encrypted and compressed separately.
*/
class BlobBuilder {
@@ -48,61 +55,139 @@ class BlobBuilder {
private static final Logging logger = new Logging(BlobBuilder.class);
/**
- * Gzip compress the given chunk data
+ * Builds blob.
*
- * @param filePath blob file full path
- * @param chunkData uncompressed chunk data
- * @return compressed chunk data
- * @throws IOException
+ * @param filePath Path of the destination file 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 bdecVersion version of blob
+ * @return {@link Blob} data
*/
- static byte[] compress(String filePath, ByteArrayOutputStream chunkData) throws IOException {
- int uncompressedSize = chunkData.size();
- ByteArrayOutputStream compressedOutputStream = new ByteArrayOutputStream(uncompressedSize);
- try (GZIPOutputStream gzipOutputStream = new GZIPOutputStream(compressedOutputStream, true)) {
- gzipOutputStream.write(chunkData.toByteArray());
- }
- int firstCompressedSize = compressedOutputStream.size();
-
- // Based on current experiment, compressing twice will give us the best compression
- // ratio and compression time combination
- int doubleCompressedSize = 0;
- if (COMPRESS_BLOB_TWICE) {
- ByteArrayOutputStream doubleCompressedOutputStream =
- new ByteArrayOutputStream(firstCompressedSize);
- try (GZIPOutputStream doubleGzipOutputStream =
- new GZIPOutputStream(doubleCompressedOutputStream, true)) {
- doubleGzipOutputStream.write(compressedOutputStream.toByteArray());
+ static Blob constructBlobAndMetadata(
+ String filePath, List>> blobData, Constants.BdecVersion bdecVersion)
+ throws IOException, NoSuchPaddingException, NoSuchAlgorithmException,
+ InvalidAlgorithmParameterException, InvalidKeyException, IllegalBlockSizeException,
+ BadPaddingException {
+ List chunksMetadataList = new ArrayList<>();
+ List chunksDataList = new ArrayList<>();
+ long curDataSize = 0L;
+ CRC32 crc = new CRC32();
+
+ // TODO: channels with different schema can't be combined even if they belongs to same table
+ for (List> channelsDataPerTable : blobData) {
+ ChannelFlushContext firstChannelFlushContext =
+ channelsDataPerTable.get(0).getChannelContext();
+
+ Flusher flusher = channelsDataPerTable.get(0).createFlusher();
+ Flusher.SerializationResult serializedChunk =
+ 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();
+
+ // 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(encryptedCompressedChunkData, paddedChunkLength);
+ int encryptedCompressedChunkDataSize = encryptedCompressedChunkData.length;
+
+ // Create chunk metadata
+ long startOffset = curDataSize;
+ ChunkMetadata chunkMetadata =
+ ChunkMetadata.builder()
+ .setOwningTableFromChannelContext(firstChannelFlushContext)
+ // 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
+ // decompression and md5 calculation on server side.
+ .setChunkLength(paddedChunkLength)
+ .setUncompressedChunkLength((int) serializedChunk.chunkEstimatedUncompressedSize)
+ .setChannelList(serializedChunk.channelsMetadataList)
+ .setChunkMD5(md5)
+ .setEncryptionKeyId(firstChannelFlushContext.getEncryptionKeyId())
+ .setEpInfo(
+ AbstractRowBuffer.buildEpInfoFromStats(
+ serializedChunk.rowCount, serializedChunk.columnEpStatsMapCombined))
+ .setFirstInsertTimeInMs(serializedChunk.chunkMinMaxInsertTimeInMs.getFirst())
+ .setLastInsertTimeInMs(serializedChunk.chunkMinMaxInsertTimeInMs.getSecond())
+ .build();
+
+ // Add chunk metadata and data to the list
+ chunksMetadataList.add(chunkMetadata);
+ chunksDataList.add(encryptedCompressedChunkData);
+ curDataSize += encryptedCompressedChunkDataSize;
+ crc.update(encryptedCompressedChunkData, 0, encryptedCompressedChunkDataSize);
+
+ logger.logInfo(
+ "Finish building chunk in blob={}, table={}, rowCount={}, startOffset={},"
+ + " estimatedUncompressedSize={}, paddedChunkLength={}, encryptedCompressedSize={},"
+ + " bdecVersion={}",
+ filePath,
+ firstChannelFlushContext.getFullyQualifiedTableName(),
+ serializedChunk.rowCount,
+ startOffset,
+ serializedChunk.chunkEstimatedUncompressedSize,
+ paddedChunkLength,
+ encryptedCompressedChunkDataSize,
+ bdecVersion);
}
- doubleCompressedSize = doubleCompressedOutputStream.size();
- compressedOutputStream = doubleCompressedOutputStream;
}
- logger.logDebug(
- "Finish compressing chunk in blob={}, uncompressedSize={}, firstCompressedSize={},"
- + " doubleCompressedSize={}",
- filePath,
- uncompressedSize,
- firstCompressedSize,
- doubleCompressedSize);
+ // Build blob file bytes
+ byte[] blobBytes =
+ buildBlob(chunksMetadataList, chunksDataList, crc.getValue(), curDataSize, bdecVersion);
+ return new Blob(blobBytes, chunksMetadataList, new BlobStats());
+ }
- return compressedOutputStream.toByteArray();
+ /**
+ * Pad the compressed data for encryption. Encryption needs padding to the
+ * ENCRYPTION_ALGORITHM_BLOCK_SIZE_BYTES to align with decryption on the Snowflake query path
+ * starting from this chunk offset.
+ *
+ * @param chunkData uncompressed chunk data
+ * @param blockSizeToAlignTo block size to align to for encryption
+ * @return padded compressed chunk data, aligned to blockSizeToAlignTo, and actual length of
+ * compressed data before padding at the end
+ * @throws IOException
+ */
+ static Pair padChunk(ByteArrayOutputStream chunkData, int blockSizeToAlignTo)
+ throws IOException {
+ int actualSize = chunkData.size();
+ int paddingSize = blockSizeToAlignTo - actualSize % blockSizeToAlignTo;
+ chunkData.write(new byte[paddingSize]);
+ return new Pair<>(chunkData.toByteArray(), actualSize);
}
/**
- * Build the blob file
+ * Build the blob file bytes
*
* @param chunksMetadataList List of chunk metadata
* @param chunksDataList List of chunk data
* @param chunksChecksum Checksum for the chunk data portion
* @param chunksDataSize Total size for the chunk data portion after compression
+ * @param bdecVersion BDEC file version
* @return The blob file as a byte array
* @throws JsonProcessingException
*/
- static byte[] build(
+ static byte[] buildBlob(
List chunksMetadataList,
List chunksDataList,
long chunksChecksum,
- long chunksDataSize)
+ long chunksDataSize,
+ Constants.BdecVersion bdecVersion)
throws IOException {
byte[] chunkMetadataListInBytes = MAPPER.writeValueAsBytes(chunksMetadataList);
@@ -120,17 +205,17 @@ static byte[] build(
ByteArrayOutputStream blob = new ByteArrayOutputStream();
if (!BLOB_NO_HEADER) {
blob.write(BLOB_EXTENSION_TYPE.getBytes());
- blob.write(BLOB_FORMAT_VERSION);
+ blob.write(bdecVersion.toByte());
blob.write(toByteArray(metadataSize + chunksDataSize));
blob.write(toByteArray(chunksChecksum));
blob.write(toByteArray(chunkMetadataListInBytes.length));
blob.write(chunkMetadataListInBytes);
}
- for (byte[] arrowData : chunksDataList) {
- blob.write(arrowData);
+ for (byte[] chunkData : chunksDataList) {
+ blob.write(chunkData);
}
- // We need to update the start offset for the EP and Arrow data in the request since
+ // We need to update the start offset for the EP and Parquet data in the request since
// some metadata was added at the beginning
for (ChunkMetadata chunkMetadata : chunksMetadataList) {
chunkMetadata.advanceStartOffset(metadataSize);
@@ -165,6 +250,19 @@ static String computeMD5(byte[] data, int length) throws NoSuchAlgorithmExceptio
MessageDigest.getInstance("MD5");
md.update(data, 0, length);
byte[] digest = md.digest();
- return DatatypeConverter.printHexBinary(digest).toLowerCase();
+ return Hex.encodeHexString(digest);
+ }
+
+ /** Blob data to store in a file and register by server side. */
+ static class Blob {
+ final byte[] blobBytes;
+ final List chunksMetadataList;
+ final BlobStats blobStats;
+
+ Blob(byte[] blobBytes, List chunksMetadataList, BlobStats blobStats) {
+ this.blobBytes = blobBytes;
+ this.chunksMetadataList = chunksMetadataList;
+ this.blobStats = blobStats;
+ }
}
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/BlobMetadata.java b/src/main/java/net/snowflake/ingest/streaming/internal/BlobMetadata.java
index e1dbdf764..fb366db9c 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/BlobMetadata.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/BlobMetadata.java
@@ -4,19 +4,52 @@
package net.snowflake.ingest.streaming.internal;
+import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.annotations.VisibleForTesting;
import java.util.List;
+import net.snowflake.ingest.utils.Constants;
+import net.snowflake.ingest.utils.ParameterProvider;
/** Metadata for a blob that sends to Snowflake as part of the register blob request */
class BlobMetadata {
private final String path;
private final String md5;
+ private final Constants.BdecVersion bdecVersion;
private final List chunks;
+ private final BlobStats blobStats;
+ private final boolean spansMixedTables;
- BlobMetadata(String path, String md5, List chunks) {
+ // used for testing only
+ @VisibleForTesting
+ BlobMetadata(String path, String md5, List chunks, BlobStats blobStats) {
+ this(
+ path,
+ md5,
+ ParameterProvider.BLOB_FORMAT_VERSION_DEFAULT,
+ chunks,
+ blobStats,
+ chunks == null ? false : chunks.size() > 1);
+ }
+
+ BlobMetadata(
+ String path,
+ String md5,
+ Constants.BdecVersion bdecVersion,
+ List chunks,
+ BlobStats blobStats,
+ boolean spansMixedTables) {
this.path = path;
this.md5 = md5;
+ this.bdecVersion = bdecVersion;
this.chunks = chunks;
+ this.blobStats = blobStats;
+ this.spansMixedTables = spansMixedTables;
+ }
+
+ @JsonIgnore
+ Constants.BdecVersion getVersion() {
+ return bdecVersion;
}
@JsonProperty("path")
@@ -33,4 +66,30 @@ String getMD5() {
List getChunks() {
return this.chunks;
}
+
+ @JsonProperty("bdec_version")
+ byte getVersionByte() {
+ return bdecVersion.toByte();
+ }
+
+ @JsonProperty("blob_stats")
+ BlobStats getBlobStats() {
+ return this.blobStats;
+ }
+
+ @JsonProperty("spans_mixed_tables")
+ boolean getSpansMixedTables() {
+ return this.spansMixedTables;
+ }
+
+ /** Create {@link BlobMetadata}. */
+ static BlobMetadata createBlobMetadata(
+ String path,
+ String md5,
+ Constants.BdecVersion bdecVersion,
+ List chunks,
+ BlobStats blobStats,
+ boolean spansMixedTables) {
+ return new BlobMetadata(path, md5, bdecVersion, chunks, blobStats, spansMixedTables);
+ }
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/BlobStats.java b/src/main/java/net/snowflake/ingest/streaming/internal/BlobStats.java
new file mode 100644
index 000000000..42994bbb4
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/BlobStats.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) 2023 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming.internal;
+
+import com.codahale.metrics.Timer;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.concurrent.TimeUnit;
+
+/** Latency stats for a blob */
+class BlobStats {
+ // flush duration cannot be calculated in the client sdk, so we pass just the start time
+ private long flushStartMs;
+
+ private long buildDurationMs;
+ private long uploadDurationMs;
+
+ @JsonProperty("build_duration_ms")
+ long getBuildDurationMs() {
+ return this.buildDurationMs;
+ }
+
+ @JsonProperty("build_duration_ms")
+ void setBuildDurationMs(Timer.Context buildLatencyContext) {
+ if (buildLatencyContext != null) {
+ this.buildDurationMs = TimeUnit.NANOSECONDS.toMillis(buildLatencyContext.stop());
+ }
+ }
+
+ @JsonProperty("upload_duration_ms")
+ long getUploadDurationMs() {
+ return this.uploadDurationMs;
+ }
+
+ @JsonProperty("upload_duration_ms")
+ void setUploadDurationMs(Timer.Context uploadLatencyContext) {
+ if (uploadLatencyContext != null) {
+ this.uploadDurationMs = TimeUnit.NANOSECONDS.toMillis(uploadLatencyContext.stop());
+ }
+ }
+
+ @JsonProperty("flush_start_ms")
+ long getFlushStartMs() {
+ return this.flushStartMs;
+ }
+
+ @JsonProperty("flush_start_ms")
+ void setFlushStartMs(long flushStartMs) {
+ this.flushStartMs = flushStartMs;
+ }
+}
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 bb4f39ec7..3926832fb 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelCache.java
@@ -13,12 +13,14 @@
* channels belong to the same table will be stored together in order to combine the channel data
* during flush. The key is a fully qualified table name and the value is a set of channels that
* belongs to this table
+ *
+ * @param type of column data ({@link ParquetChunkData})
*/
-class ChannelCache {
+class ChannelCache {
// Cache to hold all the valid channels, the key for the outer map is FullyQualifiedTableName and
// the key for the inner map is ChannelName
- private ConcurrentHashMap<
- String, ConcurrentHashMap>
+ private final ConcurrentHashMap<
+ String, ConcurrentHashMap>>
cache = new ConcurrentHashMap<>();
/**
@@ -26,15 +28,16 @@ class ChannelCache {
*
* @param channel
*/
- void addChannel(SnowflakeStreamingIngestChannelInternal channel) {
- ConcurrentHashMap channels =
+ void addChannel(SnowflakeStreamingIngestChannelInternal channel) {
+ ConcurrentHashMap> channels =
this.cache.computeIfAbsent(
channel.getFullyQualifiedTableName(), v -> new ConcurrentHashMap<>());
- SnowflakeStreamingIngestChannelInternal oldChannel = channels.put(channel.getName(), channel);
+ SnowflakeStreamingIngestChannelInternal oldChannel =
+ channels.put(channel.getName(), channel);
// Invalidate old channel if it exits to block new inserts and return error to users earlier
if (oldChannel != null) {
- oldChannel.invalidate();
+ oldChannel.invalidate("removed from cache");
}
}
@@ -43,7 +46,7 @@ void addChannel(SnowflakeStreamingIngestChannelInternal channel) {
*
* @return
*/
- Iterator>>
+ Iterator>>>
iterator() {
return this.cache.entrySet().iterator();
}
@@ -57,11 +60,11 @@ void closeAllChannels() {
/** Remove a channel in the channel cache if the channel sequencer matches */
// TODO: background cleaner to cleanup old stale channels that are not closed?
- void removeChannelIfSequencersMatch(SnowflakeStreamingIngestChannelInternal channel) {
+ void removeChannelIfSequencersMatch(SnowflakeStreamingIngestChannelInternal channel) {
cache.computeIfPresent(
channel.getFullyQualifiedTableName(),
(k, v) -> {
- SnowflakeStreamingIngestChannelInternal channelInCache = v.get(channel.getName());
+ SnowflakeStreamingIngestChannelInternal channelInCache = v.get(channel.getName());
// We need to compare the channel sequencer in case the old channel was already been
// removed
return channelInCache != null
@@ -81,12 +84,12 @@ void invalidateChannelIfSequencersMatch(
String channelName,
Long channelSequencer) {
String fullyQualifiedTableName = String.format("%s.%s.%s", dbName, schemaName, tableName);
- ConcurrentHashMap channelsMapPerTable =
+ ConcurrentHashMap> channelsMapPerTable =
cache.get(fullyQualifiedTableName);
if (channelsMapPerTable != null) {
- SnowflakeStreamingIngestChannelInternal channel = channelsMapPerTable.get(channelName);
+ SnowflakeStreamingIngestChannelInternal channel = channelsMapPerTable.get(channelName);
if (channel != null && channel.getChannelSequencer().equals(channelSequencer)) {
- channel.invalidate();
+ channel.invalidate("invalidate with matched sequencer");
}
}
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelData.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelData.java
index 013075c56..81f49d2fd 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelData.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelData.java
@@ -6,21 +6,28 @@
import java.util.HashMap;
import java.util.Map;
+import java.util.function.Supplier;
import net.snowflake.ingest.utils.ErrorCode;
+import net.snowflake.ingest.utils.Pair;
import net.snowflake.ingest.utils.SFException;
-import org.apache.arrow.vector.VectorSchemaRoot;
/**
* Contains the data and metadata returned for each channel flush, which will be used to build the
* blob and register blob request
+ *
+ * @param type of column data (Parquet {@link ParquetChunkData}
*/
-class ChannelData {
+class ChannelData {
private Long rowSequencer;
- private String offsetToken;
- private VectorSchemaRoot vectors;
+ private String endOffsetToken;
+ private String startOffsetToken;
+ private T vectors;
private float bufferSize;
- private SnowflakeStreamingIngestChannelInternal channel;
+ private int rowCount;
private Map columnEps;
+ private Pair minMaxInsertTimeInMs;
+ private ChannelFlushContext channelFlushContext;
+ private Supplier> flusherFactory;
// TODO performance test this vs in place update
/**
@@ -37,11 +44,17 @@ public static Map getCombinedColumnStatsMap(
if (left == null || right == null) {
throw new SFException(ErrorCode.INTERNAL_ERROR, "null column stats");
}
+
if (left.size() != right.size()) {
- throw new SFException(ErrorCode.INTERNAL_ERROR, "Column stats map key mismatch");
+ throw new SFException(
+ ErrorCode.INTERNAL_ERROR,
+ String.format(
+ "Column stats map key size mismatch, left=%d, right=%d, leftKeySet=%s,"
+ + " rightKeySet=%s",
+ left.size(), right.size(), left.keySet(), right.keySet()));
}
- Map result = new HashMap<>();
+ Map result = new HashMap<>();
try {
for (String key : left.keySet()) {
RowBufferStats leftStats = left.get(key);
@@ -49,11 +62,23 @@ public static Map getCombinedColumnStatsMap(
result.put(key, RowBufferStats.getCombinedStats(leftStats, rightStats));
}
} catch (NullPointerException npe) {
- throw new SFException(ErrorCode.INTERNAL_ERROR, "Column stats map key mismatch");
+ throw new SFException(npe, ErrorCode.INTERNAL_ERROR, "Column stats map key mismatch");
}
return result;
}
+ /**
+ * Combines the two paris of min/max insert timestamp together
+ *
+ * @return A new pair which the first element is min(left min, right min) and the second element
+ * is max(left max, right max)
+ */
+ public static Pair getCombinedMinMaxInsertTimeInMs(
+ Pair left, Pair right) {
+ return new Pair<>(
+ Math.min(left.getFirst(), right.getFirst()), Math.max(left.getSecond(), right.getSecond()));
+ }
+
public Map getColumnEps() {
return columnEps;
}
@@ -70,24 +95,36 @@ void setRowSequencer(Long rowSequencer) {
this.rowSequencer = rowSequencer;
}
- String getOffsetToken() {
- return this.offsetToken;
+ String getEndOffsetToken() {
+ return this.endOffsetToken;
+ }
+
+ String getStartOffsetToken() {
+ return this.startOffsetToken;
+ }
+
+ void setEndOffsetToken(String endOffsetToken) {
+ this.endOffsetToken = endOffsetToken;
}
- void setOffsetToken(String offsetToken) {
- this.offsetToken = offsetToken;
+ void setStartOffsetToken(String startOffsetToken) {
+ this.startOffsetToken = startOffsetToken;
}
- VectorSchemaRoot getVectors() {
+ T getVectors() {
return this.vectors;
}
- void setVectors(VectorSchemaRoot vectors) {
+ void setVectors(T vectors) {
this.vectors = vectors;
}
int getRowCount() {
- return this.vectors.getRowCount();
+ return this.rowCount;
+ }
+
+ void setRowCount(int rowCount) {
+ this.rowCount = rowCount;
}
float getBufferSize() {
@@ -98,16 +135,46 @@ void setBufferSize(float bufferSize) {
this.bufferSize = bufferSize;
}
- SnowflakeStreamingIngestChannelInternal getChannel() {
- return this.channel;
+ public ChannelFlushContext getChannelContext() {
+ return channelFlushContext;
+ }
+
+ public void setChannelContext(ChannelFlushContext channelFlushContext) {
+ this.channelFlushContext = channelFlushContext;
+ }
+
+ public Flusher createFlusher() {
+ return flusherFactory.get();
+ }
+
+ public void setFlusherFactory(Supplier> flusherFactory) {
+ this.flusherFactory = flusherFactory;
+ }
+
+ Pair getMinMaxInsertTimeInMs() {
+ return this.minMaxInsertTimeInMs;
}
- void setChannel(SnowflakeStreamingIngestChannelInternal channel) {
- this.channel = channel;
+ void setMinMaxInsertTimeInMs(Pair minMaxInsertTimeInMs) {
+ this.minMaxInsertTimeInMs = minMaxInsertTimeInMs;
}
@Override
public String toString() {
- return this.channel.toString();
+ return "ChannelData{"
+ + "rowSequencer="
+ + rowSequencer
+ + ", endOffsetToken='"
+ + endOffsetToken
+ + ", startOffsetToken='"
+ + startOffsetToken
+ + '\''
+ + ", bufferSize="
+ + bufferSize
+ + ", rowCount="
+ + rowCount
+ + ", channelContext="
+ + channelFlushContext
+ + '}';
}
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelFlushContext.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelFlushContext.java
new file mode 100644
index 000000000..3e5265719
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelFlushContext.java
@@ -0,0 +1,116 @@
+/*
+ * Copyright (c) 2022 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming.internal;
+
+/**
+ * Channel immutable identification and encryption attributes.
+ *
+ * It is shared with the {@link ChannelData} and used for BDEC flush as well.
+ */
+class ChannelFlushContext {
+ private final String name;
+ private final String fullyQualifiedName;
+ private final String dbName;
+ private final String schemaName;
+ private final String tableName;
+ private final String fullyQualifiedTableName;
+
+ // Sequencer for this channel, corresponding to client sequencer at server side because each
+ // connection to a channel at server side will be seen as a connection from a new client
+ private final Long channelSequencer;
+
+ // Data encryption key
+ private final String encryptionKey;
+
+ // Data encryption key id
+ private final Long encryptionKeyId;
+
+ ChannelFlushContext(
+ String name,
+ String dbName,
+ String schemaName,
+ String tableName,
+ Long channelSequencer,
+ String encryptionKey,
+ Long encryptionKeyId) {
+ this.name = name;
+ this.fullyQualifiedName = String.format("%s.%s.%s.%s", 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.channelSequencer = channelSequencer;
+ this.encryptionKey = encryptionKey;
+ this.encryptionKeyId = encryptionKeyId;
+ }
+
+ @Override
+ public String toString() {
+ return "ChannelContext{"
+ + "name='"
+ + getName()
+ + '\''
+ + ", fullyQualifiedName='"
+ + getFullyQualifiedName()
+ + '\''
+ + ", dbName='"
+ + getDbName()
+ + '\''
+ + ", schemaName='"
+ + getSchemaName()
+ + '\''
+ + ", tableName='"
+ + getTableName()
+ + '\''
+ + ", fullyQualifiedTableName='"
+ + getFullyQualifiedTableName()
+ + '\''
+ + ", channelSequencer="
+ + getChannelSequencer()
+ + ", encryptionKey='"
+ + getEncryptionKey()
+ + '\''
+ + ", encryptionKeyId="
+ + getEncryptionKeyId()
+ + '}';
+ }
+
+ String getName() {
+ return name;
+ }
+
+ String getFullyQualifiedName() {
+ return fullyQualifiedName;
+ }
+
+ String getDbName() {
+ return dbName;
+ }
+
+ String getSchemaName() {
+ return schemaName;
+ }
+
+ String getTableName() {
+ return tableName;
+ }
+
+ String getFullyQualifiedTableName() {
+ return fullyQualifiedTableName;
+ }
+
+ Long getChannelSequencer() {
+ return channelSequencer;
+ }
+
+ String getEncryptionKey() {
+ return encryptionKey;
+ }
+
+ Long getEncryptionKeyId() {
+ return encryptionKeyId;
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelMetadata.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelMetadata.java
index 3086634cb..bd072e7c3 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelMetadata.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelMetadata.java
@@ -16,7 +16,8 @@ class ChannelMetadata {
private final String channelName;
private final Long clientSequencer;
private final Long rowSequencer;
- @Nullable private final String offsetToken;
+ @Nullable private final String endOffsetToken;
+ @Nullable private final String startOffsetToken;
static Builder builder() {
return new Builder();
@@ -28,10 +29,11 @@ static class Builder {
private Long clientSequencer;
private Long rowSequencer;
@Nullable private String offsetToken; // offset token could be null
+ @Nullable private String startOffsetToken; // start offset token could be null
- Builder setOwningChannel(SnowflakeStreamingIngestChannelInternal channel) {
- this.channelName = channel.getName();
- this.clientSequencer = channel.getChannelSequencer();
+ Builder setOwningChannelFromContext(ChannelFlushContext channelFlushContext) {
+ this.channelName = channelFlushContext.getName();
+ this.clientSequencer = channelFlushContext.getChannelSequencer();
return this;
}
@@ -45,6 +47,11 @@ Builder setOffsetToken(String offsetToken) {
return this;
}
+ Builder setStartOffsetToken(String startOffsetToken) {
+ this.startOffsetToken = startOffsetToken;
+ return this;
+ }
+
ChannelMetadata build() {
return new ChannelMetadata(this);
}
@@ -58,7 +65,8 @@ private ChannelMetadata(Builder builder) {
this.channelName = builder.channelName;
this.clientSequencer = builder.clientSequencer;
this.rowSequencer = builder.rowSequencer;
- this.offsetToken = builder.offsetToken;
+ this.endOffsetToken = builder.offsetToken;
+ this.startOffsetToken = builder.startOffsetToken;
}
@JsonProperty("channel_name")
@@ -79,6 +87,18 @@ Long getRowSequencer() {
@Nullable
@JsonProperty("offset_token")
String getOffsetToken() {
- return this.offsetToken;
+ return this.endOffsetToken;
+ }
+
+ @Nullable
+ @JsonProperty("start_offset_token")
+ String getStartOffsetToken() {
+ return this.startOffsetToken;
+ }
+
+ @Nullable
+ @JsonProperty("end_offset_token")
+ String getEndOffsetToken() {
+ return this.endOffsetToken;
}
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelRegisterStatus.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelRegisterStatus.java
index 3eaa13c64..5bdc791ec 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelRegisterStatus.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelRegisterStatus.java
@@ -12,6 +12,7 @@
*/
class ChannelRegisterStatus {
private Long statusCode;
+ private String message;
private String channelName;
private Long channelSequencer;
@@ -24,6 +25,15 @@ Long getStatusCode() {
return this.statusCode;
}
+ @JsonProperty("message")
+ public String getMessage() {
+ return message == null ? StreamingIngestResponseCode.getMessageByCode(statusCode) : message;
+ }
+
+ public void setMessage(String message) {
+ this.message = message;
+ }
+
@JsonProperty("channel")
void setChannelName(String channelName) {
this.channelName = channelName;
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelRuntimeState.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelRuntimeState.java
new file mode 100644
index 000000000..8fa7720d3
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelRuntimeState.java
@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) 2022 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming.internal;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+/** Internal state of channel that is used and mutated by both channel and its buffer. */
+class ChannelRuntimeState {
+ // Indicates whether the channel is still valid
+ private volatile boolean isValid;
+
+ // The channel's current start offset token
+ private volatile String startOffsetToken;
+
+ // The channel's current end offset token
+ private volatile String endOffsetToken;
+
+ // The channel's current row sequencer
+ private final AtomicLong rowSequencer;
+
+ // First and last insert time in ms, used for end2end latency measurement
+ private Long firstInsertInMs;
+ private Long lastInsertInMs;
+
+ ChannelRuntimeState(String endOffsetToken, long rowSequencer, boolean isValid) {
+ this.endOffsetToken = endOffsetToken;
+ this.rowSequencer = new AtomicLong(rowSequencer);
+ this.isValid = isValid;
+ }
+
+ /**
+ * Returns whether the channel is in a valid state.
+ *
+ * @return whether the channel is valid
+ */
+ boolean isValid() {
+ return isValid;
+ }
+
+ /** Invalidate the channel that has this state object. */
+ void invalidate() {
+ isValid = false;
+ }
+
+ /** @return current end offset token */
+ String getEndOffsetToken() {
+ return endOffsetToken;
+ }
+
+ /** @return current start offset token */
+ String getStartOffsetToken() {
+ return startOffsetToken;
+ }
+
+ /** @return current offset token after first incrementing it by one. */
+ long incrementAndGetRowSequencer() {
+ return rowSequencer.incrementAndGet();
+ }
+
+ /** @return row sequencer of the channel. */
+ long getRowSequencer() {
+ return rowSequencer.get();
+ }
+
+ /**
+ * Updates the channel's start and end offset token.
+ *
+ * @param startOffsetToken new start offset token of the batch
+ * @param endOffsetToken new end offset token
+ */
+ void updateOffsetToken(String startOffsetToken, String endOffsetToken, int rowCount) {
+ if (rowCount == 0) {
+ this.startOffsetToken = startOffsetToken;
+ }
+ this.endOffsetToken = endOffsetToken;
+ }
+
+ /** Update the insert stats for the current row buffer whenever needed */
+ void updateInsertStats(long currentTimeInMs, int rowCount) {
+ if (rowCount == 0) {
+ this.firstInsertInMs = currentTimeInMs;
+ }
+ this.lastInsertInMs = currentTimeInMs;
+ }
+
+ /** Get the insert timestamp of the first row in the current row buffer */
+ Long getFirstInsertInMs() {
+ return this.firstInsertInMs;
+ }
+
+ /** Get the insert timestamp of the last row in the current row buffer */
+ Long getLastInsertInMs() {
+ return this.lastInsertInMs;
+ }
+}
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 8ae7bf54f..b98782ab9 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusRequest.java
@@ -4,8 +4,6 @@
package net.snowflake.ingest.streaming.internal;
-import com.fasterxml.jackson.annotation.JsonInclude;
-import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import java.util.List;
@@ -29,16 +27,12 @@ static class ChannelStatusRequestDTO {
// Client Sequencer
private final Long clientSequencer;
- // Optional Row Sequencer
- private final Long rowSequencer;
-
ChannelStatusRequestDTO(SnowflakeStreamingIngestChannelInternal channel) {
this.channelName = channel.getName();
this.databaseName = channel.getDBName();
this.schemaName = channel.getSchemaName();
this.tableName = channel.getTableName();
this.clientSequencer = channel.getChannelSequencer();
- this.rowSequencer = channel.getRowSequencer();
}
@JsonProperty("table")
@@ -65,12 +59,6 @@ String getChannelName() {
Long getClientSequencer() {
return clientSequencer;
}
-
- @JsonInclude(Include.NON_EMPTY)
- @JsonProperty("row_sequencer")
- Long getRowSequencer() {
- return rowSequencer;
- }
}
// Optional Request ID. Used for diagnostic purposes.
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusResponse.java
index 0f49d3a5f..96a24581d 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusResponse.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChannelsStatusResponse.java
@@ -8,7 +8,7 @@
import java.util.List;
/** Class used to serialize a response for the channels status endpoint */
-class ChannelsStatusResponse {
+class ChannelsStatusResponse extends StreamingIngestResponse {
static class ChannelStatusResponseDTO {
@@ -75,6 +75,7 @@ void setStatusCode(Long statusCode) {
}
@JsonProperty("status_code")
+ @Override
Long getStatusCode() {
return this.statusCode;
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ChunkMetadata.java b/src/main/java/net/snowflake/ingest/streaming/internal/ChunkMetadata.java
index a7626c4fd..7b42dbc5e 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ChunkMetadata.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ChunkMetadata.java
@@ -15,10 +15,13 @@ class ChunkMetadata {
private final String tableName;
private Long chunkStartOffset;
private final Integer chunkLength;
+ private final Integer uncompressedChunkLength;
private final List channels;
private final String chunkMD5;
private final EpInfo epInfo;
private final Long encryptionKeyId;
+ private final Long firstInsertTimeInMs;
+ private final Long lastInsertTimeInMs;
static Builder builder() {
return new Builder();
@@ -30,16 +33,21 @@ static class Builder {
private String schemaName;
private String tableName;
private Long chunkStartOffset;
- private Integer chunkLength;
+ private Integer chunkLength; // compressedChunkLength
+
+ private Integer uncompressedChunkLength;
+
private List channels;
private String chunkMD5;
private EpInfo epInfo;
private Long encryptionKeyId;
+ private Long firstInsertTimeInMs;
+ private Long lastInsertTimeInMs;
- Builder setOwningTable(SnowflakeStreamingIngestChannelInternal channel) {
- this.dbName = channel.getDBName();
- this.schemaName = channel.getSchemaName();
- this.tableName = channel.getTableName();
+ Builder setOwningTableFromChannelContext(ChannelFlushContext channelFlushContext) {
+ this.dbName = channelFlushContext.getDbName();
+ this.schemaName = channelFlushContext.getSchemaName();
+ this.tableName = channelFlushContext.getTableName();
return this;
}
@@ -58,6 +66,15 @@ Builder setChunkLength(Integer chunkLength) {
return this;
}
+ /**
+ * Currently we send estimated uncompressed size that is close to the actual parquet data size
+ * and mostly about user data but parquet encoding overhead may be slightly different.
+ */
+ public Builder setUncompressedChunkLength(Integer uncompressedChunkLength) {
+ this.uncompressedChunkLength = uncompressedChunkLength;
+ return this;
+ }
+
Builder setChannelList(List channels) {
this.channels = channels;
return this;
@@ -73,6 +90,16 @@ Builder setEncryptionKeyId(Long encryptionKeyId) {
return this;
}
+ Builder setFirstInsertTimeInMs(Long firstInsertTimeInMs) {
+ this.firstInsertTimeInMs = firstInsertTimeInMs;
+ return this;
+ }
+
+ Builder setLastInsertTimeInMs(Long lastInsertTimeInMs) {
+ this.lastInsertTimeInMs = lastInsertTimeInMs;
+ return this;
+ }
+
ChunkMetadata build() {
return new ChunkMetadata(this);
}
@@ -88,16 +115,21 @@ private ChunkMetadata(Builder builder) {
Utils.assertNotNull("chunk MD5", builder.chunkMD5);
Utils.assertNotNull("chunk ep info", builder.epInfo);
Utils.assertNotNull("encryption key id", builder.encryptionKeyId);
+ Utils.assertNotNull("chunk first insert time in ms", builder.firstInsertTimeInMs);
+ Utils.assertNotNull("chunk last insert time in ms", builder.lastInsertTimeInMs);
this.dbName = builder.dbName;
this.schemaName = builder.schemaName;
this.tableName = builder.tableName;
this.chunkStartOffset = builder.chunkStartOffset;
this.chunkLength = builder.chunkLength;
+ this.uncompressedChunkLength = builder.uncompressedChunkLength;
this.channels = builder.channels;
this.chunkMD5 = builder.chunkMD5;
this.epInfo = builder.epInfo;
this.encryptionKeyId = builder.encryptionKeyId;
+ this.firstInsertTimeInMs = builder.firstInsertTimeInMs;
+ this.lastInsertTimeInMs = builder.lastInsertTimeInMs;
}
/**
@@ -134,6 +166,11 @@ Integer getChunkLength() {
return chunkLength;
}
+ @JsonProperty("chunk_length_uncompressed")
+ public Integer getUncompressedChunkLength() {
+ return uncompressedChunkLength;
+ }
+
@JsonProperty("channels")
List getChannels() {
return this.channels;
@@ -153,4 +190,14 @@ EpInfo getEpInfo() {
Long getEncryptionKeyId() {
return this.encryptionKeyId;
}
+
+ @JsonProperty("first_insert_time_in_ms")
+ Long getFirstInsertTimeInMs() {
+ return this.firstInsertTimeInMs;
+ }
+
+ @JsonProperty("last_insert_time_in_ms")
+ Long getLastInsertTimeInMs() {
+ return this.lastInsertTimeInMs;
+ }
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ClientBufferParameters.java b/src/main/java/net/snowflake/ingest/streaming/internal/ClientBufferParameters.java
new file mode 100644
index 000000000..278d4abea
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ClientBufferParameters.java
@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) 2023 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming.internal;
+
+import net.snowflake.ingest.utils.Constants;
+import net.snowflake.ingest.utils.ParameterProvider;
+
+/** Channel's buffer relevant parameters that are set at the owning client level. */
+public class ClientBufferParameters {
+
+ private boolean enableParquetInternalBuffering;
+
+ private long maxChunkSizeInBytes;
+
+ private long maxAllowedRowSizeInBytes;
+
+ private Constants.BdecParquetCompression bdecParquetCompression;
+
+ /**
+ * Private constructor used for test methods
+ *
+ * @param enableParquetInternalBuffering flag whether buffering in internal Parquet buffers is
+ * enabled
+ * @param maxChunkSizeInBytes maximum chunk size in bytes
+ * @param maxAllowedRowSizeInBytes maximum row size in bytes
+ */
+ private ClientBufferParameters(
+ boolean enableParquetInternalBuffering,
+ long maxChunkSizeInBytes,
+ long maxAllowedRowSizeInBytes,
+ Constants.BdecParquetCompression bdecParquetCompression) {
+ this.enableParquetInternalBuffering = enableParquetInternalBuffering;
+ this.maxChunkSizeInBytes = maxChunkSizeInBytes;
+ this.maxAllowedRowSizeInBytes = maxAllowedRowSizeInBytes;
+ this.bdecParquetCompression = bdecParquetCompression;
+ }
+
+ /** @param clientInternal reference to the client object where the relevant parameters are set */
+ public ClientBufferParameters(SnowflakeStreamingIngestClientInternal clientInternal) {
+ this.enableParquetInternalBuffering =
+ clientInternal != null
+ ? clientInternal.getParameterProvider().getEnableParquetInternalBuffering()
+ : ParameterProvider.ENABLE_PARQUET_INTERNAL_BUFFERING_DEFAULT;
+ this.maxChunkSizeInBytes =
+ clientInternal != null
+ ? clientInternal.getParameterProvider().getMaxChunkSizeInBytes()
+ : ParameterProvider.MAX_CHUNK_SIZE_IN_BYTES_DEFAULT;
+ this.maxAllowedRowSizeInBytes =
+ clientInternal != null
+ ? clientInternal.getParameterProvider().getMaxAllowedRowSizeInBytes()
+ : ParameterProvider.MAX_ALLOWED_ROW_SIZE_IN_BYTES_DEFAULT;
+ this.bdecParquetCompression =
+ clientInternal != null
+ ? clientInternal.getParameterProvider().getBdecParquetCompressionAlgorithm()
+ : ParameterProvider.BDEC_PARQUET_COMPRESSION_ALGORITHM_DEFAULT;
+ }
+
+ /**
+ * @param enableParquetInternalBuffering flag whether buffering in internal Parquet buffers is
+ * enabled
+ * @param maxChunkSizeInBytes maximum chunk size in bytes
+ * @param maxAllowedRowSizeInBytes maximum row size in bytes
+ * @return ClientBufferParameters object
+ */
+ public static ClientBufferParameters test_createClientBufferParameters(
+ boolean enableParquetInternalBuffering,
+ long maxChunkSizeInBytes,
+ long maxAllowedRowSizeInBytes,
+ Constants.BdecParquetCompression bdecParquetCompression) {
+ return new ClientBufferParameters(
+ enableParquetInternalBuffering,
+ maxChunkSizeInBytes,
+ maxAllowedRowSizeInBytes,
+ bdecParquetCompression);
+ }
+
+ public boolean getEnableParquetInternalBuffering() {
+ return enableParquetInternalBuffering;
+ }
+
+ public long getMaxChunkSizeInBytes() {
+ return maxChunkSizeInBytes;
+ }
+
+ public long getMaxAllowedRowSizeInBytes() {
+ return maxAllowedRowSizeInBytes;
+ }
+
+ public Constants.BdecParquetCompression getBdecParquetCompression() {
+ return bdecParquetCompression;
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ColumnMetadata.java b/src/main/java/net/snowflake/ingest/streaming/internal/ColumnMetadata.java
index 108035a02..73b75da9e 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ColumnMetadata.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ColumnMetadata.java
@@ -11,6 +11,7 @@
/** Column metadata for each column in a Snowflake table */
class ColumnMetadata {
private String name;
+ private String internalName;
private String type;
private String logicalType;
private String physicalType;
@@ -21,9 +22,16 @@ class ColumnMetadata {
private boolean nullable;
private String collation;
+ /**
+ * The column ordinal is an internal id of the column used by server scanner for the column
+ * identification.
+ */
+ private Integer ordinal;
+
@JsonProperty("name")
void setName(String name) {
this.name = name;
+ this.internalName = LiteralQuoteUtils.unquoteColumnName(name);
}
String getName() {
@@ -111,6 +119,19 @@ boolean getNullable() {
return this.nullable;
}
+ @JsonProperty("ordinal")
+ void setOrdinal(Integer ordinal) {
+ this.ordinal = ordinal;
+ }
+
+ public Integer getOrdinal() {
+ return ordinal;
+ }
+
+ String getInternalName() {
+ return internalName;
+ }
+
@Override
public String toString() {
Map map = new HashMap<>();
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ColumnProperties.java b/src/main/java/net/snowflake/ingest/streaming/internal/ColumnProperties.java
new file mode 100644
index 000000000..856170c57
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ColumnProperties.java
@@ -0,0 +1,60 @@
+package net.snowflake.ingest.streaming.internal;
+
+/**
+ * Class that encapsulates column properties. These are the same properties showed in the output of
+ * SHOW COLUMNS. Note
+ * that this is slightly different than the internal column metadata used elsewhere in this SDK.
+ */
+public class ColumnProperties {
+ private String type;
+
+ private String logicalType;
+
+ private Integer precision;
+
+ private Integer scale;
+
+ private Integer byteLength;
+
+ private Integer length;
+
+ private boolean nullable;
+
+ ColumnProperties(ColumnMetadata columnMetadata) {
+ this.type = columnMetadata.getType();
+ this.logicalType = columnMetadata.getLogicalType();
+ this.precision = columnMetadata.getPrecision();
+ this.scale = columnMetadata.getScale();
+ this.byteLength = columnMetadata.getByteLength();
+ this.length = columnMetadata.getLength();
+ this.nullable = columnMetadata.getNullable();
+ }
+
+ public String getType() {
+ return type;
+ }
+
+ public String getLogicalType() {
+ return logicalType;
+ }
+
+ public Integer getPrecision() {
+ return precision;
+ }
+
+ public Integer getScale() {
+ return scale;
+ }
+
+ public Integer getByteLength() {
+ return byteLength;
+ }
+
+ public Integer getLength() {
+ return length;
+ }
+
+ public boolean isNullable() {
+ return nullable;
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/DataValidationUtil.java b/src/main/java/net/snowflake/ingest/streaming/internal/DataValidationUtil.java
index 190a69250..814423c28 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/DataValidationUtil.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/DataValidationUtil.java
@@ -4,620 +4,912 @@
package net.snowflake.ingest.streaming.internal;
+import static net.snowflake.ingest.streaming.internal.BinaryStringUtils.unicodeCharactersCount;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.module.SimpleModule;
+import com.fasterxml.jackson.databind.ser.std.ToStringSerializer;
import java.math.BigDecimal;
import java.math.BigInteger;
+import java.nio.charset.StandardCharsets;
+import java.time.Instant;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.OffsetDateTime;
+import java.time.OffsetTime;
+import java.time.ZoneId;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
+import java.time.format.DateTimeParseException;
+import java.util.Arrays;
+import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
-import java.util.concurrent.TimeUnit;
-import javax.xml.bind.DatatypeConverter;
+import java.util.function.Supplier;
import net.snowflake.client.jdbc.internal.google.common.collect.Sets;
-import net.snowflake.client.jdbc.internal.snowflake.common.core.SFDate;
-import net.snowflake.client.jdbc.internal.snowflake.common.core.SFTimestamp;
import net.snowflake.client.jdbc.internal.snowflake.common.core.SnowflakeDateTimeFormat;
import net.snowflake.client.jdbc.internal.snowflake.common.util.Power10;
+import net.snowflake.ingest.streaming.internal.serialization.ByteArraySerializer;
+import net.snowflake.ingest.streaming.internal.serialization.ZonedDateTimeSerializer;
import net.snowflake.ingest.utils.ErrorCode;
import net.snowflake.ingest.utils.SFException;
+import org.apache.commons.codec.DecoderException;
+import org.apache.commons.codec.binary.Hex;
/** Utility class for parsing and validating inputs based on Snowflake types */
class DataValidationUtil {
- static final int MAX_STRING_LENGTH = 16777216;
- static final BigInteger MAX_BIGINTEGER = BigInteger.valueOf(10).pow(38);
- static final BigInteger MIN_BIGINTEGER =
- BigInteger.valueOf(-1).multiply(BigInteger.valueOf(10).pow(38));
+
+ /**
+ * Seconds limit used for integer-stored timestamp scale guessing. Value needs to be aligned with
+ * the value from {@link SnowflakeDateTimeFormat#parse}
+ */
+ private static final long SECONDS_LIMIT_FOR_EPOCH = 31536000000L;
+
+ /**
+ * Milliseconds limit used for integer-stored timestamp scale guessing. Value needs to be aligned
+ * with the value from {@link SnowflakeDateTimeFormat#parse}
+ */
+ private static final long MILLISECONDS_LIMIT_FOR_EPOCH = SECONDS_LIMIT_FOR_EPOCH * 1000L;
+
+ /**
+ * Microseconds limit used for integer-stored timestamp scale guessing. Value needs to be aligned
+ * with the value from {@link SnowflakeDateTimeFormat#parse}
+ */
+ private static final long MICROSECONDS_LIMIT_FOR_EPOCH = SECONDS_LIMIT_FOR_EPOCH * 1000000L;
+
+ public static final int BYTES_8_MB = 8 * 1024 * 1024;
+ public static final int BYTES_16_MB = 2 * BYTES_8_MB;
+
+ // TODO SNOW-664249: There is a few-byte mismatch between the value sent by the user and its
+ // server-side representation. Validation leaves a small buffer for this difference.
+ static final int MAX_SEMI_STRUCTURED_LENGTH = BYTES_16_MB - 64;
+
private static final ObjectMapper objectMapper = new ObjectMapper();
- private static final SnowflakeDateTimeFormat snowflakeDateTimeFormatter =
- SnowflakeDateTimeFormat.fromSqlFormat("auto");
+
+ // The version of Jackson we are using does not support serialization of date objects from the
+ // java.time package. Here we define a module with custom java.time serializers. Additionally, we
+ // define custom serializer for byte[] because the Jackson default is to serialize it as
+ // base64-encoded string, and we would like to serialize it as JSON array of numbers.
+ static {
+ SimpleModule module = new SimpleModule();
+ module.addSerializer(byte[].class, new ByteArraySerializer());
+ module.addSerializer(ZonedDateTime.class, new ZonedDateTimeSerializer());
+ module.addSerializer(LocalTime.class, new ToStringSerializer());
+ module.addSerializer(OffsetTime.class, new ToStringSerializer());
+ module.addSerializer(LocalDate.class, new ToStringSerializer());
+ module.addSerializer(LocalDateTime.class, new ToStringSerializer());
+ module.addSerializer(OffsetDateTime.class, new ToStringSerializer());
+ objectMapper.registerModule(module);
+ }
/**
- * Expects string JSON
+ * Validates and parses input as JSON. All types in the object tree must be valid variant types,
+ * see {@link DataValidationUtil#isAllowedSemiStructuredType}.
*
- * @param input the input data, must be able to convert to String
+ * @param input Object to validate
+ * @return JSON tree representing the input
*/
- static String validateAndParseVariant(Object input) {
- String output;
- try {
- output = input.toString();
- } catch (Exception e) {
- throw new SFException(
- e, ErrorCode.INVALID_ROW, input.toString(), "Input column can't be convert to String.");
+ private static JsonNode validateAndParseSemiStructuredAsJsonTree(
+ String columnName, Object input, String snowflakeType, final long insertRowIndex) {
+ if (input instanceof String) {
+ String stringInput = (String) input;
+ verifyValidUtf8(stringInput, columnName, snowflakeType, insertRowIndex);
+ try {
+ return objectMapper.readTree(stringInput);
+ } catch (JsonProcessingException e) {
+ throw valueFormatNotAllowedException(
+ columnName, snowflakeType, "Not a valid JSON", insertRowIndex);
+ }
+ } else if (isAllowedSemiStructuredType(input)) {
+ return objectMapper.valueToTree(input);
}
- if (output.length() > MAX_STRING_LENGTH) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
- String.format(
- "Variant too long. length=%d maxLength=%d", output.length(), MAX_STRING_LENGTH));
- }
- return output;
+ throw typeNotAllowedException(
+ columnName,
+ input.getClass(),
+ snowflakeType,
+ new String[] {
+ "String",
+ "Primitive data types and their arrays",
+ "java.time.*",
+ "List",
+ "Map",
+ "T[]"
+ },
+ insertRowIndex);
}
/**
- * Expects string JSON or JsonNode
+ * Validates and parses input as JSON. All types in the object tree must be valid variant types,
+ * see {@link DataValidationUtil#isAllowedSemiStructuredType}.
*
- * @param input Must be valid string JSON or JsonNode
+ * @param input Object to validate
+ * @param insertRowIndex
+ * @return JSON string representing the input
*/
- static String validateAndParseObject(Object input) {
- String output;
- try {
- JsonNode node = objectMapper.readTree(input.toString());
- output = node.toString();
- } catch (Exception e) {
- throw new SFException(
- e, ErrorCode.INVALID_ROW, input.toString(), "Input column can't be convert to Json");
+ static String validateAndParseVariant(String columnName, Object input, long insertRowIndex) {
+ JsonNode node =
+ validateAndParseSemiStructuredAsJsonTree(columnName, input, "VARIANT", insertRowIndex);
+
+ // Missing nodes are not valid json, ingest them as NULL instead
+ if (node.isMissingNode()) {
+ return null;
}
- if (output.length() > MAX_STRING_LENGTH) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
+ String output = node.toString();
+ int stringLength = output.getBytes(StandardCharsets.UTF_8).length;
+ if (stringLength > MAX_SEMI_STRUCTURED_LENGTH) {
+ throw valueFormatNotAllowedException(
+ columnName,
+ "VARIANT",
String.format(
- "Object too large. length=%d maxLength=%d", output.length(), MAX_STRING_LENGTH));
+ "Variant too long: length=%d maxLength=%d", stringLength, MAX_SEMI_STRUCTURED_LENGTH),
+ insertRowIndex);
}
return output;
}
/**
- * Validates and parses input for TIMESTAMP_NTZ Snowflake type
+ * Validates that passed object is allowed data type for semi-structured columns (i.e. VARIANT,
+ * ARRAY, OBJECT). For non-trivial types like maps, arrays or lists, it recursively traverses the
+ * object tree and validates that all types in the tree are also allowed. Allowed Java types:
*
- * @param input String date in valid format or seconds past the epoch. Accepts fractional seconds
- * with precision up to the column's scale
- * @param metadata
- * @return TimestampWrapper with epoch seconds, fractional seconds, and epoch time in the column
- * scale
+ *
+ * - primitive types (int, long, boolean, ...)
+ *
- String
+ *
- BigInteger
+ *
- BigDecimal
+ *
- LocalTime
+ *
- OffsetTime
+ *
- LocalDate
+ *
- LocalDateTime
+ *
- OffsetDateTime
+ *
- ZonedDateTime
+ *
- Map where T is an allowed semi-structured type
+ *
- List where T is an allowed semi-structured type
+ *
- primitive arrays (char[], int[], ...)
+ *
- T[] where T is an allowed semi-structured type
+ *
+ *
+ * @param o Object to validate
+ * @return If the passed object is allowed for ingestion into semi-structured column
*/
- static TimestampWrapper validateAndParseTimestampNtzSb16(
- Object input, Map metadata) {
- try {
- int scale = Integer.parseInt(metadata.get(ArrowRowBuffer.COLUMN_SCALE));
- String valueString = getStringValue(input);
+ static boolean isAllowedSemiStructuredType(Object o) {
+ // Allow null
+ if (o == null) {
+ return true;
+ }
- long epoch;
- int fraction;
+ // Allow string
+ if (o instanceof String) {
+ return true;
+ }
- SFTimestamp timestamp = snowflakeDateTimeFormatter.parse(valueString);
- // The formatter will not correctly parse numbers with decimal values
- if (timestamp != null) {
- epoch = timestamp.getSeconds().longValue();
- fraction = getFractionFromTimestamp(timestamp);
- } else {
- /*
- * Assume the data is of the form ".". For example "10.5"
- * is interpreted as meaning 10.5 seconds past the epoch. The following logic breaks out and
- * stores the epoch seconds and fractional seconds values separately
- */
- String[] items = valueString.split("\\.");
- if (items.length != 2) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
- String.format("Unable to parse timestamp from value. value=%s", valueString));
- }
- epoch = Long.parseLong(items[0]);
- int l = items.length > 1 ? items[1].length() : 0;
- /* Fraction is in nanoseconds, but Snowflake will error if the fraction gives
- * accuracy greater than the scale
- * */
- fraction = l == 0 ? 0 : Integer.parseInt(items[1]) * (l < 9 ? Power10.intTable[9 - l] : 1);
- }
- if (fraction % Power10.intTable[9 - scale] != 0) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
- String.format(
- "Row specifies accuracy greater than column scale. fraction=%d, scale=%d",
- fraction, scale));
+ // Allow all primitive Java data types
+ if (o instanceof Long
+ || o instanceof Integer
+ || o instanceof Short
+ || o instanceof Byte
+ || o instanceof Float
+ || o instanceof Double
+ || o instanceof Boolean
+ || o instanceof Character) {
+ return true;
+ }
+
+ // Allow BigInteger and BigDecimal
+ if (o instanceof BigInteger || o instanceof BigDecimal) {
+ return true;
+ }
+
+ // Allow supported types from java.time package
+ if (o instanceof java.time.LocalTime
+ || o instanceof OffsetTime
+ || o instanceof LocalDate
+ || o instanceof LocalDateTime
+ || o instanceof ZonedDateTime
+ || o instanceof OffsetDateTime) {
+ return true;
+ }
+
+ // Map is allowed, as long as T is also a supported semi-structured type
+ if (o instanceof Map) {
+ boolean allKeysAreStrings =
+ ((Map) o).keySet().stream().allMatch(x -> x instanceof String);
+ if (!allKeysAreStrings) {
+ return false;
}
- return new TimestampWrapper(epoch, fraction, getTimeInScale(valueString, scale));
- } catch (NumberFormatException e) {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), e.getMessage());
+ boolean allValuesAreAllowed =
+ ((Map) o)
+ .values().stream().allMatch(DataValidationUtil::isAllowedSemiStructuredType);
+ return allValuesAreAllowed;
}
- }
- /**
- * Given a SFTimestamp, get the number of nanoseconds since the last whole second. This
- * corresponds to the fraction component for Timestamp types
- *
- * @param input SFTimestamp
- * @return Timestamp fraction value, the number of nanoseconds since the last whole second
- */
- private static int getFractionFromTimestamp(SFTimestamp input) {
- BigDecimal epochSecondsInNanoSeconds =
- new BigDecimal(input.getSeconds().multiply(BigInteger.valueOf(Power10.intTable[9])));
- return input.getNanosSinceEpoch().subtract(epochSecondsInNanoSeconds).intValue();
+ // Allow arrays of primitive data types
+ if (o instanceof byte[]
+ || o instanceof short[]
+ || o instanceof int[]
+ || o instanceof long[]
+ || o instanceof float[]
+ || o instanceof double[]
+ || o instanceof boolean[]
+ || o instanceof char[]) {
+ return true;
+ }
+
+ // Allow arrays of allowed semi-structured objects
+ if (o.getClass().isArray()) {
+ return Arrays.stream((Object[]) o).allMatch(DataValidationUtil::isAllowedSemiStructuredType);
+ }
+
+ // Allow lists consisting of allowed semi-structured objects
+ if (o instanceof List) {
+ return ((List) o).stream().allMatch(DataValidationUtil::isAllowedSemiStructuredType);
+ }
+
+ // If nothing matches, reject the input
+ return false;
}
/**
- * Validates and parses input for TIMESTAMP_TZ Snowflake type
+ * Validates and parses JSON array. Non-array types are converted into single-element arrays. All
+ * types in the array tree must be valid variant types, see {@link
+ * DataValidationUtil#isAllowedSemiStructuredType}.
*
- * @param input TIMESTAMP_TZ in "2021-01-01 01:00:00 +0100" format
- * @param metadata Column metadata containing scale
- * @return TimestampWrapper with epoch seconds, fractional seconds, and epoch time in the column
- * scale
+ * @param input Object to validate
+ * @param insertRowIndex
+ * @return JSON array representing the input
*/
- static TimestampWrapper validateAndParseTimestampTz(Object input, Map metadata) {
- try {
- if (input instanceof String) {
- String stringInput = (String) input;
- int scale = Integer.parseInt(metadata.get(ArrowRowBuffer.COLUMN_SCALE));
- SFTimestamp timestamp = snowflakeDateTimeFormatter.parse(stringInput);
- if (timestamp == null) {
- throw new SFException(
- ErrorCode.INVALID_ROW, input, "Cannot be converted to a timestamp_tz value");
- }
+ static String validateAndParseArray(String columnName, Object input, long insertRowIndex) {
+ JsonNode jsonNode =
+ validateAndParseSemiStructuredAsJsonTree(columnName, input, "ARRAY", insertRowIndex);
- int fraction = getFractionFromTimestamp(timestamp);
-
- BigDecimal epochTimeAtSecondsScale =
- timestamp.getNanosSinceEpoch().divide(BigDecimal.valueOf(Power10.intTable[9]));
- if (fraction % Power10.intTable[9 - scale] != 0) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
- String.format(
- "Row specifies accuracy greater than column scale. fraction=%d, scale=%d",
- fraction, scale));
- }
- TimestampWrapper wrapper =
- new TimestampWrapper(
- timestamp.getSeconds().longValue(),
- fraction,
- getTimeInScale(getStringValue(epochTimeAtSecondsScale), scale),
- timestamp);
- return wrapper;
- } else {
- throw new SFException(
- ErrorCode.INVALID_ROW, input, "Cannot be converted to a timestamp_tz value");
- }
- } catch (NumberFormatException e) {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), e.getMessage());
+ // Non-array values are ingested as single-element arrays, mimicking the Worksheets behavior
+ if (!jsonNode.isArray()) {
+ jsonNode = objectMapper.createArrayNode().add(jsonNode);
+ }
+
+ String output = jsonNode.toString();
+ // Throw an exception if the size is too large
+ int stringLength = output.getBytes(StandardCharsets.UTF_8).length;
+ if (stringLength > MAX_SEMI_STRUCTURED_LENGTH) {
+ throw valueFormatNotAllowedException(
+ columnName,
+ "ARRAY",
+ String.format(
+ "Array too large. length=%d maxLength=%d", stringLength, MAX_SEMI_STRUCTURED_LENGTH),
+ insertRowIndex);
}
+ return output;
}
/**
- * Validates the input is less than the suppoed maximum allowed string
+ * Validates and parses JSON object. Input is rejected if the value does not represent JSON object
+ * (e.g. String '{}' or Map). All types in the object tree must be valid variant types,
+ * see {@link DataValidationUtil#isAllowedSemiStructuredType}.
*
- * @param input Object to validate and parse to String
- * @param maxLengthOptional Maximum allowed length of the output String, if empty then uses
- * maximum allowed by Snowflake
- * (https://docs.snowflake.com/en/sql-reference/data-types-text.html#varchar)
+ * @param input Object to validate
+ * @param insertRowIndex
+ * @return JSON object representing the input
*/
- static String validateAndParseString(Object input, Optional maxLengthOptional) {
- int maxLength = maxLengthOptional.orElse(MAX_STRING_LENGTH);
- String output = getStringValue(input);
- if (output.length() > maxLength) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
- String.format("String too long. length=%d maxLength=%d", output.length(), maxLength));
+ static String validateAndParseObject(String columnName, Object input, long insertRowIndex) {
+ JsonNode jsonNode =
+ validateAndParseSemiStructuredAsJsonTree(columnName, input, "OBJECT", insertRowIndex);
+ if (!jsonNode.isObject()) {
+ throw valueFormatNotAllowedException(columnName, "OBJECT", "Not an object", insertRowIndex);
+ }
+
+ String output = jsonNode.toString();
+ // Throw an exception if the size is too large
+ int stringLength = output.getBytes(StandardCharsets.UTF_8).length;
+ if (stringLength > MAX_SEMI_STRUCTURED_LENGTH) {
+ throw valueFormatNotAllowedException(
+ columnName,
+ "OBJECT",
+ String.format(
+ "Object too large. length=%d maxLength=%d", stringLength, MAX_SEMI_STRUCTURED_LENGTH),
+ insertRowIndex);
}
return output;
}
/**
- * Returns a BigDecimal representation of the input. Strings of the form "1.23E4" will be treated
- * as being written in * scientific notation (e.g. 1.23 * 10^4)
- *
- * @param input
+ * Converts user input to offset date time, which is the canonical representation of dates and
+ * timestamps.
*/
- static BigDecimal validateAndParseBigDecimal(Object input) {
- BigDecimal output;
- try {
- if (input instanceof BigDecimal) {
- output = (BigDecimal) input;
- } else if (input instanceof BigInteger) {
- output = new BigDecimal((BigInteger) input);
- } else if (input instanceof String) {
- output = stringToBigDecimal((String) input);
- } else {
- output = stringToBigDecimal((input.toString()));
- }
- if (output.toBigInteger().compareTo(MAX_BIGINTEGER) >= 0
- || output.toBigInteger().compareTo(MIN_BIGINTEGER) <= 0) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
- String.format(
- "Number out of representable range, Max=%s, Min=%s",
- MIN_BIGINTEGER, MAX_BIGINTEGER));
- }
- return output;
- } catch (NumberFormatException e) {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), e.getMessage());
+ private static OffsetDateTime inputToOffsetDateTime(
+ String columnName,
+ String typeName,
+ Object input,
+ ZoneId defaultTimezone,
+ final long insertRowIndex) {
+ if (input instanceof OffsetDateTime) {
+ return (OffsetDateTime) input;
}
- }
- static BigDecimal handleScientificNotationError(String input, NumberFormatException e) {
- // Support scientific notation on string input
- String[] splitInput = input.toLowerCase().split("e");
- if (splitInput.length != 2) {
- throw e;
+ if (input instanceof ZonedDateTime) {
+ return ((ZonedDateTime) input).toOffsetDateTime();
}
- return (new BigDecimal(splitInput[0]))
- .multiply(BigDecimal.valueOf(10).pow(Integer.parseInt(splitInput[1])));
- }
- static BigDecimal stringToBigDecimal(String input) {
- try {
- return new BigDecimal(input);
- } catch (NumberFormatException e) {
- return handleScientificNotationError(input, e);
+ if (input instanceof LocalDateTime) {
+ return ((LocalDateTime) input).atZone(defaultTimezone).toOffsetDateTime();
}
- }
- /**
- * Validates the input represents a valid NUMBER(38,0) and returns a BigInteger. Accepts Java
- * number types and strings. Strings of the form "1.23E4" will be treated as being written in
- * scientific notation (e.g. 1.23 * 10^4)
- *
- * @param input
- */
- static BigInteger validateAndParseBigInteger(Object input) {
- try {
- if (input instanceof Double) {
- checkInteger((double) input);
- return new BigDecimal(input.toString()).toBigInteger();
- } else if (input instanceof Float) {
- checkInteger((float) input);
- return new BigDecimal(input.toString()).toBigInteger();
+ if (input instanceof LocalDate) {
+ return ((LocalDate) input).atStartOfDay().atZone(defaultTimezone).toOffsetDateTime();
+ }
+
+ if (input instanceof Instant) {
+ // Just like integer-stored timestamps, instants are always interpreted in UTC
+ return ((Instant) input).atZone(ZoneOffset.UTC).toOffsetDateTime();
+ }
+
+ if (input instanceof String) {
+ String stringInput = ((String) input).trim();
+ {
+ // First, try to parse ZonedDateTime
+ ZonedDateTime zoned = catchParsingError(() -> ZonedDateTime.parse(stringInput));
+ if (zoned != null) {
+ return zoned.toOffsetDateTime();
+ }
}
- BigInteger output;
- if (input instanceof BigInteger) {
- output = (BigInteger) input;
- } else if (input instanceof String) {
- BigDecimal bigDecimal = stringToBigDecimal((String) input);
- if (bigDecimal.stripTrailingZeros().scale() <= 0) {
- return bigDecimal.toBigInteger();
- } else {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), "Invalid integer");
+
+ {
+ // Next, try to parse OffsetDateTime
+ OffsetDateTime offset = catchParsingError(() -> OffsetDateTime.parse(stringInput));
+ if (offset != null) {
+ return offset;
}
- } else {
- output = new BigInteger(input.toString());
}
- if (output.compareTo(MAX_BIGINTEGER) >= 0 || output.compareTo(MIN_BIGINTEGER) <= 0) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
- String.format(
- "Number out of representable range, Max=%s, Min=%s",
- MIN_BIGINTEGER, MAX_BIGINTEGER));
+
+ {
+ // Alternatively, try to parse LocalDateTime
+ LocalDateTime localDateTime = catchParsingError(() -> LocalDateTime.parse(stringInput));
+ if (localDateTime != null) {
+ return localDateTime.atZone(defaultTimezone).toOffsetDateTime();
+ }
}
- return output;
- } catch (NumberFormatException e) {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), e.getMessage());
- }
- }
- static void checkInteger(float input) {
- if (Math.floor(input) != input) {
- throw new SFException(ErrorCode.INVALID_ROW, input, "Value must be integer");
+ {
+ // Alternatively, try to parse LocalDate
+ LocalDate localDate = catchParsingError(() -> LocalDate.parse(stringInput));
+ if (localDate != null) {
+ return localDate.atStartOfDay().atZone(defaultTimezone).toOffsetDateTime();
+ }
+ }
+
+ {
+ // Alternatively, try to parse integer-stored timestamp
+ // Just like in Snowflake, integer-stored timestamps are always in UTC
+ Instant instant = catchParsingError(() -> parseInstantGuessScale(stringInput));
+ if (instant != null) {
+ return instant.atOffset(ZoneOffset.UTC);
+ }
+ }
+
+ // Couldn't parse anything, throw an exception
+ throw valueFormatNotAllowedException(
+ columnName,
+ typeName,
+ "Not a valid value, see"
+ + " https://docs.snowflake.com/en/user-guide/data-load-snowpipe-streaming-overview"
+ + " for the list of supported formats",
+ insertRowIndex);
}
+
+ // Type is not supported, throw an exception
+ throw typeNotAllowedException(
+ columnName,
+ input.getClass(),
+ typeName,
+ new String[] {"String", "LocalDate", "LocalDateTime", "ZonedDateTime", "OffsetDateTime"},
+ insertRowIndex);
}
- static void checkInteger(double input) {
- if (Math.floor(input) != input) {
- throw new SFException(ErrorCode.INVALID_ROW, input, "Value must be integer");
+ private static T catchParsingError(Supplier op) {
+ try {
+ return op.get();
+ } catch (DateTimeParseException | NumberFormatException e) {
+ return null;
}
}
/**
- * Validates input is an integer @see {@link #validateAndParseInteger(Object)} and that it does
- * not exceed the MIN/MAX Short value
+ * Validates and parses input for TIMESTAMP_NTZ, TIMESTAMP_LTZ and TIMEATAMP_TZ Snowflake types.
+ * Allowed Java types:
*
- * @param input
+ *
+ * - String
+ *
- LocalDate
+ *
- LocalDateTime
+ *
- OffsetDateTime
+ *
- ZonedDateTime
+ *
+ *
+ * @param columnName Column name, used in validation error messages
+ * @param input String date in valid format, seconds past the epoch or java.time.* object. Accepts
+ * fractional seconds with precision up to the column's scale
+ * @param scale decimal scale of timestamp 16 byte integer
+ * @param defaultTimezone Input, which does not carry timezone information is going to be
+ * interpreted in the default timezone.
+ * @param trimTimezone Whether timezone information should be removed from the resulting date,
+ * should be true for TIMESTAMP_NTZ columns.
+ * @param insertRowIndex
+ * @return TimestampWrapper
*/
- static short validateAndParseShort(Object input) {
- Integer intValue = validateAndParseInteger(input);
- if (intValue > Short.MAX_VALUE || intValue < Short.MIN_VALUE) {
+ static TimestampWrapper validateAndParseTimestamp(
+ String columnName,
+ Object input,
+ int scale,
+ ZoneId defaultTimezone,
+ boolean trimTimezone,
+ long insertRowIndex) {
+ OffsetDateTime offsetDateTime =
+ inputToOffsetDateTime(columnName, "TIMESTAMP", input, defaultTimezone, insertRowIndex);
+
+ if (trimTimezone) {
+ offsetDateTime = offsetDateTime.withOffsetSameLocal(ZoneOffset.UTC);
+ }
+ if (offsetDateTime.getYear() < 1 || offsetDateTime.getYear() > 9999) {
throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
+ ErrorCode.INVALID_VALUE_ROW,
String.format(
- "Value exceeds min/max short. min=%s, max%s", Short.MIN_VALUE, Short.MAX_VALUE));
+ "Timestamp out of representable inclusive range of years between 1 and 9999,"
+ + " rowIndex:%d",
+ insertRowIndex));
}
- return intValue.shortValue();
+ return new TimestampWrapper(offsetDateTime, scale);
}
/**
- * Validates input is an integer @see {@link #validateAndParseInteger(Object)} and that it does
- * not exceed the MIN/MAX Byte value
+ * Converts input to string, validates that length is less than max allowed string size
+ * https://docs.snowflake.com/en/sql-reference/data-types-text.html#varchar. Allowed data types:
*
- * @param input
+ *
+ * - String
+ *
- Number
+ *
- boolean
+ *
- char
+ *
+ *
+ * @param input Object to validate and parse to String
+ * @param maxLengthOptional Maximum allowed length of the output String, if empty then uses
+ * maximum allowed by Snowflake
+ * (https://docs.snowflake.com/en/sql-reference/data-types-text.html#varchar)
+ * @param insertRowIndex
*/
- static byte validateAndParseByte(Object input) {
- Integer intValue = validateAndParseInteger(input);
- if (intValue > Byte.MAX_VALUE || intValue < Byte.MIN_VALUE) {
- throw new SFException(
- ErrorCode.INVALID_ROW,
- input.toString(),
+ static String validateAndParseString(
+ String columnName, Object input, Optional maxLengthOptional, long insertRowIndex) {
+ String output;
+ if (input instanceof String) {
+ output = (String) input;
+ verifyValidUtf8(output, columnName, "STRING", insertRowIndex);
+ } else if (input instanceof Number) {
+ output = new BigDecimal(input.toString()).stripTrailingZeros().toPlainString();
+ } else if (input instanceof Boolean || input instanceof Character) {
+ output = input.toString();
+ } else {
+ throw typeNotAllowedException(
+ columnName,
+ input.getClass(),
+ "STRING",
+ new String[] {"String", "Number", "boolean", "char"},
+ insertRowIndex);
+ }
+ byte[] utf8Bytes = output.getBytes(StandardCharsets.UTF_8);
+
+ // Strings can never be larger than 16MB
+ if (utf8Bytes.length > BYTES_16_MB) {
+ throw valueFormatNotAllowedException(
+ columnName,
+ "STRING",
String.format(
- "Value exceeds min/max byte. min=%s, max%s", Byte.MIN_VALUE, Byte.MAX_VALUE));
+ "String too long: length=%d bytes maxLength=%d bytes", utf8Bytes.length, BYTES_16_MB),
+ insertRowIndex);
}
- return intValue.byteValue();
+
+ // If max allowed length is specified (e.g. VARCHAR(10)), the number of unicode characters must
+ // not exceed this value
+ maxLengthOptional.ifPresent(
+ maxAllowedCharacters -> {
+ int actualCharacters = unicodeCharactersCount(output);
+ if (actualCharacters > maxAllowedCharacters) {
+ throw valueFormatNotAllowedException(
+ columnName,
+ "STRING",
+ String.format(
+ "String too long: length=%d characters maxLength=%d characters",
+ actualCharacters, maxAllowedCharacters),
+ insertRowIndex);
+ }
+ });
+ return output;
}
/**
- * Validates the input can be represented as an integer with value between Integer.MIN_VALUE and
- * Integer.MAX_VALUE
- *
- * @param input
+ * Returns a BigDecimal representation of the input. Strings of the form "1.23E4" will be treated
+ * as being written in * scientific notation (e.g. 1.23 * 10^4). Does not perform any size
+ * validation. Allowed Java types:
+ * byte, short, int, long
+ * float, double
+ * BigInteger, BigDecimal
+ * String
*/
- static int validateAndParseInteger(Object input) {
- try {
- if (input instanceof Integer) {
- return (int) input;
- } else if (input instanceof Long) {
- if ((long) input > Integer.MAX_VALUE || (long) input < Integer.MIN_VALUE) {
- throw new SFException(
- ErrorCode.INVALID_ROW, input.toString(), "Value greater than max integer");
- }
- return ((Long) input).intValue();
- } else if (input instanceof Double) {
- // Number must be integer
- checkInteger((double) input);
- if (((Number) input).longValue() > Integer.MAX_VALUE
- || ((Number) input).longValue() < Integer.MIN_VALUE) {
- throw new SFException(
- ErrorCode.INVALID_ROW, input.toString(), "Value greater than max integer");
- }
- return ((Number) input).intValue();
- } else if (input instanceof Float) {
- // Number must be integer
- checkInteger((float) input);
-
- if (((Number) input).longValue() > Integer.MAX_VALUE
- || ((Number) input).longValue() < Integer.MIN_VALUE) {
- throw new SFException(
- ErrorCode.INVALID_ROW, input.toString(), "Value greater than max integer");
- }
- return ((Number) input).intValue();
- } else if (input instanceof BigInteger) {
- if (((BigInteger) input).compareTo(BigInteger.valueOf(Integer.MAX_VALUE)) > 0
- || ((BigInteger) input).compareTo(BigInteger.valueOf(Integer.MIN_VALUE)) < 0) {
- throw new SFException(
- ErrorCode.INVALID_ROW, input.toString(), "Value greater than max integer");
- }
- return ((BigInteger) input).intValue();
- } else {
- try {
- return Integer.parseInt(input.toString());
- } catch (NumberFormatException e) {
- Double doubleValue = handleScientificNotationError(input.toString(), e).doubleValue();
- if (Math.floor(doubleValue) == doubleValue) {
- return doubleValue.intValue();
- } else {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), "Value must be integer");
- }
- }
+ static BigDecimal validateAndParseBigDecimal(
+ String columnName, Object input, long insertRowIndex) {
+ if (input instanceof BigDecimal) {
+ return (BigDecimal) input;
+ } else if (input instanceof BigInteger) {
+ return new BigDecimal((BigInteger) input);
+ } else if (input instanceof Byte
+ || input instanceof Short
+ || input instanceof Integer
+ || input instanceof Long) {
+ return BigDecimal.valueOf(((Number) input).longValue());
+ } else if (input instanceof Float || input instanceof Double) {
+ return BigDecimal.valueOf(((Number) input).doubleValue());
+ } else if (input instanceof String) {
+ try {
+ final String stringInput = ((String) input).trim();
+ return new BigDecimal(stringInput);
+ } catch (NumberFormatException e) {
+ throw valueFormatNotAllowedException(
+ columnName, "NUMBER", "Not a valid number", insertRowIndex);
}
- } catch (NumberFormatException err) {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), err.getMessage());
+ } else {
+ throw typeNotAllowedException(
+ columnName,
+ input.getClass(),
+ "NUMBER",
+ new String[] {
+ "int", "long", "byte", "short", "float", "double", "BigDecimal", "BigInteger", "String"
+ },
+ insertRowIndex);
}
}
/**
- * Validates the input can be represented as an integer with value between Long.MIN_VALUE and
- * Long.MAX_VALUE
+ * Returns the number of days between the epoch and the passed date. Allowed Java types:
*
- * @param input
+ *
+ * - String
+ *
- {@link LocalDate}
+ *
- {@link LocalDateTime}
+ *
- {@link OffsetDateTime}
+ *
- {@link ZonedDateTime}
+ *
- {@link Instant}
+ *
*/
- static long validateAndParseLong(Object input) {
- try {
- if (input instanceof Integer) {
- return Long.valueOf((int) input);
- } else if (input instanceof Long) {
- return (long) input;
- } else if (input instanceof Double) {
- // Number must be integer
- checkInteger((double) input);
- BigInteger bigInt = validateAndParseBigInteger(input);
- if (bigInt.compareTo(BigInteger.valueOf(Long.MAX_VALUE)) > 0
- || bigInt.compareTo(BigInteger.valueOf(Long.MIN_VALUE)) < 0) {
- throw new SFException(
- ErrorCode.INVALID_ROW, input.toString(), "Value greater than max long");
- }
- return ((Number) input).longValue();
- } else if (input instanceof Float) {
- // Number must be integer
- checkInteger((float) input);
- BigInteger bigInt = validateAndParseBigInteger(input);
- if (bigInt.compareTo(BigInteger.valueOf(Long.MAX_VALUE)) > 0
- || bigInt.compareTo(BigInteger.valueOf(Long.MIN_VALUE)) < 0) {
- throw new SFException(
- ErrorCode.INVALID_ROW, input.toString(), "Value greater than max long");
- }
- return ((Number) input).longValue();
- } else if (input instanceof BigInteger) {
- if (((BigInteger) input).compareTo(BigInteger.valueOf(Long.MAX_VALUE)) > 0
- || ((BigInteger) input).compareTo(BigInteger.valueOf(Long.MIN_VALUE)) < 0) {
- throw new SFException(
- ErrorCode.INVALID_ROW, input.toString(), "Value greater than max long");
- }
- return ((BigInteger) input).longValue();
- } else {
- try {
- return Long.parseLong(input.toString());
- } catch (NumberFormatException e) {
- Double doubleValue = handleScientificNotationError(input.toString(), e).doubleValue();
- if (Math.floor(doubleValue) == doubleValue) {
- return doubleValue.longValue();
- } else {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), "Value must be integer");
- }
- }
- }
- } catch (NumberFormatException err) {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString(), err.getMessage());
+ static int validateAndParseDate(String columnName, Object input, long insertRowIndex) {
+ OffsetDateTime offsetDateTime =
+ inputToOffsetDateTime(columnName, "DATE", input, ZoneOffset.UTC, insertRowIndex);
+
+ if (offsetDateTime.getYear() < -9999 || offsetDateTime.getYear() > 9999) {
+ throw new SFException(
+ ErrorCode.INVALID_VALUE_ROW,
+ String.format(
+ "Date out of representable inclusive range of years between -9999 and 9999,"
+ + " rowIndex:%d",
+ insertRowIndex));
}
+
+ return Math.toIntExact(offsetDateTime.toLocalDate().toEpochDay());
}
/**
- * Expects input in epoch days or "YYYY-MM-DD" format. Uses @see {@link
- * #validateAndParseInteger(Object)} to validate and parse input as integer
+ * Validates input for data type BINARY. Allowed Java types:
*
- * @param input epoch days or "YYYY-MM-DD" date format
+ *
+ * - byte[]
+ *
- String (hex-encoded)
+ *
+ *
+ * @param input Array to validate
+ * @param maxLengthOptional Max array length, defaults to 8MB, which is the max allowed length for
+ * BINARY column
+ * @param insertRowIndex
+ * @return Validated array
*/
- static int validateAndParseDate(Object input) {
- try {
- return validateAndParseInteger(input);
- } catch (SFException e) {
- // Try parsing the time from a String
- Optional timestamp =
- Optional.ofNullable(snowflakeDateTimeFormatter.parse(input.toString()));
- return timestamp
- .map(
- t ->
- // Adjust nanoseconds past the epoch to match the column scale value
- (int) TimeUnit.MILLISECONDS.toDays(SFDate.fromTimestamp(t).getTime()))
- .orElseThrow(
- () ->
- new SFException(
- ErrorCode.INVALID_ROW, input, "Cannot be converted to a Date value"));
- }
- }
-
- static byte[] validateAndParseBinary(Object input) {
+ static byte[] validateAndParseBinary(
+ String columnName, Object input, Optional maxLengthOptional, long insertRowIndex) {
+ byte[] output;
if (input instanceof byte[]) {
- return (byte[]) input;
+ // byte[] is a mutable object, we need to create a defensive copy to protect against
+ // concurrent modifications of the array, which could lead to mismatch between data
+ // and metadata
+ byte[] originalInputArray = (byte[]) input;
+ output = new byte[originalInputArray.length];
+ System.arraycopy(originalInputArray, 0, output, 0, originalInputArray.length);
+ } else if (input instanceof String) {
+ try {
+ String stringInput = ((String) input).trim();
+ output = Hex.decodeHex(stringInput);
+ } catch (DecoderException e) {
+ throw valueFormatNotAllowedException(
+ columnName, "BINARY", "Not a valid hex string", insertRowIndex);
+ }
+ } else {
+ throw typeNotAllowedException(
+ columnName,
+ input.getClass(),
+ "BINARY",
+ new String[] {"byte[]", "String"},
+ insertRowIndex);
}
- try {
- return DatatypeConverter.parseHexBinary(input.toString());
- } catch (IllegalArgumentException e) {
- throw new SFException(ErrorCode.INVALID_ROW, input, e.getMessage());
+
+ int maxLength = maxLengthOptional.orElse(BYTES_8_MB);
+ if (output.length > maxLength) {
+ throw valueFormatNotAllowedException(
+ columnName,
+ "BINARY",
+ String.format("Binary too long: length=%d maxLength=%d", output.length, maxLength),
+ insertRowIndex);
}
+ return output;
}
/**
- * @param input Seconds past the epoch. or String Time representation
- * @param metadata
+ * Returns the number of units since 00:00, depending on the scale (scale=0: seconds, scale=3:
+ * milliseconds, scale=9: nanoseconds). Allowed Java types:
+ *
+ *
+ * - String
+ *
- {@link LocalTime}
+ *
- {@link OffsetTime}
+ *
*/
- static BigInteger validateAndParseTime(Object input, Map metadata) {
- int scale = Integer.parseInt(metadata.get(ArrowRowBuffer.COLUMN_SCALE));
- return getTimeInScale(getStringValue(input), scale);
+ static BigInteger validateAndParseTime(
+ String columnName, Object input, int scale, long insertRowIndex) {
+ if (input instanceof LocalTime) {
+ LocalTime localTime = (LocalTime) input;
+ return BigInteger.valueOf(localTime.toNanoOfDay()).divide(Power10.sb16Table[9 - scale]);
+ } else if (input instanceof OffsetTime) {
+ return validateAndParseTime(
+ columnName, ((OffsetTime) input).toLocalTime(), scale, insertRowIndex);
+ } else if (input instanceof String) {
+ String stringInput = ((String) input).trim();
+ {
+ // First, try to parse LocalTime
+ LocalTime localTime = catchParsingError(() -> LocalTime.parse(stringInput));
+ if (localTime != null) {
+ return validateAndParseTime(columnName, localTime, scale, insertRowIndex);
+ }
+ }
+
+ {
+ // Alternatively, try to parse OffsetTime
+ OffsetTime offsetTime = catchParsingError((() -> OffsetTime.parse(stringInput)));
+ if (offsetTime != null) {
+ return validateAndParseTime(columnName, offsetTime.toLocalTime(), scale, insertRowIndex);
+ }
+ }
+
+ {
+ // Alternatively, try to parse integer-stored time
+ Instant parsedInstant = catchParsingError(() -> parseInstantGuessScale(stringInput));
+ if (parsedInstant != null) {
+ return validateAndParseTime(
+ columnName,
+ LocalDateTime.ofInstant(parsedInstant, ZoneOffset.UTC).toLocalTime(),
+ scale,
+ insertRowIndex);
+ }
+ }
+
+ throw valueFormatNotAllowedException(
+ columnName,
+ "TIME",
+ "Not a valid time, see"
+ + " https://docs.snowflake.com/en/user-guide/data-load-snowpipe-streaming-overview"
+ + " for the list of supported formats",
+ insertRowIndex);
+
+ } else {
+ throw typeNotAllowedException(
+ columnName,
+ input.getClass(),
+ "TIME",
+ new String[] {"String", "LocalTime", "OffsetTime"},
+ insertRowIndex);
+ }
}
/**
- * Returns the time past the epoch in the given scale.
+ * Attempts to parse integer-stored date from string input. Tries to guess the scale according to
+ * the rules documented at
+ * https://docs.snowflake.com/en/user-guide/date-time-input-output.html#auto-detection-of-integer-stored-date-time-and-timestamp-values.
*
- * @param value
- * @param scale Value between 0 and 9. For seconds scale = 0, for milliseconds scale = 3, for
- * nanoseconds scale = 9
+ * @param input String to parse, must represent a valid long
+ * @return Instant representing the input
+ * @throws NumberFormatException If the input in not a valid long
*/
- static BigInteger getTimeInScale(String value, int scale) {
+ private static Instant parseInstantGuessScale(String input) {
+ BigInteger epochNanos;
try {
- BigDecimal decVal = new BigDecimal(value);
- BigDecimal epochScale = decVal.multiply(BigDecimal.valueOf(Power10.intTable[scale]));
- return epochScale.toBigInteger();
+ long val = Long.parseLong(input);
+
+ if (val > -SECONDS_LIMIT_FOR_EPOCH && val < SECONDS_LIMIT_FOR_EPOCH) {
+ epochNanos = BigInteger.valueOf(val).multiply(Power10.sb16Table[9]);
+ } else if (val > -MILLISECONDS_LIMIT_FOR_EPOCH && val < MILLISECONDS_LIMIT_FOR_EPOCH) {
+ epochNanos = BigInteger.valueOf(val).multiply(Power10.sb16Table[6]);
+ } else if (val > -MICROSECONDS_LIMIT_FOR_EPOCH && val < MICROSECONDS_LIMIT_FOR_EPOCH) {
+ epochNanos = BigInteger.valueOf(val).multiply(Power10.sb16Table[3]);
+ } else {
+ epochNanos = BigInteger.valueOf(val);
+ }
} catch (NumberFormatException e) {
- // Try parsing the time from a String
- Optional timestamp =
- Optional.ofNullable(snowflakeDateTimeFormatter.parse(value));
- return timestamp
- .map(
- t ->
- // Adjust nanoseconds past the epoch to match the column scale value
- t.getNanosSinceEpoch()
- .divide(BigDecimal.valueOf(Power10.intTable[9 - scale]))
- .toBigInteger())
- .orElseThrow(
- () ->
- new SFException(
- ErrorCode.INVALID_ROW, value, "Cannot be converted to a time value"));
+ // The input is bigger than max long value, treat it as nano-seconds directly
+ epochNanos = new BigInteger(input);
}
- }
-
- static String getStringValue(Object value) {
- String stringValue;
- if (value instanceof String) {
- stringValue = (String) value;
- } else if (value instanceof BigDecimal) {
- stringValue = ((BigDecimal) value).toPlainString();
- } else {
- stringValue = value.toString();
- }
- return stringValue;
+ return Instant.ofEpochSecond(
+ epochNanos.divide(Power10.sb16Table[9]).longValue(),
+ epochNanos.remainder(Power10.sb16Table[9]).longValue());
}
/**
- * Converts input to double value.
+ * Converts input to double value. Allowed Java types:
+ *
+ *
+ * - Number
+ *
- String
+ *
*
* @param input
+ * @param insertRowIndex
*/
- static double validateAndParseReal(Object input) {
- double doubleValue;
- try {
- if (input instanceof Double) {
- doubleValue = (double) input;
- } else if (input instanceof Float) {
- doubleValue = Double.parseDouble(input.toString());
- } else if (input instanceof BigDecimal) {
- doubleValue = ((BigDecimal) input).doubleValue();
- } else {
- doubleValue = Double.parseDouble((String) input);
+ static double validateAndParseReal(String columnName, Object input, long insertRowIndex) {
+ if (input instanceof Float) {
+ return Double.parseDouble(input.toString());
+ } else if (input instanceof Number) {
+ return ((Number) input).doubleValue();
+ } else if (input instanceof String) {
+ String stringInput = ((String) input).trim();
+ try {
+ return Double.parseDouble(stringInput);
+ } catch (NumberFormatException err) {
+ stringInput = stringInput.toLowerCase();
+ switch (stringInput) {
+ case "nan":
+ return Double.NaN;
+ case "inf":
+ return Double.POSITIVE_INFINITY;
+ case "-inf":
+ return Double.NEGATIVE_INFINITY;
+ default:
+ throw valueFormatNotAllowedException(
+ columnName, "REAL", "Not a valid decimal number", insertRowIndex);
+ }
}
- } catch (NumberFormatException err) {
- throw new SFException(ErrorCode.INVALID_ROW, input.toString());
}
-
- return doubleValue;
+ throw typeNotAllowedException(
+ columnName, input.getClass(), "REAL", new String[] {"Number", "String"}, insertRowIndex);
}
/**
* Validate and parse input to integer output, 1=true, 0=false. String values converted to boolean
* according to https://docs.snowflake.com/en/sql-reference/functions/to_boolean.html#usage-notes
+ * Allowed Java types:
*
- * @param input
+ *
+ * - boolean
+ *
- String
+ *
- java.lang.Number
+ *
+ *
+ * @param input Input to be converted
* @return 1 or 0 where 1=true, 0=false
*/
- static int validateAndParseBoolean(Object input) {
- int output = 0;
+ static int validateAndParseBoolean(String columnName, Object input, long insertRowIndex) {
if (input instanceof Boolean) {
- output = (boolean) input ? 1 : 0;
+ return (boolean) input ? 1 : 0;
} else if (input instanceof Number) {
- output = ((Number) input).doubleValue() > 0 ? 1 : 0;
- } else {
- output = convertStringToBoolean((String) input) ? 1 : 0;
+ return new BigDecimal(input.toString()).compareTo(BigDecimal.ZERO) == 0 ? 0 : 1;
+ } else if (input instanceof String) {
+ return convertStringToBoolean(columnName, (String) input, insertRowIndex) ? 1 : 0;
+ }
+
+ throw typeNotAllowedException(
+ columnName,
+ input.getClass(),
+ "BOOLEAN",
+ new String[] {"boolean", "Number", "String"},
+ insertRowIndex);
+ }
+
+ static void checkValueInRange(
+ BigDecimal bigDecimalValue, int scale, int precision, final long insertRowIndex) {
+ if (bigDecimalValue.abs().compareTo(BigDecimal.TEN.pow(precision - scale)) >= 0) {
+ throw new SFException(
+ ErrorCode.INVALID_FORMAT_ROW,
+ String.format(
+ "Number out of representable exclusive range of (-1e%s..1e%s), rowIndex:%d",
+ precision - scale, precision - scale, insertRowIndex));
}
- return output;
}
static Set allowedBooleanStringsLowerCased =
Sets.newHashSet("1", "0", "yes", "no", "y", "n", "t", "f", "true", "false", "on", "off");
- static boolean convertStringToBoolean(String value) {
- String lowerCasedValue = value.toLowerCase();
- if (!allowedBooleanStringsLowerCased.contains(lowerCasedValue)) {
- throw new SFException(ErrorCode.INVALID_ROW, value);
- }
- return "1".equals(lowerCasedValue)
- || "yes".equals(lowerCasedValue)
- || "y".equals(lowerCasedValue)
- || "t".equals(lowerCasedValue)
- || "true".equals(lowerCasedValue)
- || "on".equals(lowerCasedValue);
+ private static boolean convertStringToBoolean(
+ String columnName, String value, final long insertRowIndex) {
+ String normalizedInput = value.toLowerCase().trim();
+ if (!allowedBooleanStringsLowerCased.contains(normalizedInput)) {
+ throw valueFormatNotAllowedException(
+ columnName,
+ "BOOLEAN",
+ "Not a valid boolean, see"
+ + " https://docs.snowflake.com/en/sql-reference/data-types-logical.html#conversion-to-boolean"
+ + " for the list of supported formats",
+ insertRowIndex);
+ }
+ return "1".equals(normalizedInput)
+ || "yes".equals(normalizedInput)
+ || "y".equals(normalizedInput)
+ || "t".equals(normalizedInput)
+ || "true".equals(normalizedInput)
+ || "on".equals(normalizedInput);
+ }
+
+ /**
+ * Create exception that a Java type cannot be ingested into a specific Snowflake column type
+ *
+ * @param javaType Java type failing the validation
+ * @param snowflakeType Target Snowflake column type
+ * @param allowedJavaTypes Java types supported for the Java type
+ */
+ private static SFException typeNotAllowedException(
+ String columnName,
+ Class javaType,
+ String snowflakeType,
+ String[] allowedJavaTypes,
+ final long insertRowIndex) {
+ return new SFException(
+ ErrorCode.INVALID_FORMAT_ROW,
+ String.format(
+ "Object of type %s cannot be ingested into Snowflake column %s of type %s, rowIndex:%d",
+ javaType.getName(), columnName, snowflakeType, insertRowIndex),
+ String.format(
+ String.format("Allowed Java types: %s", String.join(", ", allowedJavaTypes))));
+ }
+
+ /**
+ * Create exception when the Java type is correct, but the value is invalid (e.g. boolean cannot
+ * be parsed from a string)
+ *
+ * Note: Do not log actual Object Value
+ *
+ * @param columnName Column Name
+ * @param snowflakeType Snowflake column type
+ * @param reason Reason why value format is not allowed.
+ * @param rowIndex Index of the Input row primarily for debugging purposes.
+ * @return SFException is thrown
+ */
+ private static SFException valueFormatNotAllowedException(
+ String columnName, String snowflakeType, String reason, final long rowIndex) {
+ return new SFException(
+ ErrorCode.INVALID_VALUE_ROW,
+ String.format(
+ "Value cannot be ingested into Snowflake column %s of type %s, rowIndex:%d, reason:"
+ + " %s",
+ columnName, snowflakeType, rowIndex, reason));
+ }
+
+ /**
+ * Validates that a string is valid UTF-8 string. It catches situations like unmatched high/low
+ * UTF-16 surrogate, for example.
+ */
+ private static void verifyValidUtf8(
+ String input, String columnName, String dataType, final long insertRowIndex) {
+ String roundTripStr =
+ new String(input.getBytes(StandardCharsets.UTF_8), StandardCharsets.UTF_8);
+ if (!input.equals(roundTripStr)) {
+ throw valueFormatNotAllowedException(
+ columnName, dataType, "Invalid Unicode string", insertRowIndex);
+ }
}
}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelResponse.java b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelResponse.java
new file mode 100644
index 000000000..fce91f1d2
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelResponse.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming.internal;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/** Response for a {@link net.snowflake.ingest.streaming.DropChannelRequest}. */
+class DropChannelResponse extends StreamingIngestResponse {
+ private Long statusCode;
+ private String message;
+ private String dbName;
+ private String schemaName;
+ private String tableName;
+ private String channelName;
+
+ @JsonProperty("status_code")
+ void setStatusCode(Long statusCode) {
+ this.statusCode = statusCode;
+ }
+
+ @Override
+ Long getStatusCode() {
+ return this.statusCode;
+ }
+
+ @JsonProperty("message")
+ void setMessage(String message) {
+ this.message = message;
+ }
+
+ String getMessage() {
+ return this.message;
+ }
+
+ @JsonProperty("database")
+ void setDBName(String dbName) {
+ this.dbName = dbName;
+ }
+
+ String getDBName() {
+ return this.dbName;
+ }
+
+ @JsonProperty("schema")
+ void setSchemaName(String schemaName) {
+ this.schemaName = schemaName;
+ }
+
+ String getSchemaName() {
+ return this.schemaName;
+ }
+
+ @JsonProperty("table")
+ void setTableName(String tableName) {
+ this.tableName = tableName;
+ }
+
+ String getTableName() {
+ return this.tableName;
+ }
+
+ @JsonProperty("channel")
+ void setChannelName(String channelName) {
+ this.channelName = channelName;
+ }
+
+ String getChannelName() {
+ return this.channelName;
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelVersionRequest.java b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelVersionRequest.java
new file mode 100644
index 000000000..abc6b8067
--- /dev/null
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/DropChannelVersionRequest.java
@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
+ */
+
+package net.snowflake.ingest.streaming.internal;
+
+import net.snowflake.ingest.streaming.DropChannelRequest;
+
+/**
+ * Same as DropChannelRequest but allows specifying a client sequencer to drop a specific version.
+ */
+class DropChannelVersionRequest extends DropChannelRequest {
+ private final Long clientSequencer;
+
+ public DropChannelVersionRequest(DropChannelRequestBuilder builder, long clientSequencer) {
+ super(builder);
+ this.clientSequencer = clientSequencer;
+ }
+
+ Long getClientSequencer() {
+ return this.clientSequencer;
+ }
+}
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/EpInfo.java b/src/main/java/net/snowflake/ingest/streaming/internal/EpInfo.java
index 32f77ceac..e6e6a4d9d 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/EpInfo.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/EpInfo.java
@@ -1,7 +1,11 @@
package net.snowflake.ingest.streaming.internal;
+import static net.snowflake.ingest.utils.Constants.EP_NDV_UNKNOWN;
+
import com.fasterxml.jackson.annotation.JsonProperty;
import java.util.Map;
+import net.snowflake.ingest.utils.ErrorCode;
+import net.snowflake.ingest.utils.SFException;
/** Class used to serialize/deserialize EP information */
class EpInfo {
@@ -17,6 +21,30 @@ class EpInfo {
this.columnEps = columnEps;
}
+ /** Some basic verification logic to make sure the EP info is correct */
+ public void verifyEpInfo() {
+ for (Map.Entry entry : columnEps.entrySet()) {
+ String colName = entry.getKey();
+ FileColumnProperties colEp = entry.getValue();
+ // Make sure the null count should always smaller than the total row count
+ if (colEp.getNullCount() > rowCount) {
+ throw new SFException(
+ ErrorCode.INTERNAL_ERROR,
+ String.format(
+ "Null count bigger than total row count on col=%s, nullCount=%s, rowCount=%s",
+ colName, colEp.getNullCount(), rowCount));
+ }
+
+ // Make sure the NDV should always be -1
+ if (colEp.getDistinctValues() != EP_NDV_UNKNOWN) {
+ throw new SFException(
+ ErrorCode.INTERNAL_ERROR,
+ String.format(
+ "Unexpected NDV on col=%s, value=%d", colName, colEp.getDistinctValues()));
+ }
+ }
+ }
+
@JsonProperty("rows")
long getRowCount() {
return rowCount;
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FileColumnProperties.java b/src/main/java/net/snowflake/ingest/streaming/internal/FileColumnProperties.java
index 9322a44ff..a305a52f7 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/FileColumnProperties.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/FileColumnProperties.java
@@ -1,11 +1,14 @@
package net.snowflake.ingest.streaming.internal;
+import static net.snowflake.ingest.streaming.internal.BinaryStringUtils.truncateBytesAsHex;
+
import com.fasterxml.jackson.annotation.JsonProperty;
import java.math.BigInteger;
import java.util.Objects;
/** Audit register endpoint/FileColumnPropertyDTO property list. */
class FileColumnProperties {
+ private int columnOrdinal;
private String minStrValue;
private String maxStrValue;
@@ -38,7 +41,11 @@ class FileColumnProperties {
// Default value to use for min/max int when all data in the given column is NULL
public static final BigInteger DEFAULT_MIN_MAX_INT_VAL_FOR_EP = BigInteger.valueOf(0);
+ // Default value to use for min/max real when all data in the given column is NULL
+ public static final Double DEFAULT_MIN_MAX_REAL_VAL_FOR_EP = 0d;
+
FileColumnProperties(RowBufferStats stats) {
+ this.setColumnOrdinal(stats.getOrdinal());
this.setCollation(stats.getCollationDefinitionString());
this.setMaxIntValue(
stats.getCurrentMaxIntValue() == null
@@ -48,17 +55,44 @@ class FileColumnProperties {
stats.getCurrentMinIntValue() == null
? DEFAULT_MIN_MAX_INT_VAL_FOR_EP
: stats.getCurrentMinIntValue());
- this.setMinRealValue(stats.getCurrentMinRealValue());
- this.setMaxRealValue(stats.getCurrentMaxRealValue());
+ this.setMinRealValue(
+ stats.getCurrentMinRealValue() == null
+ ? DEFAULT_MIN_MAX_REAL_VAL_FOR_EP
+ : stats.getCurrentMinRealValue());
+ this.setMaxRealValue(
+ stats.getCurrentMaxRealValue() == null
+ ? DEFAULT_MIN_MAX_REAL_VAL_FOR_EP
+ : stats.getCurrentMaxRealValue());
this.setMaxLength(stats.getCurrentMaxLength());
- this.setMaxStrNonCollated(stats.getCurrentMaxStrValue());
- this.setMinStrNonCollated(stats.getCurrentMinStrValue());
- this.setMaxStrValue(stats.getCurrentMaxColStrValue());
- this.setMinStrValue(stats.getCurrentMinColStrValue());
+
+ this.setMaxStrNonCollated(null);
+ this.setMinStrNonCollated(null);
+
+ // current hex-encoded min value, truncated down to 32 bytes
+ if (stats.getCurrentMinStrValue() != null) {
+ String truncatedAsHex = truncateBytesAsHex(stats.getCurrentMinStrValue(), false);
+ this.setMinStrValue(truncatedAsHex);
+ }
+
+ // current hex-encoded max value, truncated up to 32 bytes
+ if (stats.getCurrentMaxStrValue() != null) {
+ String truncatedAsHex = truncateBytesAsHex(stats.getCurrentMaxStrValue(), true);
+ this.setMaxStrValue(truncatedAsHex);
+ }
+
this.setNullCount(stats.getCurrentNullCount());
this.setDistinctValues(stats.getDistinctValues());
}
+ @JsonProperty("columnId")
+ public int getColumnOrdinal() {
+ return columnOrdinal;
+ }
+
+ public void setColumnOrdinal(int columnOrdinal) {
+ this.columnOrdinal = columnOrdinal;
+ }
+
// Annotation required in order to have package private fields serialized
@JsonProperty("minStrValue")
String getMinStrValue() {
@@ -171,6 +205,7 @@ void setMaxStrNonCollated(String maxStrNonCollated) {
@Override
public String toString() {
final StringBuilder sb = new StringBuilder("{");
+ sb.append("\"columnOrdinal\": ").append(columnOrdinal);
if (minIntValue != null) {
sb.append(", \"minIntValue\": ").append(minIntValue);
sb.append(", \"maxIntValue\": ").append(maxIntValue);
@@ -196,7 +231,8 @@ public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
FileColumnProperties that = (FileColumnProperties) o;
- return distinctValues == that.distinctValues
+ return Objects.equals(columnOrdinal, that.columnOrdinal)
+ && distinctValues == that.distinctValues
&& nullCount == that.nullCount
&& maxLength == that.maxLength
&& Objects.equals(minStrValue, that.minStrValue)
@@ -213,6 +249,7 @@ public boolean equals(Object o) {
@Override
public int hashCode() {
return Objects.hash(
+ columnOrdinal,
minStrValue,
maxStrValue,
collation,
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 60c72463d..30dadbe66 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java
@@ -5,27 +5,26 @@
package net.snowflake.ingest.streaming.internal;
import static net.snowflake.ingest.utils.Constants.BLOB_EXTENSION_TYPE;
-import static net.snowflake.ingest.utils.Constants.CPU_IO_TIME_RATIO;
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;
import static net.snowflake.ingest.utils.Constants.THREAD_SHUTDOWN_TIMEOUT_IN_SEC;
+import static net.snowflake.ingest.utils.Utils.getStackTrace;
import com.codahale.metrics.Timer;
-import java.io.ByteArrayOutputStream;
+import com.google.common.annotations.VisibleForTesting;
import java.io.IOException;
import java.lang.management.ManagementFactory;
-import java.nio.file.Path;
-import java.nio.file.Paths;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.Calendar;
-import java.util.HashMap;
+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;
@@ -33,27 +32,20 @@
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
-import java.util.zip.CRC32;
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.Cryptor;
import net.snowflake.ingest.utils.ErrorCode;
import net.snowflake.ingest.utils.Logging;
import net.snowflake.ingest.utils.Pair;
import net.snowflake.ingest.utils.SFException;
import net.snowflake.ingest.utils.Utils;
-import org.apache.arrow.util.VisibleForTesting;
-import org.apache.arrow.vector.VectorLoader;
-import org.apache.arrow.vector.VectorSchemaRoot;
-import org.apache.arrow.vector.VectorUnloader;
-import org.apache.arrow.vector.ipc.ArrowStreamWriter;
-import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
/**
* Responsible for flushing data from client to Snowflake tables. When a flush is triggered, it will
@@ -62,25 +54,30 @@
* build the blob using the ChannelData returned from previous step
* upload the blob to stage
* register the blob to the targeted Snowflake table
+ *
+ * @param type of column data ({@link ParquetChunkData})
*/
-class FlushService {
+class FlushService {
+
+ // The max number of upload retry attempts to the stage
+ private static final int DEFAULT_MAX_UPLOAD_RETRIES = 5;
// Static class to save the list of channels that are used to build a blob, which is mainly used
// to invalidate all the channels when there is a failure
- static class BlobData {
- private final String filePath;
- private final List> data;
+ static class BlobData {
+ private final String path;
+ private final List>> data;
- BlobData(String filePath, List> data) {
- this.filePath = filePath;
+ BlobData(String path, List>> data) {
+ this.path = path;
this.data = data;
}
- String getFilePath() {
- return filePath;
+ String getPath() {
+ return path;
}
- List> getData() {
+ List>> getData() {
return data;
}
}
@@ -91,7 +88,7 @@ List> getData() {
private final AtomicLong counter;
// The client that owns this flush service
- private final SnowflakeStreamingIngestClientInternal owningClient;
+ private final SnowflakeStreamingIngestClientInternal owningClient;
// Thread to schedule the flush job
@VisibleForTesting ScheduledExecutorService flushWorker;
@@ -103,13 +100,13 @@ List> getData() {
@VisibleForTesting ExecutorService buildUploadWorkers;
// Reference to the channel cache
- private final ChannelCache channelCache;
+ private final ChannelCache channelCache;
// Reference to the Streaming Ingest stage
private final StreamingIngestStage targetStage;
// Reference to register service
- private final RegisterService registerService;
+ private final RegisterService registerService;
// Indicates whether we need to schedule a flush
@VisibleForTesting volatile boolean isNeedFlush;
@@ -123,6 +120,9 @@ List> getData() {
// A map which stores the timer for each blob in order to capture the flush latency
private final Map latencyTimerContextMap;
+ // blob encoding version
+ private final Constants.BdecVersion bdecVersion;
+
/**
* Constructor for TESTING that takes (usually mocked) StreamingIngestStage
*
@@ -131,19 +131,20 @@ List> getData() {
* @param isTestMode
*/
FlushService(
- SnowflakeStreamingIngestClientInternal client,
- ChannelCache cache,
+ SnowflakeStreamingIngestClientInternal client,
+ ChannelCache cache,
StreamingIngestStage targetStage, // For testing
boolean isTestMode) {
this.owningClient = client;
this.channelCache = cache;
this.targetStage = targetStage;
this.counter = new AtomicLong(0);
- this.registerService = new RegisterService(client, isTestMode);
+ this.registerService = new RegisterService<>(client, isTestMode);
this.isNeedFlush = false;
this.lastFlushTime = System.currentTimeMillis();
this.isTestMode = isTestMode;
- this.latencyTimerContextMap = new HashMap<>();
+ this.latencyTimerContextMap = new ConcurrentHashMap<>();
+ this.bdecVersion = this.owningClient.getParameterProvider().getBlobFormatVersion();
createWorkers();
}
@@ -155,7 +156,7 @@ List> getData() {
* @param isTestMode
*/
FlushService(
- SnowflakeStreamingIngestClientInternal client, ChannelCache cache, boolean isTestMode) {
+ SnowflakeStreamingIngestClientInternal client, ChannelCache cache, boolean isTestMode) {
this.owningClient = client;
this.channelCache = cache;
try {
@@ -165,17 +166,20 @@ List> getData() {
client.getRole(),
client.getHttpClient(),
client.getRequestBuilder(),
- client.getName());
+ client.getName(),
+ DEFAULT_MAX_UPLOAD_RETRIES,
+ client.getProxyProperties());
} catch (SnowflakeSQLException | IOException err) {
throw new SFException(err, ErrorCode.UNABLE_TO_CONNECT_TO_STAGE);
}
- this.registerService = new RegisterService(client, isTestMode);
+ this.registerService = new RegisterService<>(client, isTestMode);
this.counter = new AtomicLong(0);
this.isNeedFlush = false;
this.lastFlushTime = System.currentTimeMillis();
this.isTestMode = isTestMode;
- this.latencyTimerContextMap = new HashMap<>();
+ this.latencyTimerContextMap = new ConcurrentHashMap<>();
+ this.bdecVersion = this.owningClient.getParameterProvider().getBlobFormatVersion();
createWorkers();
}
@@ -256,12 +260,13 @@ private CompletableFuture registerFuture() {
*/
CompletableFuture flush(boolean isForce) {
long timeDiffMillis = System.currentTimeMillis() - this.lastFlushTime;
+
if (isForce
|| (!DISABLE_BACKGROUND_FLUSH
- && !this.isTestMode
+ && !isTestMode()
&& (this.isNeedFlush
|| timeDiffMillis
- >= this.owningClient.getParameterProvider().getBufferFlushIntervalInMs()))) {
+ >= this.owningClient.getParameterProvider().getCachedMaxClientLagInMs()))) {
return this.statsFuture()
.thenCompose((v) -> this.distributeFlush(isForce, timeDiffMillis))
@@ -278,7 +283,25 @@ private void createWorkers() {
this.flushWorker = Executors.newSingleThreadScheduledExecutor(flushThreadFactory);
this.flushWorker.scheduleWithFixedDelay(
() -> {
- flush(false);
+ flush(false)
+ .exceptionally(
+ e -> {
+ String errorMessage =
+ String.format(
+ "Background flush task failed, client=%s, exception=%s, detail=%s,"
+ + " trace=%s.",
+ this.owningClient.getName(),
+ e.getCause(),
+ e.getCause().getMessage(),
+ getStackTrace(e.getCause()));
+ logger.logError(errorMessage);
+ if (this.owningClient.getTelemetryService() != null) {
+ this.owningClient
+ .getTelemetryService()
+ .reportClientFailure(this.getClass().getSimpleName(), errorMessage);
+ }
+ return null;
+ });
},
0,
this.owningClient.getParameterProvider().getBufferFlushCheckIntervalInMs(),
@@ -290,13 +313,14 @@ private void createWorkers() {
this.registerWorker = Executors.newSingleThreadExecutor(registerThreadFactory);
// Create threads for building and uploading blobs
- // Size: number of available processors * (1 + IO time/CPU time), currently the
- // CPU_IO_TIME_RATIO is 1 under the assumption that build and upload will take similar time
+ // Size: number of available processors * (1 + IO time/CPU time)
ThreadFactory buildUploadThreadFactory =
new ThreadFactoryBuilder().setNameFormat("ingest-build-upload-thread-%d").build();
int buildUploadThreadCount =
Math.min(
- Runtime.getRuntime().availableProcessors() * (1 + CPU_IO_TIME_RATIO), MAX_THREAD_COUNT);
+ Runtime.getRuntime().availableProcessors()
+ * (1 + this.owningClient.getParameterProvider().getIOTimeCpuRatio()),
+ MAX_THREAD_COUNT);
this.buildUploadWorkers =
Executors.newFixedThreadPool(buildUploadThreadCount, buildUploadThreadFactory);
@@ -312,69 +336,156 @@ private void createWorkers() {
* off a build blob work when certain size has reached or we have reached the end
*/
void distributeFlushTasks() {
- Iterator>>
+ Iterator<
+ Map.Entry<
+ String, ConcurrentHashMap>>>
itr = this.channelCache.iterator();
- List>> blobs = new ArrayList<>();
-
- while (itr.hasNext()) {
- List> blobData = new ArrayList<>();
- float totalBufferSize = 0;
-
- // Distribute work at table level, create a new blob if reaching the blob size limit
- while (itr.hasNext() && totalBufferSize <= MAX_BLOB_SIZE_IN_BYTES) {
- ConcurrentHashMap table =
- itr.next().getValue();
- List channelsDataPerTable = new ArrayList<>();
- // TODO: we could do parallel stream to get the channelData if needed
- for (SnowflakeStreamingIngestChannelInternal channel : table.values()) {
- if (channel.isValid()) {
- ChannelData data = channel.getData();
- if (data != null) {
- channelsDataPerTable.add(data);
- totalBufferSize += data.getBufferSize();
- }
- }
+ List, CompletableFuture>> blobs = new ArrayList<>();
+ List> leftoverChannelsDataPerTable = new ArrayList<>();
+
+ while (itr.hasNext() || !leftoverChannelsDataPerTable.isEmpty()) {
+ List>> blobData = new ArrayList<>();
+ float totalBufferSizeInBytes = 0F;
+ final String blobPath = getBlobPath(this.targetStage.getClientPrefix());
+
+ // Distribute work at table level, split the blob if reaching the blob size limit or the
+ // channel has different encryption key ids
+ while (itr.hasNext() || !leftoverChannelsDataPerTable.isEmpty()) {
+ List> channelsDataPerTable = Collections.synchronizedList(new ArrayList<>());
+ if (!leftoverChannelsDataPerTable.isEmpty()) {
+ channelsDataPerTable.addAll(leftoverChannelsDataPerTable);
+ leftoverChannelsDataPerTable.clear();
+ } else if (blobData.size()
+ >= this.owningClient
+ .getParameterProvider()
+ .getMaxChunksInBlobAndRegistrationRequest()) {
+ // Create a new blob if the current one already contains max allowed number of chunks
+ logger.logInfo(
+ "Max allowed number of chunks in the current blob reached. chunkCount={}"
+ + " maxChunkCount={} currentBlobPath={}",
+ blobData.size(),
+ this.owningClient.getParameterProvider().getMaxChunksInBlobAndRegistrationRequest(),
+ blobPath);
+ break;
+ } else {
+ ConcurrentHashMap> table =
+ itr.next().getValue();
+ // Use parallel stream since getData could be the performance bottleneck when we have a
+ // high number of channels
+ table.values().parallelStream()
+ .forEach(
+ channel -> {
+ if (channel.isValid()) {
+ ChannelData data = channel.getData(blobPath);
+ if (data != null) {
+ channelsDataPerTable.add(data);
+ }
+ }
+ });
}
+
if (!channelsDataPerTable.isEmpty()) {
- blobData.add(channelsDataPerTable);
+ int idx = 0;
+ float totalBufferSizePerTableInBytes = 0F;
+ while (idx < channelsDataPerTable.size()) {
+ ChannelData