[SNOW-559751] Add role switching to SimpleIngestManager

src/main/java/net/snowflake/ingest/SimpleIngestManager.java

@@ -73,7 +73,8 @@ public static class Builder {
     // Hostname to connect to, default will be RequestBuilder#DEFAULT_HOST
     private String hostName;
 
-    private PrivateKey privateKey;
+    // the role name we want to use to authenticate
+    private String role;
 
     /**
      * getAccount - returns the name of the account this builder will inject into the IngestManager
@@ -164,6 +165,27 @@ public String getUserAgentSuffix() {
       return userAgentSuffix;
     }
 
+    /**
+     * Sets the name of the role that will be used for authorization. If not set, the default role
+     * assigned to the user will be used.
+     *
+     * @param role the role name we'll be using for auth
+     * @return the current builder with the role name set
+     */
+    public Builder setRole(String role) {
+      this.role = role;
+      return this;
+    }
+
+    /**
+     * Get the set role name.
+     *
+     * <p>It can be null or empty.
+     */
+    public String getRole() {
+      return role;
+    }
+
     /* Sets the user agent suffix as part of this SimpleIngestManager Instance */
     public Builder setUserAgentSuffix(String userAgentSuffix) {
       this.userAgentSuffix = userAgentSuffix;
@@ -188,9 +210,9 @@ public Builder setHostName(String hostName) {
     public SimpleIngestManager build() {
       if (isNullOrEmpty(hostName)) {
         return new SimpleIngestManager(
-            account, user, pipe, DEFAULT_HOST_SUFFIX, keypair, userAgentSuffix);
+            account, user, pipe, DEFAULT_HOST_SUFFIX, keypair, userAgentSuffix, role);
       }
-      return new SimpleIngestManager(account, user, pipe, hostName, keypair, userAgentSuffix);
+      return new SimpleIngestManager(account, user, pipe, hostName, keypair, userAgentSuffix, role);
     }
   }
 
@@ -334,6 +356,36 @@ public SimpleIngestManager(
     this.builder = new RequestBuilder(account, user, hostName, keyPair, userAgentSuffix);
   }
 
+  /**
+   * Using this constructor for Builder pattern. KeyPair can be passed in now since we have
+   * made @see {@link Utils#createKeyPairFromPrivateKey(PrivateKey)} public
+   *
+   * @param account The account into which we're loading Note: account should not include region or
+   *     cloud provider info. e.g. if host is testaccount.us-east-1.azure .snowflakecomputing.com,
+   *     account should be testaccount. If this is the case, you should use the constructor that
+   *     accepts hostname as argument
+   * @param user the user performing this load
+   * @param pipe the fully qualified name of the pipe
+   * @param hostName the hostname
+   * @param keyPair keyPair associated with the private key used for authentication. See @see {@link
+   *     Utils#createKeyPairFromPrivateKey} to generate KP from p8Key
+   * @param userAgentSuffix user agent suffix we want to add.
+   */
+  public SimpleIngestManager(
+      String account,
+      String user,
+      String pipe,
+      String hostName,
+      KeyPair keyPair,
+      String userAgentSuffix,
+      String role) {
+    // call our initializer method
+    init(account, user, pipe, keyPair);
+
+    // create the request builder
+    this.builder = new RequestBuilder(account, user, hostName, keyPair, userAgentSuffix, role);
+  }
+
   /**
    * Constructs a SimpleIngestManager for a given user in a specific account In addition, this also
    * takes takes the target table and source stage Finally, it also requires a valid private key
@@ -388,6 +440,28 @@ public SimpleIngestManager(
         new RequestBuilder(account, user, keyPair, schemeName, hostName, port, userAgentSuffix);
   }
 
+  /* Another flavor of constructor which supports userAgentSuffix and role */
+  public SimpleIngestManager(
+      String account,
+      String user,
+      String pipe,
+      PrivateKey privateKey,
+      String schemeName,
+      String hostName,
+      int port,
+      String userAgentSuffix,
+      String role)
+      throws NoSuchAlgorithmException, InvalidKeySpecException {
+    KeyPair keyPair = Utils.createKeyPairFromPrivateKey(privateKey);
+    // call our initializer method
+    init(account, user, pipe, keyPair);
+
+    // make the request builder we'll use to build messages to the service
+    builder =
+        new RequestBuilder(
+            account, user, keyPair, schemeName, hostName, port, userAgentSuffix, role);
+  }
+
   // ========= Constructors End =========
 
   /**

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

@@ -65,6 +65,9 @@ public class RequestBuilder {
 
   private final String userAgentSuffix;
 
+  // the actual role name
+  private final String role;
+
   /* Member variables End */
 
   /* Static constants Begin */
@@ -136,6 +139,8 @@ public class RequestBuilder {
   public static final String SF_HEADER_AUTHORIZATION_TOKEN_TYPE =
       "X-Snowflake-Authorization-Token-Type";
 
+  public static final String SF_HEADER_ROLE = "X-Snowflake-Role";
+
   public static final String JWT_TOKEN_TYPE = "KEYPAIR_JWT";
 
   public static final String HTTP_HEADER_CONTENT_TYPE_JSON = "application/json";
@@ -167,6 +172,33 @@ public RequestBuilder(
     this(accountName, userName, keyPair, DEFAULT_SCHEME, hostName, DEFAULT_PORT, userAgentSuffix);
   }
 
+  /**
+   * RequestBuilder constructor which uses default schemes, host and port.
+   *
+   * @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 userAgentSuffix - The suffix part of HTTP Header User-Agent
+   * @param role - the role for which the actions are performed.
+   */
+  public RequestBuilder(
+      String accountName,
+      String userName,
+      String hostName,
+      KeyPair keyPair,
+      String userAgentSuffix,
+      String role) {
+    this(
+        accountName,
+        userName,
+        keyPair,
+        DEFAULT_SCHEME,
+        hostName,
+        DEFAULT_PORT,
+        userAgentSuffix,
+        role);
+  }
+
   /**
    * Constructor to use if not intended to use userAgentSuffix. i.e. User-Agent HTTP header suffix
    * part is null, (The default one is still valid, check out #defaultUserAgent)
@@ -207,6 +239,30 @@ public RequestBuilder(
       String hostName,
       int portNum,
       String userAgentSuffix) {
+    this(accountName, userName, keyPair, schemeName, hostName, portNum, userAgentSuffix, null);
+  }
+
+  /**
+   * 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 keyPair - our auth credentials
+   * @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 role - the role for which the actions are performed.
+   */
+  public RequestBuilder(
+      String accountName,
+      String userName,
+      KeyPair keyPair,
+      String schemeName,
+      String hostName,
+      int portNum,
+      String userAgentSuffix,
+      String role) {
     // none of these arguments should be null
     if (accountName == null || userName == null || keyPair == null) {
       throw new IllegalArgumentException();
@@ -224,16 +280,18 @@ public RequestBuilder(
     this.scheme = schemeName;
     this.host = hostName;
     this.userAgentSuffix = userAgentSuffix;
+    this.role = role;
 
     LOGGER.info(
-        "Creating a RequestBuilder with arguments : "
-            + "Account : {}, User : {}, Scheme : {}, Host : {}, Port : {}, userAgentSuffix: {}",
+        "Creating a RequestBuilder with arguments : Account : {}, User : {}, Scheme : {}, Host :"
+            + " {}, Port : {}, userAgentSuffix: {}, Role: {}",
         account,
         user,
         this.scheme,
         this.host,
         this.port,
-        this.userAgentSuffix);
+        this.userAgentSuffix,
+        this.role);
   }
 
   /**
@@ -615,12 +673,28 @@ private static void addToken(HttpUriRequest request, String token) {
     request.setHeader(SF_HEADER_AUTHORIZATION_TOKEN_TYPE, JWT_TOKEN_TYPE);
   }
 
-  private static void addHeaders(HttpUriRequest request, String token, String userAgentSuffix) {
+  /**
+   * addToken - adds a X-Snowflake-Role to a request
+   *
+   * @param request the URI request
+   * @param role the role name to add
+   */
+  private static void addRoleHeader(HttpUriRequest request, String role) {
+    if (!isNullOrEmpty(role)) {
+      request.setHeader(SF_HEADER_ROLE, role);
+    }
+  }
+
+  private static void addHeaders(
+      HttpUriRequest request, String token, String userAgentSuffix, String role) {
     addUserAgent(request, userAgentSuffix);
 
     // Add the auth token
     addToken(request, token);
 
+    // Add the role header
+    addRoleHeader(request, role);
+
     // Add Accept header
     request.setHeader(HttpHeaders.ACCEPT, HTTP_HEADER_CONTENT_TYPE_JSON);
   }
@@ -668,7 +742,7 @@ public HttpPost generateInsertRequest(
     // Make the post request
     HttpPost post = new HttpPost(insertURI);
 
-    addHeaders(post, securityManager.getToken(), this.userAgentSuffix);
+    addHeaders(post, securityManager.getToken(), this.userAgentSuffix, this.role);
 
     // the entity for the containing the json
     final StringEntity entity =
@@ -697,7 +771,7 @@ public HttpGet generateHistoryRequest(
     // make the get request
     HttpGet get = new HttpGet(historyURI);
 
-    addHeaders(get, securityManager.getToken(), this.userAgentSuffix);
+    addHeaders(get, securityManager.getToken(), this.userAgentSuffix, this.role);
 
     return get;
   }
@@ -723,7 +797,11 @@ public HttpGet generateHistoryRangeRequest(
 
     HttpGet get = new HttpGet(historyRangeURI);
 
-    addHeaders(get, securityManager.getToken(), this.userAgentSuffix /*User agent information*/);
+    addHeaders(
+        get,
+        securityManager.getToken(),
+        this.userAgentSuffix /*User agent information*/,
+        this.role);
 
     return get;
   }
@@ -751,7 +829,11 @@ public HttpPost generateStreamingIngestPostRequest(
     // Make the post request
     HttpPost post = new HttpPost(uri);
 
-    addHeaders(post, securityManager.getToken(), this.userAgentSuffix /*User agent information*/);
+    addHeaders(
+        post,
+        securityManager.getToken(),
+        this.userAgentSuffix /*User agent information*/,
+        this.role);
 
     // The entity for the containing the json
     final StringEntity entity = new StringEntity(payload, ContentType.APPLICATION_JSON);
@@ -772,7 +854,7 @@ 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);
+    addHeaders(post, securityManager.getToken(), this.userAgentSuffix, this.role);
     return post;
   }
 
@@ -809,7 +891,7 @@ 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);
+    addHeaders(get, securityManager.getToken(), this.userAgentSuffix, this.role);
     return get;
   }