Add support for session timeout and token refresh in Iceberg REST cat…

…alog Allow configuring authentication session time-to-live and token refresh mechanics in RestSessionCatalog. This enables Trino to pass these properties to RestSessionCatalog instead of using default values. - `iceberg.rest-catalog.session-timeout` sets the duration for keeping an authentication session in cache. - `iceberg.rest-catalog.oauth2.token-refresh-enabled` controls whether a token should be refreshed if its expiration time is available.
trinodb · Feb 26, 2025 · 4c95b4c · 4c95b4c
1 parent a5f5d61
commit 4c95b4c
Show file tree

Hide file tree

Showing 7 changed files with 55 additions and 2 deletions.
diff --git a/docs/src/main/sphinx/object-storage/metastores.md b/docs/src/main/sphinx/object-storage/metastores.md
@@ -487,6 +487,8 @@ following properties:
 * - `iceberg.rest-catalog.session`
   - Session information included when communicating with the REST Catalog.
     Options are `NONE` or `USER` (default: `NONE`).
+* - `iceberg.rest-catalog.session-timeout`
+  - [Duration](prop-type-duration) to keep authentication session in cache. Defaults to `1h`.
 * - `iceberg.rest-catalog.oauth2.token`
   - The bearer token used for interactions with the server. A `token` or
     `credential` is required for `OAUTH2` security. Example: `AbCdEf123456`
@@ -499,6 +501,9 @@ following properties:
     when using `credential`.
 * - `iceberg.rest-catalog.oauth2.server-uri`
   - The endpoint to retrieve access token from OAuth2 Server.
+* - `iceberg.rest-catalog.oauth2.token-refresh-enabled`
+  - Controls whether a token should be refreshed if information about its expiration time is available.
+    Defaults to `true`
 * - `iceberg.rest-catalog.vended-credentials-enabled`
   - Use credentials provided by the REST backend for file system access.
     Defaults to `false`.

diff --git a/...-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogConfig.java b/...-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogConfig.java
@@ -19,10 +19,12 @@
 import io.airlift.units.Duration;
 import io.airlift.units.MinDuration;
 import jakarta.validation.constraints.NotNull;
+import org.apache.iceberg.CatalogProperties;
 
 import java.net.URI;
 import java.util.Optional;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.MINUTES;
 
 @DefunctConfig("iceberg.rest-catalog.parent-namespace")
@@ -46,6 +48,7 @@ public enum SessionType
     private boolean nestedNamespaceEnabled;
     private Security security = Security.NONE;
     private SessionType sessionType = SessionType.NONE;
+    private Duration sessionTimeout = new Duration(CatalogProperties.AUTH_SESSION_TIMEOUT_MS_DEFAULT, MILLISECONDS);
     private boolean vendedCredentialsEnabled;
     private boolean viewEndpointsEnabled = true;
     private boolean sigV4Enabled;
@@ -135,6 +138,21 @@ public IcebergRestCatalogConfig setSessionType(SessionType sessionType)
         return this;
     }
 
