From e52bc1f0cb502c0066eed0050b2eb7c268d7ea8e Mon Sep 17 00:00:00 2001 From: David Wang Date: Sat, 9 Dec 2023 01:21:00 +0000 Subject: [PATCH] SNOW-984877 Update to support customzed url parser and RequestBuilder --- .../connection/RequestBuilderFactory.java | 14 +++++ ...SnowflakeStreamingIngestClientFactory.java | 25 +++++++- ...nowflakeStreamingIngestClientInternal.java | 63 +++++++++++++++---- 3 files changed, 88 insertions(+), 14 deletions(-) create mode 100644 src/main/java/net/snowflake/ingest/connection/RequestBuilderFactory.java diff --git a/src/main/java/net/snowflake/ingest/connection/RequestBuilderFactory.java b/src/main/java/net/snowflake/ingest/connection/RequestBuilderFactory.java new file mode 100644 index 000000000..3d81ee6be --- /dev/null +++ b/src/main/java/net/snowflake/ingest/connection/RequestBuilderFactory.java @@ -0,0 +1,14 @@ +package net.snowflake.ingest.connection; + +import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient; +import net.snowflake.ingest.utils.SnowflakeURL; + +public class RequestBuilderFactory { + public RequestBuilder build(SnowflakeURL url, + String userName, + Object credential, + CloseableHttpClient httpClient, + String clientName) { + return new RequestBuilder(url, userName, credential, httpClient, clientName); + } +} diff --git a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java index 2af794cd2..6548d13e4 100644 --- a/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java +++ b/src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java @@ -6,6 +6,8 @@ import java.util.Map; import java.util.Properties; + +import net.snowflake.ingest.connection.RequestBuilderFactory; import net.snowflake.ingest.streaming.internal.SnowflakeStreamingIngestClientInternal; import net.snowflake.ingest.utils.Constants; import net.snowflake.ingest.utils.SnowflakeURL; @@ -28,6 +30,10 @@ public static class Builder { // Allows client to override some default parameter values private Map parameterOverrides; + private SnowflakeURL snowflakeURL; + + private RequestBuilderFactory requestBuilderFactory; + private Builder(String name) { this.name = name; } @@ -37,6 +43,16 @@ public Builder setProperties(Properties prop) { return this; } + public Builder setSnowflakeURL(SnowflakeURL snowflakeURL) { + this.snowflakeURL = snowflakeURL; + return this; + } + + public Builder setRequestBuilderFactory(RequestBuilderFactory requestBuilderFactory) { + this.requestBuilderFactory = requestBuilderFactory; + return this; + } + public Builder setParameterOverrides(Map parameterOverrides) { this.parameterOverrides = parameterOverrides; return this; @@ -47,8 +63,15 @@ public SnowflakeStreamingIngestClient build() { Utils.assertNotNull("connection properties", this.prop); Properties prop = Utils.createProperties(this.prop); - SnowflakeURL accountURL = new SnowflakeURL(prop.getProperty(Constants.ACCOUNT_URL)); + SnowflakeURL accountURL = this.snowflakeURL; + if (accountURL == null) { + accountURL = new SnowflakeURL(prop.getProperty(Constants.ACCOUNT_URL)); + } + if (requestBuilderFactory != null) { + return new SnowflakeStreamingIngestClientInternal<>( + this.name, accountURL, prop, this.parameterOverrides, requestBuilderFactory); + } return new SnowflakeStreamingIngestClientInternal<>( this.name, accountURL, prop, this.parameterOverrides); } diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java index a52149331..b81b853f4 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java @@ -59,10 +59,7 @@ import javax.management.ObjectName; import net.snowflake.client.core.SFSessionProperty; import net.snowflake.client.jdbc.internal.apache.http.impl.client.CloseableHttpClient; -import net.snowflake.ingest.connection.IngestResponseException; -import net.snowflake.ingest.connection.OAuthCredential; -import net.snowflake.ingest.connection.RequestBuilder; -import net.snowflake.ingest.connection.TelemetryService; +import net.snowflake.ingest.connection.*; import net.snowflake.ingest.streaming.OpenChannelRequest; import net.snowflake.ingest.streaming.SnowflakeStreamingIngestChannel; import net.snowflake.ingest.streaming.SnowflakeStreamingIngestClient; @@ -151,13 +148,36 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea * @param parameterOverrides parameters we override in case we want to set different values */ SnowflakeStreamingIngestClientInternal( - String name, - SnowflakeURL accountURL, - Properties prop, - CloseableHttpClient httpClient, - boolean isTestMode, - RequestBuilder requestBuilder, - Map parameterOverrides) { + String name, + SnowflakeURL accountURL, + Properties prop, + CloseableHttpClient httpClient, + boolean isTestMode, + RequestBuilder requestBuilder, + Map parameterOverrides) { + this(name, accountURL, prop, httpClient, isTestMode, requestBuilder, parameterOverrides, new RequestBuilderFactory()); + } + /** + * Constructor + * + * @param name the name of the client + * @param accountURL Snowflake account url + * @param prop connection properties + * @param httpClient http client for sending request + * @param isTestMode whether we're under test mode + * @param requestBuilder http request builder + * @param parameterOverrides parameters we override in case we want to set different values + * @param requestBuilderFactory instance of RequestBuilderFactory + */ + SnowflakeStreamingIngestClientInternal( + String name, + SnowflakeURL accountURL, + Properties prop, + CloseableHttpClient httpClient, + boolean isTestMode, + RequestBuilder requestBuilder, + Map parameterOverrides, + RequestBuilderFactory requestBuilderFactory) { this.parameterProvider = new ParameterProvider(parameterOverrides, prop); this.name = name; @@ -190,8 +210,7 @@ public class SnowflakeStreamingIngestClientInternal implements SnowflakeStrea prop.getProperty(Constants.OAUTH_CLIENT_SECRET), prop.getProperty(Constants.OAUTH_REFRESH_TOKEN)); } - this.requestBuilder = - new RequestBuilder( + this.requestBuilder = requestBuilderFactory.build( accountURL, prop.get(USER).toString(), credential, @@ -236,6 +255,24 @@ public SnowflakeStreamingIngestClientInternal( this(name, accountURL, prop, null, false, null, parameterOverrides); } + /** + * Default Constructor + * + * @param name the name of the client + * @param accountURL Snowflake account url + * @param prop connection properties + * @param parameterOverrides map of parameters to override for this client + * @param requestBuilderFactory instance of RequestBuilderFactory + */ + public SnowflakeStreamingIngestClientInternal( + String name, + SnowflakeURL accountURL, + Properties prop, + Map parameterOverrides, + RequestBuilderFactory requestBuilderFactory) { + this(name, accountURL, prop, null, false, null, parameterOverrides, requestBuilderFactory); + } + /** * Constructor for TEST ONLY *