Skip to content

Commit d5e80d0

Browse files
eliziarioeliziario
authored
OAuth backend Refactor (#7209)
1 parent e87604e commit d5e80d0

File tree

12 files changed

+165
-240
lines changed

12 files changed

+165
-240
lines changed

‎airbyte-oauth/src/main/java/io/airbyte/oauth/BaseOAuthConfig.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ protected JsonNode getDestinationOAuthParamConfig(final UUID workspaceId, final
6060
* Throws an exception if the client ID cannot be extracted. Subclasses should override this to
6161
* parse the config differently.
6262
*
63-
* @return
63+
* @return The configured Client ID used for this oauth flow
6464
*/
6565
protected String getClientIdUnsafe(final JsonNode oauthConfig) {
6666
if (oauthConfig.get("client_id") != null) {
@@ -74,7 +74,7 @@ protected String getClientIdUnsafe(final JsonNode oauthConfig) {
7474
* Throws an exception if the client secret cannot be extracted. Subclasses should override this to
7575
* parse the config differently.
7676
*
77-
* @return
77+
* @return The configured client secret for this OAuthFlow
7878
*/
7979
protected String getClientSecretUnsafe(final JsonNode oauthConfig) {
8080
if (oauthConfig.get("client_secret") != null) {

‎airbyte-oauth/src/main/java/io/airbyte/oauth/BaseOAuthFlow.java

+103-16
Original file line numberDiff line numberDiff line change
@@ -6,38 +6,80 @@
66

77
import com.fasterxml.jackson.databind.JsonNode;
88
import com.google.common.collect.ImmutableMap;
9+
import com.google.gson.Gson;
10+
import com.google.gson.reflect.TypeToken;
911
import io.airbyte.commons.json.Jsons;
1012
import io.airbyte.config.persistence.ConfigNotFoundException;
1113
import io.airbyte.config.persistence.ConfigRepository;
1214
import java.io.IOException;
15+
import java.lang.reflect.Type;
1316
import java.net.URI;
17+
import java.net.URISyntaxException;
1418
import java.net.URLEncoder;
1519
import java.net.http.HttpClient;
1620
import java.net.http.HttpClient.Version;
1721
import java.net.http.HttpRequest;
1822
import java.net.http.HttpResponse;
1923
import java.nio.charset.StandardCharsets;
24+
import java.util.HashMap;
2025
import java.util.Map;
2126
import java.util.UUID;
27+
import java.util.function.Function;
2228
import java.util.function.Supplier;
2329
import org.apache.commons.lang3.RandomStringUtils;
30+
import org.apache.http.client.utils.URIBuilder;
2431

2532
/*
2633
* Class implementing generic oAuth 2.0 flow.
2734
*/
2835
public abstract class BaseOAuthFlow extends BaseOAuthConfig {
2936

30-
private final HttpClient httpClient;
37+
/**
38+
* Simple enum of content type strings and their respective encoding functions used for POSTing the
39+
* access token request
40+
*/
41+
public enum TOKEN_REQUEST_CONTENT_TYPE {
42+
43+
URL_ENCODED("application/x-www-form-urlencoded", BaseOAuthFlow::toUrlEncodedString),
44+
JSON("application/json", BaseOAuthFlow::toJson);
45+
46+
String contentType;
47+
Function<Map<String, String>, String> converter;
48+
49+
TOKEN_REQUEST_CONTENT_TYPE(String contentType, Function<Map<String, String>, String> converter) {
50+
this.contentType = contentType;
51+
this.converter = converter;
52+
}
53+
54+
}
55+
56+
protected final HttpClient httpClient;
57+
private final TOKEN_REQUEST_CONTENT_TYPE tokenReqContentType;
3158
private final Supplier<String> stateSupplier;
3259

3360
public BaseOAuthFlow(final ConfigRepository configRepository) {
3461
this(configRepository, HttpClient.newBuilder().version(Version.HTTP_1_1).build(), BaseOAuthFlow::generateRandomState);
3562
}
3663

37-
public BaseOAuthFlow(final ConfigRepository configRepository, final HttpClient httpClient, final Supplier<String> stateSupplier) {
64+
public BaseOAuthFlow(ConfigRepository configRepository, TOKEN_REQUEST_CONTENT_TYPE tokenReqContentType) {
65+
this(configRepository,
66+
HttpClient.newBuilder().version(Version.HTTP_1_1).build(),
67+
BaseOAuthFlow::generateRandomState,
68+
tokenReqContentType);
69+
}
70+
71+
public BaseOAuthFlow(ConfigRepository configRepository, HttpClient httpClient, Supplier<String> stateSupplier) {
72+
this(configRepository, httpClient, stateSupplier, TOKEN_REQUEST_CONTENT_TYPE.URL_ENCODED);
73+
}
74+
75+
public BaseOAuthFlow(ConfigRepository configRepository,
76+
HttpClient httpClient,
77+
Supplier<String> stateSupplier,
78+
TOKEN_REQUEST_CONTENT_TYPE tokenReqContentType) {
3879
super(configRepository);
3980
this.httpClient = httpClient;
4081
this.stateSupplier = stateSupplier;
82+
this.tokenReqContentType = tokenReqContentType;
4183
}
4284

4385
@Override
@@ -54,6 +96,31 @@ public String getDestinationConsentUrl(final UUID workspaceId, final UUID destin
5496
return formatConsentUrl(destinationDefinitionId, getClientIdUnsafe(oAuthParamConfig), redirectUrl);
5597
}
5698

99+
protected String formatConsentUrl(String clientId,
100+
String redirectUrl,
101+
String host,
102+
String path,
103+
String scope,
104+
String responseType)
105+
throws IOException {
106+
final URIBuilder builder = new URIBuilder()
107+
.setScheme("https")
108+
.setHost(host)
109+
.setPath(path)
110+
// required
111+
.addParameter("client_id", clientId)
112+
.addParameter("redirect_uri", redirectUrl)
113+
.addParameter("state", getState())
114+
// optional
115+
.addParameter("response_type", responseType)
116+
.addParameter("scope", scope);
117+
try {
118+
return builder.build().toString();
119+
} catch (URISyntaxException e) {
120+
throw new IOException("Failed to format Consent URL for OAuth flow", e);
121+
}
122+
}
123+
57124
/**
58125
* Depending on the OAuth flow implementation, the URL to grant user's consent may differ,
59126
* especially in the query parameters to be provided. This function should generate such consent URL
@@ -84,7 +151,8 @@ public Map<String, Object> completeSourceOAuth(
84151
getClientIdUnsafe(oAuthParamConfig),
85152
getClientSecretUnsafe(oAuthParamConfig),
86153
extractCodeParameter(queryParams),
87-
redirectUrl);
154+
redirectUrl,
155+
oAuthParamConfig);
88156
}
89157

90158
@Override
@@ -98,20 +166,26 @@ public Map<String, Object> completeDestinationOAuth(final UUID workspaceId,
98166
getClientIdUnsafe(oAuthParamConfig),
99167
getClientSecretUnsafe(oAuthParamConfig),
100168
extractCodeParameter(queryParams),
101-
redirectUrl);
169+
redirectUrl, oAuthParamConfig);
102170
}
103171

104-
private Map<String, Object> completeOAuthFlow(final String clientId, final String clientSecret, final String authCode, final String redirectUrl)
172+
private Map<String, Object> completeOAuthFlow(final String clientId,
173+
final String clientSecret,
174+
final String authCode,
175+
final String redirectUrl,
176+
JsonNode oAuthParamConfig)
105177
throws IOException {
178+
var accessTokenUrl = getAccessTokenUrl(oAuthParamConfig);
106179
final HttpRequest request = HttpRequest.newBuilder()
107-
.POST(HttpRequest.BodyPublishers.ofString(toUrlEncodedString(getAccessTokenQueryParameters(clientId, clientSecret, authCode, redirectUrl))))
108-
.uri(URI.create(getAccessTokenUrl()))
109-
.header("Content-Type", "application/x-www-form-urlencoded")
180+
.POST(HttpRequest.BodyPublishers
181+
.ofString(tokenReqContentType.converter.apply(getAccessTokenQueryParameters(clientId, clientSecret, authCode, redirectUrl))))
182+
.uri(URI.create(accessTokenUrl))
183+
.header("Content-Type", tokenReqContentType.contentType)
110184
.build();
111185
// TODO: Handle error response to report better messages
112186
try {
113-
final HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());;
114-
return extractRefreshToken(Jsons.deserialize(response.body()));
187+
final HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
188+
return extractRefreshToken(Jsons.deserialize(response.body()), accessTokenUrl);
115189
} catch (final InterruptedException e) {
116190
throw new IOException("Failed to complete OAuth flow", e);
117191
}
@@ -146,13 +220,20 @@ protected String extractCodeParameter(Map<String, Object> queryParams) throws IO
146220
/**
147221
* Returns the URL where to retrieve the access token from.
148222
*/
149-
protected abstract String getAccessTokenUrl();
223+
protected abstract String getAccessTokenUrl(JsonNode oAuthParamConfig);
150224

151-
/**
152-
* Once the auth code is exchange for a refresh token, the oauth flow implementation can extract and
153-
* returns the values of fields to be used in the connector's configurations.
154-
*/
155-
protected abstract Map<String, Object> extractRefreshToken(JsonNode data) throws IOException;
225+
protected Map<String, Object> extractRefreshToken(final JsonNode data, String accessTokenUrl) throws IOException {
226+
final Map<String, Object> result = new HashMap<>();
227+
if (data.has("refresh_token")) {
228+
result.put("refresh_token", data.get("refresh_token").asText());
229+
} else if (data.has("access_token")) {
230+
result.put("access_token", data.get("access_token").asText());
231+
} else {
232+
throw new IOException(String.format("Missing 'refresh_token' in query params from %s", accessTokenUrl));
233+
}
234+
return Map.of("credentials", result);
235+
236+
}
156237

157238
private static String urlEncode(final String s) {
158239
try {
@@ -173,4 +254,10 @@ private static String toUrlEncodedString(final Map<String, String> body) {
173254
return result.toString();
174255
}
175256

257+
protected static String toJson(final Map<String, String> body) {
258+
final Gson gson = new Gson();
259+
Type gsonType = new TypeToken<Map<String, String>>() {}.getType();
260+
return gson.toJson(body, gsonType);
261+
}
262+
176263
}

‎airbyte-oauth/src/main/java/io/airbyte/oauth/OAuthImplementationFactory.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import io.airbyte.oauth.flows.google.GoogleAnalyticsOAuthFlow;
1515
import io.airbyte.oauth.flows.google.GoogleSearchConsoleOAuthFlow;
1616
import java.util.Map;
17+
import java.util.UUID;
1718

1819
public class OAuthImplementationFactory {
1920

@@ -31,7 +32,7 @@ public OAuthImplementationFactory(final ConfigRepository configRepository) {
3132
.build();
3233
}
3334

34-
public OAuthFlowImplementation create(final String imageName) {
35+
public OAuthFlowImplementation create(final String imageName, final UUID workspaceId) {
3536
if (OAUTH_FLOW_MAPPING.containsKey(imageName)) {
3637
return OAUTH_FLOW_MAPPING.get(imageName);
3738
} else {

‎airbyte-oauth/src/main/java/io/airbyte/oauth/flows/AsanaOAuthFlow.java

+1-13
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import com.fasterxml.jackson.databind.JsonNode;
88
import com.google.common.annotations.VisibleForTesting;
99
import com.google.common.collect.ImmutableMap;
10-
import io.airbyte.commons.json.Jsons;
1110
import io.airbyte.config.persistence.ConfigRepository;
1211
import io.airbyte.oauth.BaseOAuthFlow;
1312
import java.io.IOException;
@@ -50,7 +49,7 @@ protected String formatConsentUrl(UUID definitionId, String clientId, String red
5049
}
5150

5251
@Override
53-
protected String getAccessTokenUrl() {
52+
protected String getAccessTokenUrl(JsonNode oAuthParamConfig) {
5453
return ACCESS_TOKEN_URL;
5554
}
5655

@@ -62,15 +61,4 @@ protected Map<String, String> getAccessTokenQueryParameters(String clientId, Str
6261
.build();
6362
}
6463

65-
@Override
66-
protected Map<String, Object> extractRefreshToken(JsonNode data) throws IOException {
67-
System.out.println(Jsons.serialize(data));
68-
if (data.has("refresh_token")) {
69-
final String refreshToken = data.get("refresh_token").asText();
70-
return Map.of("credentials", Map.of("refresh_token", refreshToken));
71-
} else {
72-
throw new IOException(String.format("Missing 'refresh_token' in query params from %s", ACCESS_TOKEN_URL));
73-
}
74-
}
75-
7664
}

‎airbyte-oauth/src/main/java/io/airbyte/oauth/flows/FacebookMarketingOAuthFlow.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -54,12 +54,12 @@ protected String formatConsentUrl(final UUID definitionId, final String clientId
5454
}
5555

5656
@Override
57-
protected String getAccessTokenUrl() {
57+
protected String getAccessTokenUrl(JsonNode oAuthParamConfig) {
5858
return ACCESS_TOKEN_URL;
5959
}
6060

6161
@Override
62-
protected Map<String, Object> extractRefreshToken(final JsonNode data) throws IOException {
62+
protected Map<String, Object> extractRefreshToken(final JsonNode data, String accessTokenUrl) throws IOException {
6363
// Facebook does not have refresh token but calls it "long lived access token" instead:
6464
// see https://developers.facebook.com/docs/facebook-login/access-tokens/refreshing
6565
if (data.has("access_token")) {

‎airbyte-oauth/src/main/java/io/airbyte/oauth/flows/google/GoogleAdsOAuthFlow.java

-8
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,7 @@
88
import com.google.common.annotations.VisibleForTesting;
99
import com.google.common.base.Preconditions;
1010
import io.airbyte.config.persistence.ConfigRepository;
11-
import java.io.IOException;
1211
import java.net.http.HttpClient;
13-
import java.util.Map;
1412
import java.util.function.Supplier;
1513

1614
public class GoogleAdsOAuthFlow extends GoogleOAuthFlow {
@@ -46,10 +44,4 @@ protected String getClientSecretUnsafe(final JsonNode config) {
4644
return super.getClientSecretUnsafe(config.get("credentials"));
4745
}
4846

49-
@Override
50-
protected Map<String, Object> extractRefreshToken(final JsonNode data) throws IOException {
51-
// the config object containing refresh token is nested inside the "credentials" object
52-
return Map.of("credentials", super.extractRefreshToken(data));
53-
}
54-
5547
}

‎airbyte-oauth/src/main/java/io/airbyte/oauth/flows/google/GoogleAnalyticsOAuthFlow.java

-8
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,7 @@
88
import com.google.common.annotations.VisibleForTesting;
99
import com.google.common.base.Preconditions;
1010
import io.airbyte.config.persistence.ConfigRepository;
11-
import java.io.IOException;
1211
import java.net.http.HttpClient;
13-
import java.util.Map;
1412
import java.util.function.Supplier;
1513

1614
public class GoogleAnalyticsOAuthFlow extends GoogleOAuthFlow {
@@ -45,10 +43,4 @@ protected String getClientSecretUnsafe(final JsonNode config) {
4543
return super.getClientSecretUnsafe(config.get("credentials"));
4644
}
4745

48-
@Override
49-
protected Map<String, Object> extractRefreshToken(final JsonNode data) throws IOException {
50-
// the config object containing refresh token is nested inside the "credentials" object
51-
return Map.of("credentials", super.extractRefreshToken(data));
52-
}
53-
5446
}

‎airbyte-oauth/src/main/java/io/airbyte/oauth/flows/google/GoogleOAuthFlow.java

+1-16
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import java.io.IOException;
1313
import java.net.URISyntaxException;
1414
import java.net.http.HttpClient;
15-
import java.util.HashMap;
1615
import java.util.Map;
1716
import java.util.UUID;
1817
import java.util.function.Supplier;
@@ -64,7 +63,7 @@ protected String formatConsentUrl(final UUID definitionId, final String clientId
6463
protected abstract String getScope();
6564

6665
@Override
67-
protected String getAccessTokenUrl() {
66+
protected String getAccessTokenUrl(JsonNode oAuthParamConfig) {
6867
return ACCESS_TOKEN_URL;
6968
}
7069

@@ -82,18 +81,4 @@ protected Map<String, String> getAccessTokenQueryParameters(final String clientI
8281
.build();
8382
}
8483

85-
@Override
86-
protected Map<String, Object> extractRefreshToken(final JsonNode data) throws IOException {
87-
final Map<String, Object> result = new HashMap<>();
88-
if (data.has("access_token")) {
89-
result.put("access_token", data.get("access_token").asText());
90-
}
91-
if (data.has("refresh_token")) {
92-
result.put("refresh_token", data.get("refresh_token").asText());
93-
} else {
94-
throw new IOException(String.format("Missing 'refresh_token' in query params from %s", ACCESS_TOKEN_URL));
95-
}
96-
return result;
97-
}
98-
9984
}

‎airbyte-oauth/src/main/java/io/airbyte/oauth/flows/google/GoogleSearchConsoleOAuthFlow.java

+7-2
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import io.airbyte.config.persistence.ConfigRepository;
1111
import java.io.IOException;
1212
import java.net.http.HttpClient;
13+
import java.util.HashMap;
1314
import java.util.Map;
1415
import java.util.function.Supplier;
1516

@@ -47,9 +48,13 @@ protected String getClientSecretUnsafe(final JsonNode config) {
4748
}
4849

4950
@Override
50-
protected Map<String, Object> extractRefreshToken(final JsonNode data) throws IOException {
51+
protected Map<String, Object> extractRefreshToken(final JsonNode data, String accessTokenUrl) throws IOException {
5152
// the config object containing refresh token is nested inside the "authorization" object
52-
return Map.of("authorization", super.extractRefreshToken(data));
53+
final Map<String, Object> result = new HashMap<>();
54+
if (data.has("refresh_token")) {
55+
result.put("refresh_token", data.get("refresh_token").asText());
56+
}
57+
return Map.of("authorization", result);
5358
}
5459

5560
}

0 commit comments

Comments
 (0)