+    @NotNull
+    @MinDuration("0ms")
+    public Duration getSessionTimeout()
+    {
+        return sessionTimeout;
+    }
+
+    @Config("iceberg.rest-catalog.session-timeout")
+    @ConfigDescription("Duration to keep authentication session in cache")
+    public IcebergRestCatalogConfig setSessionTimeout(Duration sessionTimeout)
+    {
+        this.sessionTimeout = sessionTimeout;
+        return this;
+    }
+
     public boolean isVendedCredentialsEnabled()
     {
         return vendedCredentialsEnabled;

diff --git a/...rino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityConfig.java b/...rino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityConfig.java
@@ -17,6 +17,7 @@
 import io.airlift.configuration.ConfigDescription;
 import io.airlift.configuration.ConfigSecuritySensitive;
 import jakarta.validation.constraints.AssertTrue;
+import org.apache.iceberg.rest.auth.OAuth2Properties;
 
 import java.net.URI;
 import java.util.Optional;
@@ -27,6 +28,7 @@ public class OAuth2SecurityConfig
     private String scope;
     private String token;
     private URI serverUri;
+    private boolean tokenRefreshEnabled = OAuth2Properties.TOKEN_REFRESH_ENABLED_DEFAULT;
 
     public Optional<String> getCredential()
     {
@@ -82,6 +84,19 @@ public OAuth2SecurityConfig setServerUri(URI serverUri)
         return this;
     }
 
+    public boolean isTokenRefreshEnabled()
+    {
+        return tokenRefreshEnabled;
+    }
+
+    @Config("iceberg.rest-catalog.oauth2.token-refresh-enabled")
+    @ConfigDescription("Controls whether a token should be refreshed if information about its expiration time is available")
+    public OAuth2SecurityConfig setTokenRefreshEnabled(boolean tokenRefreshEnabled)
+    {
+        this.tokenRefreshEnabled = tokenRefreshEnabled;
+        return this;
+    }
+
     @AssertTrue(message = "OAuth2 requires a credential or token")
     public boolean credentialOrTokenPresent()
     {

diff --git a/...-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityProperties.java b/...-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/OAuth2SecurityProperties.java
@@ -42,6 +42,7 @@ public OAuth2SecurityProperties(OAuth2SecurityConfig securityConfig)
                 value -> propertiesBuilder.put(OAuth2Properties.TOKEN, value));
         securityConfig.getServerUri().ifPresent(
                 value -> propertiesBuilder.put(OAuth2Properties.OAUTH2_SERVER_URI, value.toString()));
+        propertiesBuilder.put(OAuth2Properties.TOKEN_REFRESH_ENABLED, String.valueOf(securityConfig.isTokenRefreshEnabled()));
 
         this.securityProperties = propertiesBuilder.buildOrThrow();
     }

diff --git a/...rg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoIcebergRestCatalogFactory.java b/...rg/src/main/java/io/trino/plugin/iceberg/catalog/rest/TrinoIcebergRestCatalogFactory.java
@@ -18,6 +18,7 @@
 import com.google.common.collect.Maps;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import com.google.inject.Inject;
+import io.airlift.units.Duration;
 import io.trino.cache.EvictableCacheBuilder;
 import io.trino.plugin.hive.NodeVersion;
 import io.trino.plugin.iceberg.IcebergConfig;
@@ -42,6 +43,7 @@
 
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static org.apache.iceberg.CatalogProperties.AUTH_SESSION_TIMEOUT_MS;
 import static org.apache.iceberg.rest.auth.OAuth2Properties.CREDENTIAL;
 import static org.apache.iceberg.rest.auth.OAuth2Properties.TOKEN;
 
@@ -56,6 +58,7 @@ public class TrinoIcebergRestCatalogFactory
     private final Optional<String> warehouse;
     private final boolean nestedNamespaceEnabled;
     private final SessionType sessionType;
+    private final Duration sessionTimeout;
     private final boolean vendedCredentialsEnabled;
     private final boolean viewEndpointsEnabled;
     private final SecurityProperties securityProperties;
