You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by bl...@apache.org on 2023/12/10 21:57:38 UTC
(iceberg) branch main updated: Core: Add REST catalog table session cache (#8920)
This is an automated email from the ASF dual-hosted git repository.
blue pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/main by this push:
new 4090a88600 Core: Add REST catalog table session cache (#8920)
4090a88600 is described below
commit 4090a8860061f58748e3faa6804094f90d3575f3
Author: Eduard Tudenhoefner <et...@gmail.com>
AuthorDate: Sun Dec 10 22:57:32 2023 +0100
Core: Add REST catalog table session cache (#8920)
The purpose of caching auth session for tables is mainly so that we can
stop refreshing the session when it isn't used anymore
---
.../apache/iceberg/rest/RESTSessionCatalog.java | 56 +++++++++++++++++-----
.../org/apache/iceberg/rest/TestRESTCatalog.java | 2 +-
2 files changed, 45 insertions(+), 13 deletions(-)
diff --git a/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java b/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java
index a2c2b7dd70..5f660f0f4f 100644
--- a/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java
+++ b/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java
@@ -92,6 +92,7 @@ import org.apache.iceberg.rest.responses.LoadViewResponse;
import org.apache.iceberg.rest.responses.OAuthTokenResponse;
import org.apache.iceberg.rest.responses.UpdateNamespacePropertiesResponse;
import org.apache.iceberg.util.EnvironmentUtil;
+import org.apache.iceberg.util.Pair;
import org.apache.iceberg.util.PropertyUtil;
import org.apache.iceberg.util.ThreadPools;
import org.apache.iceberg.view.BaseView;
@@ -123,6 +124,7 @@ public class RESTSessionCatalog extends BaseViewSessionCatalog
private final Function<Map<String, String>, RESTClient> clientBuilder;
private final BiFunction<SessionContext, Map<String, String>, FileIO> ioBuilder;
private Cache<String, AuthSession> sessions = null;
+ private Cache<String, AuthSession> tableSessions = null;
private Cache<TableOperations, FileIO> fileIOCloser;
private AuthSession catalogAuth = null;
private boolean keepTokenRefreshed = true;
@@ -197,6 +199,7 @@ public class RESTSessionCatalog extends BaseViewSessionCatalog
Map<String, String> baseHeaders = configHeaders(mergedProps);
this.sessions = newSessionCache(mergedProps);
+ this.tableSessions = newSessionCache(mergedProps);
this.keepTokenRefreshed =
PropertyUtil.propertyAsBoolean(
mergedProps,
@@ -242,7 +245,15 @@ public class RESTSessionCatalog extends BaseViewSessionCatalog
AuthSession session =
sessions.get(
context.sessionId(),
- id -> newSession(context.credentials(), context.properties(), catalogAuth));
+ id -> {
+ Pair<String, Supplier<AuthSession>> newSession =
+ newSession(context.credentials(), context.properties(), catalogAuth);
+ if (null != newSession) {
+ return newSession.second().get();
+ }
+
+ return null;
+ });
return session != null ? session : catalogAuth;
}
@@ -859,7 +870,12 @@ public class RESTSessionCatalog extends BaseViewSessionCatalog
}
private AuthSession tableSession(Map<String, String> tableConf, AuthSession parent) {
- AuthSession session = newSession(tableConf, tableConf, parent);
+ Pair<String, Supplier<AuthSession>> newSession = newSession(tableConf, tableConf, parent);
+ if (null == newSession) {
+ return parent;
+ }
+
+ AuthSession session = tableSessions.get(newSession.first(), id -> newSession.second().get());
return session != null ? session : parent;
}
@@ -889,30 +905,46 @@ public class RESTSessionCatalog extends BaseViewSessionCatalog
return configResponse;
}
- private AuthSession newSession(
+ private Pair<String, Supplier<AuthSession>> newSession(
Map<String, String> credentials, Map<String, String> properties, AuthSession parent) {
if (credentials != null) {
// use the bearer token without exchanging
if (credentials.containsKey(OAuth2Properties.TOKEN)) {
- return AuthSession.fromAccessToken(
- client,
- tokenRefreshExecutor(),
+ return Pair.of(
credentials.get(OAuth2Properties.TOKEN),
- expiresAtMillis(properties),
- parent);
+ () ->
+ AuthSession.fromAccessToken(
+ client,
+ tokenRefreshExecutor(),
+ credentials.get(OAuth2Properties.TOKEN),
+ expiresAtMillis(properties),
+ parent));
}
if (credentials.containsKey(OAuth2Properties.CREDENTIAL)) {
// fetch a token using the client credentials flow
- return AuthSession.fromCredential(
- client, tokenRefreshExecutor(), credentials.get(OAuth2Properties.CREDENTIAL), parent);
+ return Pair.of(
+ credentials.get(OAuth2Properties.CREDENTIAL),
+ () ->
+ AuthSession.fromCredential(
+ client,
+ tokenRefreshExecutor(),
+ credentials.get(OAuth2Properties.CREDENTIAL),
+ parent));
}
for (String tokenType : TOKEN_PREFERENCE_ORDER) {
if (credentials.containsKey(tokenType)) {
// exchange the token for an access token using the token exchange flow
- return AuthSession.fromTokenExchange(
- client, tokenRefreshExecutor(), credentials.get(tokenType), tokenType, parent);
+ return Pair.of(
+ credentials.get(tokenType),
+ () ->
+ AuthSession.fromTokenExchange(
+ client,
+ tokenRefreshExecutor(),
+ credentials.get(tokenType),
+ tokenType,
+ parent));
}
}
}
diff --git a/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java b/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java
index 3de9c7b1d3..d0634be16c 100644
--- a/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java
+++ b/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java
@@ -1162,7 +1162,7 @@ public class TestRESTCatalog extends CatalogTests<RESTCatalog> {
// if the table returned a bearer token, there will be no token request
if (!tableConfig.containsKey("token")) {
// client credentials or token exchange to get a table token
- Mockito.verify(adapter, times(2))
+ Mockito.verify(adapter, times(1))
.execute(
eq(HTTPMethod.POST),
eq("v1/oauth/tokens"),