Skip to content

Commit afbe584

Browse files
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 <[email protected]>
1 parent a3a6501 commit afbe584

File tree

2 files changed

+23
-4
lines changed

2 files changed

+23
-4
lines changed

airbyte-oauth/src/main/java/io/airbyte/oauth/flows/SourceSnowflakeOAuthFlow.java

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,17 +45,27 @@ protected String formatConsentUrl(UUID definitionId,
4545
JsonNode inputOAuthConfiguration)
4646
throws IOException {
4747
try {
48-
return new URIBuilder(String.format(AUTHORIZE_URL, extractUrl(inputOAuthConfiguration)))
48+
String consentUrl = new URIBuilder(String.format(AUTHORIZE_URL, extractUrl(inputOAuthConfiguration)))
4949
.addParameter("client_id", clientId)
5050
.addParameter("redirect_uri", redirectUrl)
5151
.addParameter("response_type", "code")
5252
.addParameter("state", getState())
5353
.build().toString();
54+
String providedRole = extractRole(inputOAuthConfiguration);
55+
return providedRole.isEmpty()
56+
? consentUrl
57+
: getConsentUrlWithScopeRole(consentUrl, providedRole);
5458
} catch (final URISyntaxException e) {
5559
throw new IOException("Failed to format Consent URL for OAuth flow", e);
5660
}
5761
}
5862

63+
private static String getConsentUrlWithScopeRole(String consentUrl, String providedRole) throws URISyntaxException {
64+
return new URIBuilder(consentUrl)
65+
.addParameter("scope", "session:role:" + providedRole)
66+
.build().toString();
67+
}
68+
5969
@Override
6070
protected String getAccessTokenUrl(JsonNode inputOAuthConfiguration) {
6171
return String.format(ACCESS_TOKEN_URL, extractUrl(inputOAuthConfiguration));
@@ -141,4 +151,9 @@ private String extractUrl(JsonNode inputOAuthConfiguration) {
141151
return url == null ? "snowflakecomputing.com" : url.asText();
142152
}
143153

154+
private String extractRole(JsonNode inputOAuthConfiguration) {
155+
var role = inputOAuthConfiguration.get("role");
156+
return role == null ? "" : role.asText();
157+
}
158+
144159
}

airbyte-oauth/src/test/java/io/airbyte/oauth/flows/SnowflakeOAuthFlowTest.java

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,17 @@
1515
@SuppressWarnings("PMD.JUnitTestsShouldIncludeAssert")
1616
class SnowflakeOAuthFlowTest extends BaseOAuthFlowTest {
1717

18+
public static final String STRING = "string";
19+
public static final String TYPE = "type";
20+
1821
@Override
1922
protected BaseOAuthFlow getOAuthFlow() {
2023
return new SourceSnowflakeOAuthFlow(getConfigRepository(), getHttpClient(), this::getConstantState);
2124
}
2225

2326
@Override
2427
protected String getExpectedConsentUrl() {
25-
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";
28+
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";
2629
}
2730

2831
@Override
@@ -35,7 +38,7 @@ protected Map<String, String> getExpectedOutput() {
3538

3639
@Override
3740
protected JsonNode getCompleteOAuthOutputSpecification() {
38-
return getJsonSchema(Map.of("access_token", Map.of("type", "string"), "refresh_token", Map.of("type", "string")));
41+
return getJsonSchema(Map.of("access_token", Map.of(TYPE, STRING), "refresh_token", Map.of(TYPE, STRING)));
3942
}
4043

4144
@Override
@@ -58,12 +61,13 @@ protected JsonNode getOAuthParamConfig() {
5861
protected JsonNode getInputOAuthConfiguration() {
5962
return Jsons.jsonNode(ImmutableMap.builder()
6063
.put("host", "account.aws.snowflakecomputing.com")
64+
.put("role", "some_role")
6165
.build());
6266
}
6367

6468
@Override
6569
protected JsonNode getUserInputFromConnectorConfigSpecification() {
66-
return getJsonSchema(Map.of("host", Map.of("type", "string")));
70+
return getJsonSchema(Map.of("host", Map.of(TYPE, STRING), "role", Map.of(TYPE, STRING)));
6771
}
6872

6973
@Test

0 commit comments

Comments
 (0)