Skip to content

Commit 4816bf3

Browse files
adutraadutra
authored
AWS, Core, GCP: Auth Manager API enablement (#12197)
1 parent 57ec405 commit 4816bf3

File tree

14 files changed

+603
-713
lines changed

14 files changed

+603
-713
lines changed

Diff for: ‎aws/src/main/java/org/apache/iceberg/aws/RESTSigV4Signer.java

+3
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,10 @@
5151
* <p>See <a
5252
* href="https://docs.aws.amazon.com/general/latest/gr/signing-aws-api-requests.html">Signing AWS
5353
* API requests</a> for details about the protocol.
54+
*
55+
* @deprecated since 1.9.0, will be removed in 1.10.0; use {@link RESTSigV4AuthManager} instead.
5456
*/
57+
@Deprecated
5558
public class RESTSigV4Signer implements HttpRequestInterceptor {
5659
static final String EMPTY_BODY_SHA256 =
5760
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";

Diff for: ‎aws/src/main/java/org/apache/iceberg/aws/s3/VendedCredentialsProvider.java

+13-14
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,10 @@
2828
import org.apache.iceberg.relocated.com.google.common.base.Strings;
2929
import org.apache.iceberg.rest.ErrorHandlers;
3030
import org.apache.iceberg.rest.HTTPClient;
31-
import org.apache.iceberg.rest.HTTPHeaders;
3231
import org.apache.iceberg.rest.RESTClient;
33-
import org.apache.iceberg.rest.auth.DefaultAuthSession;
34-
import org.apache.iceberg.rest.auth.OAuth2Properties;
35-
import org.apache.iceberg.rest.auth.OAuth2Util;
32+
import org.apache.iceberg.rest.auth.AuthManager;
33+
import org.apache.iceberg.rest.auth.AuthManagers;
34+
import org.apache.iceberg.rest.auth.AuthSession;
3635
import org.apache.iceberg.rest.credentials.Credential;
3736
import org.apache.iceberg.rest.responses.LoadCredentialsResponse;
3837
import software.amazon.awssdk.auth.credentials.AwsCredentials;
@@ -48,6 +47,8 @@ public class VendedCredentialsProvider implements AwsCredentialsProvider, SdkAut
4847
private volatile HTTPClient client;
4948
private final Map<String, String> properties;
5049
private final CachedSupplier<AwsCredentials> credentialCache;
50+
private AuthManager authManager;
51+
private AuthSession authSession;
5152

5253
private VendedCredentialsProvider(Map<String, String> properties) {
5354
Preconditions.checkArgument(null != properties, "Invalid properties: null");
@@ -66,8 +67,10 @@ public AwsCredentials resolveCredentials() {
6667

6768
@Override
6869
public void close() {
69-
IoUtils.closeQuietly(client, null);
70-
credentialCache.close();
70+
IoUtils.closeQuietlyV2(authSession, null);
71+
IoUtils.closeQuietlyV2(authManager, null);
72+
IoUtils.closeQuietlyV2(client, null);
73+
IoUtils.closeQuietlyV2(credentialCache, null);
7174
}
7275

7376
public static VendedCredentialsProvider create(Map<String, String> properties) {
@@ -78,14 +81,10 @@ private RESTClient httpClient() {
7881
if (null == client) {
7982
synchronized (this) {
8083
if (null == client) {
81-
DefaultAuthSession authSession =
82-
DefaultAuthSession.of(
83-
HTTPHeaders.of(OAuth2Util.authHeaders(properties.get(OAuth2Properties.TOKEN))));
84-
client =
85-
HTTPClient.builder(properties)
86-
.uri(properties.get(URI))
87-
.withAuthSession(authSession)
88-
.build();
84+
authManager = AuthManagers.loadAuthManager("s3-credentials-refresh", properties);
85+
HTTPClient httpClient = HTTPClient.builder(properties).uri(properties.get(URI)).build();
86+
authSession = authManager.catalogSession(httpClient, properties);
87+
client = httpClient.withAuthSession(authSession);
8988
}
9089
}
9190
}

Diff for: ‎aws/src/main/java/org/apache/iceberg/aws/s3/signer/S3V4RestSignerClient.java

+41-130
Original file line numberDiff line numberDiff line change
@@ -20,15 +20,12 @@
2020

2121
import com.github.benmanes.caffeine.cache.Cache;
2222
import com.github.benmanes.caffeine.cache.Caffeine;
23-
import com.github.benmanes.caffeine.cache.RemovalListener;
2423
import java.io.IOException;
2524
import java.io.InputStream;
2625
import java.net.URI;
27-
import java.time.Duration;
2826
import java.util.Collections;
2927
import java.util.List;
3028
import java.util.Map;
31-
import java.util.concurrent.ScheduledExecutorService;
3229
import java.util.concurrent.TimeUnit;
3330
import java.util.function.Consumer;
3431
import java.util.function.Supplier;
@@ -42,13 +39,12 @@
4239
import org.apache.iceberg.rest.HTTPClient;
4340
import org.apache.iceberg.rest.RESTClient;
4441
import org.apache.iceberg.rest.ResourcePaths;
45-
import org.apache.iceberg.rest.auth.AuthConfig;
42+
import org.apache.iceberg.rest.auth.AuthManager;
43+
import org.apache.iceberg.rest.auth.AuthManagers;
44+
import org.apache.iceberg.rest.auth.AuthSession;
4645
import org.apache.iceberg.rest.auth.OAuth2Properties;
4746
import org.apache.iceberg.rest.auth.OAuth2Util;
48-
import org.apache.iceberg.rest.auth.OAuth2Util.AuthSession;
49-
import org.apache.iceberg.rest.responses.OAuthTokenResponse;
5047
import org.apache.iceberg.util.PropertyUtil;
51-
import org.apache.iceberg.util.ThreadPools;
5248
import org.immutables.value.Value;
5349
import org.slf4j.Logger;
5450
import org.slf4j.LoggerFactory;
@@ -64,7 +60,7 @@
6460

6561
@Value.Immutable
6662
public abstract class S3V4RestSignerClient
67-
extends AbstractAws4Signer<AwsS3V4SignerParams, Aws4PresignerParams> {
63+
extends AbstractAws4Signer<AwsS3V4SignerParams, Aws4PresignerParams> implements AutoCloseable {
6864

6965
private static final Logger LOG = LoggerFactory.getLogger(S3V4RestSignerClient.class);
7066
public static final String S3_SIGNER_URI = "s3.signer.uri";
@@ -81,13 +77,14 @@ public abstract class S3V4RestSignerClient
8177
private static final String SCOPE = "sign";
8278

8379
@SuppressWarnings("immutables:incompat")
84-
private static volatile ScheduledExecutorService tokenRefreshExecutor;
80+
private volatile AuthManager authManager;
8581

86-
@SuppressWarnings("immutables:incompat")
87-
private static volatile RESTClient httpClient;
82+
@SuppressWarnings({"immutables:incompat", "VisibilityModifier"})
83+
@VisibleForTesting
84+
static volatile RESTClient httpClient;
8885

8986
@SuppressWarnings("immutables:incompat")
90-
private static volatile Cache<String, AuthSession> authSessionCache;
87+
private volatile AuthSession authSession;
9188

9289
public abstract Map<String, String> properties();
9390

@@ -138,52 +135,6 @@ boolean keepTokenRefreshed() {
138135
OAuth2Properties.TOKEN_REFRESH_ENABLED_DEFAULT);
139136
}
140137

141-
@VisibleForTesting
142-
ScheduledExecutorService tokenRefreshExecutor() {
143-
if (!keepTokenRefreshed()) {
144-
return null;
145-
}
146-
147-
if (null == tokenRefreshExecutor) {
148-
synchronized (S3V4RestSignerClient.class) {
149-
if (null == tokenRefreshExecutor) {
150-
tokenRefreshExecutor = ThreadPools.newScheduledPool("s3-signer-token-refresh", 1);
151-
}
152-
}
153-
}
154-
155-
return tokenRefreshExecutor;
156-
}
157-
158-
private Cache<String, AuthSession> authSessionCache() {
159-
if (null == authSessionCache) {
160-
synchronized (S3V4RestSignerClient.class) {
161-
if (null == authSessionCache) {
162-
long expirationIntervalMs =
163-
PropertyUtil.propertyAsLong(
164-
properties(),
165-
CatalogProperties.AUTH_SESSION_TIMEOUT_MS,
166-
CatalogProperties.AUTH_SESSION_TIMEOUT_MS_DEFAULT);
167-
168-
authSessionCache =
169-
Caffeine.newBuilder()
170-
.expireAfterAccess(Duration.ofMillis(expirationIntervalMs))
171-
.removalListener(
172-
(RemovalListener<String, AuthSession>)
173-
(id, auth, cause) -> {
174-
if (null != auth) {
175-
LOG.trace("Stopping refresh for AuthSession");
176-
auth.stopRefreshing();
177-
}
178-
})
179-
.build();
180-
}
181-
}
182-
}
183-
184-
return authSessionCache;
185-
}
186-
187138
private RESTClient httpClient() {
188139
if (null == httpClient) {
189140
synchronized (S3V4RestSignerClient.class) {
@@ -200,86 +151,40 @@ private RESTClient httpClient() {
200151
return httpClient;
201152
}
202153

203-
private AuthSession authSession() {
204-
String token = token().get();
205-
if (null != token) {
206-
return authSessionCache()
207-
.get(
208-
token,
209-
id -> {
210-
// this client will be reused for token refreshes; it must contain an empty auth
211-
// session in order to avoid interfering with refreshed tokens
212-
RESTClient refreshClient =
213-
httpClient().withAuthSession(org.apache.iceberg.rest.auth.AuthSession.EMPTY);
214-
return AuthSession.fromAccessToken(
215-
refreshClient,
216-
tokenRefreshExecutor(),
217-
token,
218-
expiresAtMillis(properties()),
219-
new AuthSession(
220-
ImmutableMap.of(),
221-
AuthConfig.builder()
222-
.token(token)
223-
.credential(credential())
224-
.scope(SCOPE)
225-
.oauth2ServerUri(oauth2ServerUri())
226-
.optionalOAuthParams(optionalOAuthParams())
227-
.build()));
228-
});
229-
}
230-
231-
if (credentialProvided()) {
232-
return authSessionCache()
233-
.get(
234-
credential(),
235-
id -> {
236-
AuthSession session =
237-
new AuthSession(
238-
ImmutableMap.of(),
239-
AuthConfig.builder()
240-
.credential(credential())
241-
.scope(SCOPE)
242-
.oauth2ServerUri(oauth2ServerUri())
243-
.optionalOAuthParams(optionalOAuthParams())
244-
.build());
245-
long startTimeMillis = System.currentTimeMillis();
246-
// this client will be reused for token refreshes; it must contain an empty auth
247-
// session in order to avoid interfering with refreshed tokens
248-
RESTClient refreshClient =
249-
httpClient().withAuthSession(org.apache.iceberg.rest.auth.AuthSession.EMPTY);
250-
OAuthTokenResponse authResponse =
251-
OAuth2Util.fetchToken(
252-
refreshClient,
253-
session.headers(),
254-
credential(),
255-
SCOPE,
256-
oauth2ServerUri(),
257-
optionalOAuthParams());
258-
return AuthSession.fromTokenResponse(
259-
refreshClient, tokenRefreshExecutor(), authResponse, startTimeMillis, session);
260-
});
154+
@VisibleForTesting
155+
AuthSession authSession() {
156+
if (null == authSession) {
157+
synchronized (S3V4RestSignerClient.class) {
158+
if (null == authSession) {
159+
authManager = AuthManagers.loadAuthManager("s3-signer", properties());
160+
ImmutableMap.Builder<String, String> properties =
161+
ImmutableMap.<String, String>builder()
162+
.putAll(properties())
163+
.putAll(optionalOAuthParams())
164+
.put(OAuth2Properties.OAUTH2_SERVER_URI, oauth2ServerUri())
165+
.put(OAuth2Properties.TOKEN_REFRESH_ENABLED, String.valueOf(keepTokenRefreshed()))
166+
.put(OAuth2Properties.SCOPE, SCOPE);
167+
String token = token().get();
168+
if (null != token) {
169+
properties.put(OAuth2Properties.TOKEN, token);
170+
}
171+
172+
if (credentialProvided()) {
173+
properties.put(OAuth2Properties.CREDENTIAL, credential());
174+
}
175+
176+
authSession = authManager.tableSession(httpClient(), properties.buildKeepingLast());
177+
}
178+
}
261179
}
262180

263-
return AuthSession.empty();
181+
return authSession;
264182
}
265183

266184
private boolean credentialProvided() {
267185
return null != credential() && !credential().isEmpty();
268186
}
269187

270-
private Long expiresAtMillis(Map<String, String> properties) {
271-
if (properties.containsKey(OAuth2Properties.TOKEN_EXPIRES_IN_MS)) {
272-
long expiresInMillis =
273-
PropertyUtil.propertyAsLong(
274-
properties,
275-
OAuth2Properties.TOKEN_EXPIRES_IN_MS,
276-
OAuth2Properties.TOKEN_EXPIRES_IN_MS_DEFAULT);
277-
return System.currentTimeMillis() + expiresInMillis;
278-
} else {
279-
return null;
280-
}
281-
}
282-
283188
@Value.Check
284189
protected void check() {
285190
Preconditions.checkArgument(
@@ -377,6 +282,12 @@ public SdkHttpFullRequest sign(
377282
return mutableRequest.build();
378283
}
379284

285+
@Override
286+
public void close() throws Exception {
287+
IoUtils.closeQuietlyV2(authSession, null);
288+
IoUtils.closeQuietlyV2(authManager, null);
289+
}
290+
380291
/**
381292
* Only add body for DeleteObjectsRequest. Refer to
382293
* https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html#API_DeleteObjects_RequestSyntax

0 commit comments

Comments
 (0)