@@ -89,6 +92,7 @@ public TrinoIcebergRestCatalogFactory(
         this.warehouse = restConfig.getWarehouse();
         this.nestedNamespaceEnabled = restConfig.isNestedNamespaceEnabled();
         this.sessionType = restConfig.getSessionType();
+        this.sessionTimeout = restConfig.getSessionTimeout();
         this.vendedCredentialsEnabled = restConfig.isVendedCredentialsEnabled();
         this.viewEndpointsEnabled = restConfig.isViewEndpointsEnabled();
         this.securityProperties = requireNonNull(securityProperties, "securityProperties is null");
@@ -119,6 +123,7 @@ public synchronized TrinoCatalog create(ConnectorIdentity identity)
             prefix.ifPresent(prefix -> properties.put("prefix", prefix));
             properties.put("view-endpoints-supported", Boolean.toString(viewEndpointsEnabled));
             properties.put("trino-version", trinoVersion);
+            properties.put(AUTH_SESSION_TIMEOUT_MS, String.valueOf(sessionTimeout.toMillis()));
             properties.putAll(securityProperties.get());
             properties.putAll(awsProperties.get());
 

diff --git a/...berg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogConfig.java b/...berg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestIcebergRestCatalogConfig.java
@@ -15,13 +15,15 @@
 
 import com.google.common.collect.ImmutableMap;
 import io.airlift.units.Duration;
+import org.apache.iceberg.CatalogProperties;
 import org.junit.jupiter.api.Test;
 
 import java.util.Map;
 
 import static io.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
 import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
 import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.MINUTES;
 
 public class TestIcebergRestCatalogConfig
@@ -35,6 +37,7 @@ public void testDefaults()
                 .setWarehouse(null)
                 .setNestedNamespaceEnabled(false)
                 .setSessionType(IcebergRestCatalogConfig.SessionType.NONE)
+                .setSessionTimeout(new Duration(CatalogProperties.AUTH_SESSION_TIMEOUT_MS_DEFAULT, MILLISECONDS))
                 .setSecurity(IcebergRestCatalogConfig.Security.NONE)
                 .setVendedCredentialsEnabled(false)
                 .setViewEndpointsEnabled(true)
@@ -53,6 +56,7 @@ public void testExplicitPropertyMappings()
                 .put("iceberg.rest-catalog.nested-namespace-enabled", "true")
                 .put("iceberg.rest-catalog.security", "OAUTH2")
                 .put("iceberg.rest-catalog.session", "USER")
+                .put("iceberg.rest-catalog.session-timeout", "100ms")
                 .put("iceberg.rest-catalog.vended-credentials-enabled", "true")
                 .put("iceberg.rest-catalog.view-endpoints-enabled", "false")
                 .put("iceberg.rest-catalog.sigv4-enabled", "true")
@@ -66,6 +70,7 @@ public void testExplicitPropertyMappings()
                 .setWarehouse("test_warehouse_identifier")
                 .setNestedNamespaceEnabled(true)
                 .setSessionType(IcebergRestCatalogConfig.SessionType.USER)
+                .setSessionTimeout(new Duration(100, MILLISECONDS))
                 .setSecurity(IcebergRestCatalogConfig.Security.OAUTH2)
                 .setVendedCredentialsEnabled(true)
                 .setViewEndpointsEnabled(false)

diff --git a/...-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestOAuth2SecurityConfig.java b/...-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/rest/TestOAuth2SecurityConfig.java
@@ -14,6 +14,7 @@
 package io.trino.plugin.iceberg.catalog.rest;
 
 import com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.rest.auth.OAuth2Properties;
 import org.junit.jupiter.api.Test;
 
 import java.net.URI;
@@ -33,7 +34,8 @@ public void testDefaults()
                 .setCredential(null)
                 .setToken(null)
                 .setScope(null)
-                .setServerUri(null));
+                .setServerUri(null)
+                .setTokenRefreshEnabled(OAuth2Properties.TOKEN_REFRESH_ENABLED_DEFAULT));
     }
 
     @Test
@@ -44,13 +46,15 @@ public void testExplicitPropertyMappings()
                 .put("iceberg.rest-catalog.oauth2.credential", "credential")
                 .put("iceberg.rest-catalog.oauth2.scope", "scope")
                 .put("iceberg.rest-catalog.oauth2.server-uri", "http://localhost:8080/realms/iceberg/protocol/openid-connect/token")
+                .put("iceberg.rest-catalog.oauth2.token-refresh-enabled", "false")
                 .buildOrThrow();
 
         OAuth2SecurityConfig expected = new OAuth2SecurityConfig()
                 .setCredential("credential")
                 .setToken("token")
                 .setScope("scope")
-                .setServerUri(URI.create("http://localhost:8080/realms/iceberg/protocol/openid-connect/token"));
+                .setServerUri(URI.create("http://localhost:8080/realms/iceberg/protocol/openid-connect/token"))
+                .setTokenRefreshEnabled(false);
         assertThat(expected.credentialOrTokenPresent()).isTrue();
         assertThat(expected.scopePresentOnlyWithCredential()).isFalse();
         assertFullMapping(properties, expected);