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 index 81891603ffa4..6f91f24e9c10 100644 --- 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 index e217b569aa67..97a663162893 100644 --- 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,6 +15,9 @@ @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); @@ -22,7 +25,7 @@ protected BaseOAuthFlow getOAuthFlow() { @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 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