Snowflake source oauth: add ability to support different roles during…

… oauth (#15654) * Snowflake source: add ability to support different roles during oauth * Snowflake source: add ability to support different roles during oauth * Snowflake source: handle case with empty role fields * Snowflake source: bump version for testing on dev * Snowflake source: bump version for testing on dev * auto-bump connector version [ci skip] * Snowflake source: bump version for testing on dev * Snowflake source: revert changes related to source not to core Co-authored-by: Octavia Squidington III <octavia-squidington-iii@users.noreply.github.com>
airbytehq · Aug 19, 2022 · afbe584 · afbe584
1 parent a3a6501
commit afbe584
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git a/airbyte-oauth/src/main/java/io/airbyte/oauth/flows/SourceSnowflakeOAuthFlow.java b/airbyte-oauth/src/main/java/io/airbyte/oauth/flows/SourceSnowflakeOAuthFlow.java
@@ -45,17 +45,27 @@ protected String formatConsentUrl(UUID definitionId,
                                     JsonNode inputOAuthConfiguration)
       throws IOException {
     try {
-      return new URIBuilder(String.format(AUTHORIZE_URL, extractUrl(inputOAuthConfiguration)))
+      String consentUrl = new URIBuilder(String.format(AUTHORIZE_URL, extractUrl(inputOAuthConfiguration)))
           .addParameter("client_id", clientId)
           .addParameter("redirect_uri", redirectUrl)
           .addParameter("response_type", "code")
           .addParameter("state", getState())
           .build().toString();
+      String providedRole = extractRole(inputOAuthConfiguration);
+      return providedRole.isEmpty()
+          ? consentUrl
+          : getConsentUrlWithScopeRole(consentUrl, providedRole);
     } catch (final URISyntaxException e) {
       throw new IOException("Failed to format Consent URL for OAuth flow", e);
     }
   }
 
+  private static String getConsentUrlWithScopeRole(String consentUrl, String providedRole) throws URISyntaxException {
+    return new URIBuilder(consentUrl)
+        .addParameter("scope", "session:role:" + providedRole)
+        .build().toString();
+  }
+
   @Override
   protected String getAccessTokenUrl(JsonNode inputOAuthConfiguration) {
     return String.format(ACCESS_TOKEN_URL, extractUrl(inputOAuthConfiguration));
@@ -141,4 +151,9 @@ private String extractUrl(JsonNode inputOAuthConfiguration) {
     return url == null ? "snowflakecomputing.com" : url.asText();
   }
 
+  private String extractRole(JsonNode inputOAuthConfiguration) {
+    var role = inputOAuthConfiguration.get("role");
+    return role == null ? "" : role.asText();
+  }
+
 }
diff --git a/airbyte-oauth/src/test/java/io/airbyte/oauth/flows/SnowflakeOAuthFlowTest.java b/airbyte-oauth/src/test/java/io/airbyte/oauth/flows/SnowflakeOAuthFlowTest.java
@@ -15,14 +15,17 @@
 @SuppressWarnings("PMD.JUnitTestsShouldIncludeAssert")
 class SnowflakeOAuthFlowTest extends BaseOAuthFlowTest {
 
+  public static final String STRING = "string";
+  public static final String TYPE = "type";
+
   @Override
   protected BaseOAuthFlow getOAuthFlow() {
     return new SourceSnowflakeOAuthFlow(getConfigRepository(), getHttpClient(), this::getConstantState);
   }
 
   @Override
   protected String getExpectedConsentUrl() {
-    return "https://account.aws.snowflakecomputing.com/oauth/authorize?client_id=test_client_id&redirect_uri=https%3A%2F%2Fairbyte.io&response_type=code&state=state";
+    return "https://account.aws.snowflakecomputing.com/oauth/authorize?client_id=test_client_id&redirect_uri=https%3A%2F%2Fairbyte.io&response_type=code&state=state&scope=session%3Arole%3Asome_role";
   }
 
   @Override
@@ -35,7 +38,7 @@ protected Map<String, String> getExpectedOutput() {
 
   @Override
   protected JsonNode getCompleteOAuthOutputSpecification() {
-    return getJsonSchema(Map.of("access_token", Map.of("type", "string"), "refresh_token", Map.of("type", "string")));
+    return getJsonSchema(Map.of("access_token", Map.of(TYPE, STRING), "refresh_token", Map.of(TYPE, STRING)));
   }
 
   @Override
@@ -58,12 +61,13 @@ protected JsonNode getOAuthParamConfig() {
   protected JsonNode getInputOAuthConfiguration() {
     return Jsons.jsonNode(ImmutableMap.builder()
         .put("host", "account.aws.snowflakecomputing.com")
+        .put("role", "some_role")
         .build());
   }
 
   @Override
   protected JsonNode getUserInputFromConnectorConfigSpecification() {
-    return getJsonSchema(Map.of("host", Map.of("type", "string")));
+    return getJsonSchema(Map.of("host", Map.of(TYPE, STRING), "role", Map.of(TYPE, STRING)));
   }
 
   @Test