From 66e161437558dfc8098edcb591efabfc464ace2a Mon Sep 17 00:00:00 2001 From: Kamil Bregula Date: Fri, 1 Apr 2022 16:05:10 +0200 Subject: [PATCH] [SNOW-559751] Add role swithcing to SimpleIngestManager --- .../snowflake/ingest/SimpleIngestManager.java | 80 +++++++++++++- .../ingest/connection/RequestBuilder.java | 102 ++++++++++++++++-- .../net/snowflake/ingest/SimpleIngestIT.java | 49 +++++++-- .../java/net/snowflake/ingest/TestUtils.java | 15 ++- 4 files changed, 219 insertions(+), 27 deletions(-) diff --git a/src/main/java/net/snowflake/ingest/SimpleIngestManager.java b/src/main/java/net/snowflake/ingest/SimpleIngestManager.java index 1064c25ab..fa004a153 100644 --- a/src/main/java/net/snowflake/ingest/SimpleIngestManager.java +++ b/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. + * + *

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 ========= /** diff --git a/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java b/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java index 44d846b94..998d860d5 100644 --- a/src/main/java/net/snowflake/ingest/connection/RequestBuilder.java +++ b/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; } diff --git a/src/test/java/net/snowflake/ingest/SimpleIngestIT.java b/src/test/java/net/snowflake/ingest/SimpleIngestIT.java index 02c32aca2..aa7c69069 100644 --- a/src/test/java/net/snowflake/ingest/SimpleIngestIT.java +++ b/src/test/java/net/snowflake/ingest/SimpleIngestIT.java @@ -249,7 +249,7 @@ public void testUserAgentSuffixForInsertFileAPI() throws Exception { final String userAgentSuffix = "kafka-provider/NONE"; // create ingest manager - SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix); + SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix, null); // create a file wrapper StagedFileWrapper myFile = new StagedFileWrapper(TEST_FILE_NAME, null); @@ -274,7 +274,8 @@ public void testUserAgentSuffixForInsertFileAPI() throws Exception { // Passing null and empty string would also work // create ingest manager - SimpleIngestManager nullUserAgentSuffixIngestManager = TestUtils.getManager(pipeName, null); + SimpleIngestManager nullUserAgentSuffixIngestManager = + TestUtils.getManager(pipeName, null, null); HttpPost nullAdditionalUserAgent = nullUserAgentSuffixIngestManager @@ -283,7 +284,8 @@ public void testUserAgentSuffixForInsertFileAPI() throws Exception { UUID.randomUUID(), pipeName, Collections.singletonList(myFile), false); verifyDefaultUserAgent(nullAdditionalUserAgent.getAllHeaders(), false, null); - SimpleIngestManager emptyUserAgentSuffixIngestManager = TestUtils.getManager(pipeName, null); + SimpleIngestManager emptyUserAgentSuffixIngestManager = + TestUtils.getManager(pipeName, null, null); HttpPost emptyAdditionalUserAgent = emptyUserAgentSuffixIngestManager .getRequestBuilder() @@ -331,9 +333,18 @@ private void verifyDefaultUserAgent( @Ignore @Test - public void testConfigureClientHappyCase() throws Exception { + public void testConfigureClientHappyCaseNullRole() throws Exception { final String userAgentSuffix = "kafka-provider/NONE"; - SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix); + SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix, null); + ConfigureClientResponse configureClientResponse = manager.configureClient(null); + assertEquals(0L, configureClientResponse.getClientSequencer().longValue()); + } + + @Ignore + @Test + public void testConfigureClientHappyCaseExplicitRole() throws Exception { + final String userAgentSuffix = "kafka-provider/NONE"; + SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix, TestUtils.getRole()); ConfigureClientResponse configureClientResponse = manager.configureClient(null); assertEquals(0L, configureClientResponse.getClientSequencer().longValue()); } @@ -342,7 +353,7 @@ public void testConfigureClientHappyCase() throws Exception { @Test public void testConfigureClientNoPipeFound() throws Exception { final String userAgentSuffix = "kafka-provider/NONE"; - SimpleIngestManager manager = TestUtils.getManager("nopipe", userAgentSuffix); + SimpleIngestManager manager = TestUtils.getManager("nopipe", userAgentSuffix, null); try { manager.configureClient(null); } catch (IngestResponseException exception) { @@ -353,11 +364,27 @@ public void testConfigureClientNoPipeFound() throws Exception { } } + @Ignore + @Test + public void testConfigureClientNoRoleFound() throws Exception { + final String userAgentSuffix = "kafka-provider/NONE"; + SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix, "INVALID-ROLE"); + try { + manager.configureClient(null); + } catch (IngestResponseException exception) { + assertEquals(404, exception.getErrorCode()); + assertEquals( + "Requested role 'INVALID-ROLE' does not exist or is not assigned to you. Login using" + + " connect string specifying another role (e.g. PUBLIC).", + exception.getErrorBody().getMessage()); + } + } + @Ignore @Test public void testGetClientStatusHappyCase() throws Exception { final String userAgentSuffix = "kafka-provider/NONE"; - SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix); + SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix, null); manager.configureClient(null); ClientStatusResponse clientStatusResponse = manager.getClientStatus(null); assertEquals(0L, clientStatusResponse.getClientSequencer().longValue()); @@ -368,7 +395,7 @@ public void testGetClientStatusHappyCase() throws Exception { @Test public void testGetClientStatusNoPipeFound() throws Exception { final String userAgentSuffix = "kafka-provider/NONE"; - SimpleIngestManager manager = TestUtils.getManager("nopipe", userAgentSuffix); + SimpleIngestManager manager = TestUtils.getManager("nopipe", userAgentSuffix, null); try { manager.getClientStatus(null); } catch (IngestResponseException exception) { @@ -385,7 +412,7 @@ public void testIngestFilesWithClientInfo() throws Exception { // first lets call configure client API final String userAgentSuffix = "kafka-provider/NONE"; - SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix); + SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix, null); ConfigureClientResponse configureClientResponse = manager.configureClient(null); assertEquals(0L, configureClientResponse.getClientSequencer().longValue()); @@ -421,7 +448,7 @@ public void testIngestFilesWithClientInfoWithOldClientSequencer() throws Excepti // first lets call configure client API final String userAgentSuffix = "kafka-provider/NONE"; - SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix); + SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix, null); ConfigureClientResponse configureClientResponse = manager.configureClient(null); assertEquals(0L, configureClientResponse.getClientSequencer().longValue()); final long oldClientSequencer = configureClientResponse.getClientSequencer(); @@ -482,7 +509,7 @@ public void testIngestFilesWithClientInfoWithOldClientSequencer() throws Excepti public void testIngestFilesWithClientInfoWithNoClientSequencer() throws Exception { // first lets call configure client API final String userAgentSuffix = "kafka-provider/NONE"; - SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix); + SimpleIngestManager manager = TestUtils.getManager(pipeName, userAgentSuffix, null); // put TestUtils.executeQuery("put file://" + testFilePath + " @" + stageName); diff --git a/src/test/java/net/snowflake/ingest/TestUtils.java b/src/test/java/net/snowflake/ingest/TestUtils.java index 95dd01d2d..b9b8ae90a 100644 --- a/src/test/java/net/snowflake/ingest/TestUtils.java +++ b/src/test/java/net/snowflake/ingest/TestUtils.java @@ -164,6 +164,13 @@ public static KeyPair getKeyPair() throws Exception { return keyPair; } + public static String getRole() throws Exception { + if (profile == null) { + init(); + } + return role; + } + public static Properties getProperties() throws Exception { if (profile == null) { init(); @@ -247,11 +254,12 @@ public static SimpleIngestManager getManager(String pipe) throws Exception { * @param pipe pipe name * @param userAgentSuffix suffix we want to add in all request header of user-agent to the * snowpipe API. + * @param role the role that will be used for authorization * @return ingest manager object * @throws Exception */ - public static SimpleIngestManager getManager(String pipe, final String userAgentSuffix) - throws Exception { + public static SimpleIngestManager getManager( + String pipe, final String userAgentSuffix, String role) throws Exception { if (profile == null) init(); return new SimpleIngestManager( account, @@ -261,7 +269,8 @@ public static SimpleIngestManager getManager(String pipe, final String userAgent scheme, host, port, - userAgentSuffix); + userAgentSuffix, + role); } /**