SNOW-984877 Update to support customized url and add customer name in…

… request header (#656)

* SNOW-984877 Update to support customzed url parser and RequestBuilder

* update to set account name header by argument

* update format and compiling issue

* address comment

src/main/java/net/snowflake/ingest/connection/RequestBuilder.java

@@ -54,6 +54,10 @@ public class RequestBuilder {
   // whatever the actual host is
   private final String host;
 
+  private final String accountName;
+
+  private final boolean addAccountNameInRequest;
+
   private final String userAgentSuffix;
 
   // Reference to the telemetry service
@@ -123,6 +127,8 @@ public class RequestBuilder {
 
   public static final String HTTP_HEADER_CONTENT_TYPE_JSON = "application/json";
 
+  private static final String SF_HEADER_ACCOUNT_NAME = "Snowflake-Account";
+
   /**
    * RequestBuilder - general usage constructor
    *
@@ -206,6 +212,36 @@ public RequestBuilder(
         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,
+      boolean addAccountNameInRequest) {
+    this(
+        url.getAccount(),
+        userName,
+        credential,
+        url.getScheme(),
+        url.getUrlWithoutPort(),
+        url.getPort(),
+        null,
+        null,
+        httpClient,
+        clientName,
+        addAccountNameInRequest);
+  }
+
   /**
    * RequestBuilder - constructor used by streaming ingest
    *
@@ -259,6 +295,46 @@ public RequestBuilder(
       SecurityManager securityManager,
       CloseableHttpClient httpClient,
       String clientName) {
+    this(accountName,
+        userName,
+        credential,
+        schemeName,
+        hostName,
+        portNum,
+        userAgentSuffix,
+        securityManager,
+        httpClient,
+        clientName,
+        false);
+  }
+
+  /**
+   * 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
+   * @param addAccountNameInRequest if ture, add account name in request header
+   */
+  public RequestBuilder(
+      String accountName,
+      String userName,
+      Object credential,
+      String schemeName,
+      String hostName,
+      int portNum,
+      String userAgentSuffix,
+      SecurityManager securityManager,
+      CloseableHttpClient httpClient,
+      String clientName,
+      boolean addAccountNameInRequest) {
     // none of these arguments should be null
     if (accountName == null || userName == null || credential == null) {
       throw new IllegalArgumentException();
@@ -281,6 +357,8 @@ public RequestBuilder(
     this.scheme = schemeName;
     this.host = hostName;
     this.userAgentSuffix = userAgentSuffix;
+    this.accountName = accountName;
+    this.addAccountNameInRequest = addAccountNameInRequest;
 
     // create our security/token manager
     if (securityManager == null) {
@@ -570,6 +648,9 @@ private static void addUserAgent(HttpUriRequest request, String userAgentSuffix)
   public void addToken(HttpUriRequest request) {
     request.setHeader(HttpHeaders.AUTHORIZATION, BEARER_PARAMETER + securityManager.getToken());
     request.setHeader(SF_HEADER_AUTHORIZATION_TOKEN_TYPE, this.securityManager.getTokenType());
+    if(addAccountNameInRequest) {
+      request.setHeader(SF_HEADER_ACCOUNT_NAME, accountName);
+    }
   }
 
   /**

src/main/java/net/snowflake/ingest/streaming/SnowflakeStreamingIngestClientFactory.java

@@ -6,6 +6,7 @@
 
 import java.util.Map;
 import java.util.Properties;
+
 import net.snowflake.ingest.streaming.internal.SnowflakeStreamingIngestClientInternal;
 import net.snowflake.ingest.utils.Constants;
 import net.snowflake.ingest.utils.SnowflakeURL;
@@ -28,6 +29,12 @@ public static class Builder {
     // Allows client to override some default parameter values
     private Map<String, Object> parameterOverrides;
 
+    // preset SnowflakeURL instance
+    private SnowflakeURL snowflakeURL;
+
+    // flag to specify if we need to add account name in the request header
+    private boolean addAccountNameInRequest;
+
     private Builder(String name) {
       this.name = name;
     }
@@ -37,6 +44,16 @@ public Builder setProperties(Properties prop) {
       return this;
     }
 
+    public Builder setSnowflakeURL(SnowflakeURL snowflakeURL) {
+      this.snowflakeURL = snowflakeURL;
+      return this;
+    }
+
+    public Builder setAddAccountNameInRequest(boolean addAccountNameInRequest) {
+      this.addAccountNameInRequest = addAccountNameInRequest;
+      return this;
+    }
+
     public Builder setParameterOverrides(Map<String, Object> parameterOverrides) {
       this.parameterOverrides = parameterOverrides;
       return this;
@@ -47,10 +64,17 @@ 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 (addAccountNameInRequest) {
+        return new SnowflakeStreamingIngestClientInternal<>(
+          this.name, accountURL, prop, this.parameterOverrides, addAccountNameInRequest);
+      }
       return new SnowflakeStreamingIngestClientInternal<>(
-          this.name, accountURL, prop, this.parameterOverrides);
+        this.name, accountURL, prop, this.parameterOverrides);
     }
   }
 }

src/main/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java

@@ -158,6 +158,30 @@ public class SnowflakeStreamingIngestClientInternal<T> implements SnowflakeStrea
       boolean isTestMode,
       RequestBuilder requestBuilder,
       Map<String, Object> parameterOverrides) {
+    this(name, accountURL, prop, httpClient, isTestMode, requestBuilder, parameterOverrides, false);
+  }
+
+  /**
+   * 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 addAccountNameInRequest if true, will add account name in request header
+   */
+  SnowflakeStreamingIngestClientInternal(
+      String name,
+      SnowflakeURL accountURL,
+      Properties prop,
+      CloseableHttpClient httpClient,
+      boolean isTestMode,
+      RequestBuilder requestBuilder,
+      Map<String, Object> parameterOverrides,
+      boolean addAccountNameInRequest) {
     this.parameterProvider = new ParameterProvider(parameterOverrides, prop);
 
     this.name = name;
@@ -190,13 +214,13 @@ public class SnowflakeStreamingIngestClientInternal<T> implements SnowflakeStrea
                 prop.getProperty(Constants.OAUTH_CLIENT_SECRET),
                 prop.getProperty(Constants.OAUTH_REFRESH_TOKEN));
       }
-      this.requestBuilder =
-          new RequestBuilder(
+      this.requestBuilder = new RequestBuilder(
               accountURL,
               prop.get(USER).toString(),
               credential,
               this.httpClient,
-              String.format("%s_%s", this.name, System.currentTimeMillis()));
+              String.format("%s_%s", this.name, System.currentTimeMillis()),
+              addAccountNameInRequest);
 
       logger.logInfo("Using {} for authorization", this.requestBuilder.getAuthType());
 
@@ -236,6 +260,31 @@ 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 addAccountNameInRequest if true, add account name in request header
+   */
+  public SnowflakeStreamingIngestClientInternal(
+      String name,
+      SnowflakeURL accountURL,
+      Properties prop,
+      Map<String, Object> parameterOverrides,
+      boolean addAccountNameInRequest) {
+    this(name,
+        accountURL,
+        prop,
+        null,
+        false,
+        null,
+        parameterOverrides,
+        addAccountNameInRequest);
+  }
+
   /**
    * Constructor for TEST ONLY
    *