This is an automated email from the ASF dual-hosted git repository.
diqiu50 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new 16db87bf6f Feature/trino connector oauth2 user forwarding (#12045)
16db87bf6f is described below
commit 16db87bf6f4c8c7c0b3aa3791043e821d03aa64f
Author: Mark Hoerth <[email protected]>
AuthorDate: Mon Aug 17 00:31:12 2026 -0700
Feature/trino connector oauth2 user forwarding (#12045)
### Summary
Adds per-user authorization and credential vending support to the
Gravitino Trino connector for Iceberg catalogs backed by the Iceberg
REST Catalog (IRC).
Changes
Two commits:
* oauth2 per-user token forwarding — Extends the connector's session
forwarding (forwardUser=true) to support authType=oauth2, not just
simple. When enabled, the end user's forwarded IdP token is presented to
the Gravitino server, so authorization runs against the real end user
instead of a shared service identity. Adds StaticUserTokenProvider,
widens the auth guards in GravitinoConnector and GravitinoAuthProvider,
and keys the per-user client cache by auth type.
* per-user credential vending for IRC-backed Iceberg — Configures the
generated Trino Iceberg REST catalog so that an IRC-backed Gravitino
catalog requests per-user, per-table vended credentials. Sets
security=OAUTH2, session=USER, and vended-credentials-enabled=true, and
maps the warehouse through, so Trino's Iceberg REST client forwards the
end-user token to the IRC and obtains scoped storage credentials over
the REST protocol.
### Testing
Unit tests added for the oauth2 forwarding path
(TestGravitinoAuthProvider) and the REST backend config mapping
(TestIcebergCatalogPropertyConverter). End-to-end verification against a
live IRC is in progress.
### Motivation
Trino is a primary query engine for enterprise customers who require
real per-user authorization and scoped credentials rather than a shared
service identity.
---------
Co-authored-by: Mark Hoerth <[email protected]>
Co-authored-by: yuhui <[email protected]>
Co-authored-by: Claude Opus 5 (1M context) <[email protected]>
---
docs/trino-connector/authentication.md | 211 +++++++++++++++++----
.../service/IcebergCatalogWrapperManager.java | 4 +
.../authentication/AuthenticationFilter.java | 11 ++
.../trino/connector/GravitinoConnector.java | 79 ++++++--
.../iceberg/IcebergCatalogPropertyConverter.java | 12 +-
.../connector/security/GravitinoAuthProvider.java | 58 +++++-
.../security/StaticUserTokenProvider.java | 54 ++++++
.../TestGravitinoConnectorForwardUser.java | 171 ++++++++++++++++-
.../TestIcebergCatalogPropertyConverter.java | 27 +++
.../security/TestGravitinoAuthProvider.java | 59 ++++++
10 files changed, 620 insertions(+), 66 deletions(-)
diff --git a/docs/trino-connector/authentication.md
b/docs/trino-connector/authentication.md
index 5ec6de6018..9793e7749f 100644
--- a/docs/trino-connector/authentication.md
+++ b/docs/trino-connector/authentication.md
@@ -11,6 +11,8 @@ The Gravitino Trino connector supports authenticating to the
Gravitino server us
If `gravitino.client.authType` is not set, the connector operates in
no-authentication mode and connects to the Gravitino server without any
credentials.
+## Authentication Types
+
### Simple Authentication
Simple authentication uses a username to authenticate with the Gravitino
server.
@@ -90,32 +92,6 @@ gravitino.client.oauth2.scope=gravitino
| `gravitino.client.oauth2.path` | OAuth2 token endpoint path
| (none) | Yes if authType is `oauth2` |
| `gravitino.client.oauth2.scope` | OAuth2 scope
| (none) | Yes if authType is `oauth2` |
-### Kerberos Authentication
-
-Kerberos authentication uses Kerberos tickets to authenticate with the
Gravitino server.
-
-**Configuration in `etc/catalog/gravitino.properties`:**
-
-```properties
-connector.name=gravitino
-gravitino.metalake=metalake
-gravitino.uri=http://localhost:8090
-
-# Kerberos authentication with keytab
-gravitino.client.authType=kerberos
-gravitino.client.kerberos.principal=user@REALM
-gravitino.client.kerberos.keytabFilePath=/path/to/user.keytab
-```
-
-**Configuration properties:**
-
-| Property | Description
| Default value | Required
|
-|--------------------------------------------|-----------------------------------------------------------------|---------------|-----------------------------------------|
-| `gravitino.client.authType` | Authentication type: `simple`,
`basic`, `oauth2`, or `kerberos` | (none) | Yes (to enable Kerberos)
|
-| `gravitino.client.kerberos.principal` | Kerberos principal
| (none) | Yes if authType is `kerberos`
|
-| `gravitino.client.kerberos.keytabFilePath` | Path to keytab file
| (none) | No (uses ticket cache if not
specified) |
-
-
### Example: Connecting to OAuth-Protected Gravitino Server
This example shows how to configure the Trino connector to connect to a
Gravitino server protected by OAuth authentication.
@@ -151,11 +127,36 @@ gravitino.client.oauth2.scope=test
SHOW CATALOGS;
```
-### Session Credential Forwarding
+### Kerberos Authentication
+
+Kerberos authentication uses Kerberos tickets to authenticate with the
Gravitino server.
+
+**Configuration in `etc/catalog/gravitino.properties`:**
+
+```properties
+connector.name=gravitino
+gravitino.metalake=metalake
+gravitino.uri=http://localhost:8090
+
+# Kerberos authentication with keytab
+gravitino.client.authType=kerberos
+gravitino.client.kerberos.principal=user@REALM
+gravitino.client.kerberos.keytabFilePath=/path/to/user.keytab
+```
+
+**Configuration properties:**
+
+| Property | Description
| Default value | Required
| Since version |
+|----------------------------------------------|---------------------------------------------------------------------|-----------------|-------------------------------------------|-----------------|
+| `gravitino.client.authType` | Authentication type:
`simple`, `basic`, `oauth2`, or `kerberos` | (none) | Yes (to
enable Kerberos) | 1.3.0 |
+| `gravitino.client.kerberos.principal` | Kerberos principal
| (none) | Yes if authType is
`kerberos` | 1.3.0 |
+| `gravitino.client.kerberos.keytabFilePath` | Path to keytab file
| (none) | No (uses ticket cache
if not specified) | 1.3.0 |
-Setting `gravitino.client.session.forwardUser=true` with `authType=simple`
creates a dedicated Gravitino client per Trino session user, so each user is
visible in the Gravitino audit log instead of the shared `gravitino.user`.
+## Session Credential Forwarding
-**Configuration:**
+Setting `gravitino.client.session.forwardUser=true` creates a dedicated
Gravitino client per Trino session user, so each user is visible in the
Gravitino audit log instead of the shared `gravitino.user` or service identity.
It is supported with `authType=simple` and `authType=oauth2`.
+
+**Configuration (`authType=simple`):**
```properties
connector.name=gravitino
@@ -166,22 +167,160 @@ gravitino.client.authType=simple
gravitino.client.session.forwardUser=true
```
+With `authType=simple`, the Trino session username is forwarded to Gravitino
as the simple-auth identity.
+
+**Configuration (`authType=oauth2`):**
+
+```properties
+connector.name=gravitino
+gravitino.metalake=metalake
+gravitino.uri=http://localhost:8090
+
+gravitino.client.authType=oauth2
+gravitino.client.oauth2.serverUri=http://oauth-server:8080
+gravitino.client.oauth2.credential=client_id:client_secret
+gravitino.client.oauth2.path=oauth2/token
+gravitino.client.oauth2.scope=gravitino
+gravitino.client.session.forwardUser=true
+```
+
+With `authType=oauth2`, the end user's IdP access token is presented to
Gravitino directly instead of the shared client-credentials identity. This
requires the Trino coordinator to populate the session's extra-credentials with
the caller's access token under the key `token`; the connector reads it from
there, and `buildForSession` fails with a clear error if it's missing.
+
+Whether the coordinator can populate this extra-credential depends on the
Trino distribution:
+
+- **Starburst Enterprise** supports this via
`http-server.authentication.type=DELEGATED-OAUTH2` — see [OAuth 2.0 token
pass-through](https://docs.starburst.io/latest/security/oauth2-passthrough.html).
+- **Open-source Trino does not support this yet.** There is no equivalent
coordinator-side mechanism to forward the caller's OAuth2 token into the
connector session; see [trinodb/trino discussion
#24403](https://github.com/trinodb/trino/discussions/24403) and [issue
#27917](https://github.com/trinodb/trino/issues/27917) tracking this feature
request upstream.
+
+The `gravitino.client.oauth2.*` properties above still configure the shared
bootstrap/admin client used for catalog discovery — they are unrelated to the
per-user forwarded token.
+
+For an Iceberg catalog with `catalog-backend=rest` (backed by an Iceberg REST
Catalog), the connector does not set
`iceberg.rest-catalog.security`/`iceberg.rest-catalog.session` on its own —
that catalog's own `gravitino.client.*` config is unrelated to how its
underlying Iceberg REST catalog authenticates. To also forward the end user's
token to the REST catalog itself, set
`trino.bypass.iceberg.rest-catalog.security=OAUTH2` and
`trino.bypass.iceberg.rest-catalog.session=USER` explicitl [...]
+
**Configuration properties:**
-| Property | Description
|
Default value | Required |
-|-----------------------------------------------------------|--------------------------------------------------------------------------------------------|---------------|----------|
-| `gravitino.client.session.forwardUser` | When `true` with
`authType=simple`, forwards the Trino session user to Gravitino per-query |
`false` | No |
-| `gravitino.client.session.cache.maxSize` | Maximum number
of per-user sessions to keep in the cache |
`500` | No |
-| `gravitino.client.session.cache.expireAfterAccessSeconds` | Seconds before
an idle per-user session is evicted from the cache |
`3600` | No |
+| Property | Description
| Default value | Required | Since version |
+|--------------------------------------------------------------|--------------------------------------------------------------------------------------------------|-----------------|------------|-----------------|
+| `gravitino.client.session.forwardUser` | When `true`
with `authType=simple` or `authType=oauth2`, forwards the Trino session
user/token to Gravitino per-query | `false` | No | 1.3.0
|
+| `gravitino.client.session.cache.maxSize` | Maximum
number of per-user sessions to keep in the cache
| `500` | No | 1.3.0 |
+| `gravitino.client.session.cache.expireAfterAccessSeconds` | Seconds
before an idle per-user session is evicted from the cache
| `3600` | No | 1.3.0 |
+
+### Example: OAuth2 Per-User Token Forwarding
+
+This example walks through a full setup where each Trino user's own OAuth2
access token is
+forwarded to Gravitino and to an Iceberg REST catalog (IRC), instead of a
single shared service
+identity.
+
+**1. Trino coordinator: forward the logged-in user's token to connectors.**
This is the
+prerequisite that makes `authType=oauth2` forwarding possible at all — the
coordinator must
+populate the session's extra-credentials with the caller's access token under
the key `token`.
+
+- **Starburst Enterprise**: set the following in `etc/config.properties`:
+
+ ```properties
+ http-server.authentication.type=DELEGATED-OAUTH2
+ ```
+
+ See [OAuth 2.0 token
pass-through](https://docs.starburst.io/latest/security/oauth2-passthrough.html)
+ for details, including its limitation that pass-through tokens are not
refreshed and must
+ outlive the query.
+
+- **Open-source Trino**: there is currently no equivalent coordinator setting.
Track
+ [trinodb/trino discussion
#24403](https://github.com/trinodb/trino/discussions/24403) and
+ [issue #27917](https://github.com/trinodb/trino/issues/27917) for this
feature request. Until
+ it lands upstream, this connector's `authType=oauth2` forwardUser path
requires a Trino
+ distribution that provides this extra-credential itself.
+
+**2. Gravitino server: enable OAuth2** (in `conf/gravitino.conf`):
+
+```properties
+gravitino.authenticators=oauth
+gravitino.authenticator.oauth.serviceAudience=account
+gravitino.authenticator.oauth.jwksUri=http://your-idp/realms/gravitino/protocol/openid-connect/certs
+gravitino.authenticator.oauth.tokenValidatorClass=org.apache.gravitino.server.authentication.JwksTokenValidator
+gravitino.authenticator.oauth.principalFields=preferred_username,email,sub
+```
+
+**3. Trino connector: enable OAuth2 forwarding** (in
`etc/catalog/gravitino.properties`):
+
+```properties
+connector.name=gravitino
+gravitino.metalake=my_metalake
+gravitino.uri=http://localhost:8090
+
+gravitino.client.authType=oauth2
+gravitino.client.oauth2.serverUri=http://your-idp
+gravitino.client.oauth2.credential=service-account-id:service-account-secret
+gravitino.client.oauth2.path=realms/gravitino/protocol/openid-connect/token
+gravitino.client.oauth2.scope=email
+gravitino.client.session.forwardUser=true
+```
+
+The `gravitino.client.oauth2.*` properties configure the shared service
identity used for catalog
+discovery; the per-user forwarded token (from step 1) is what each query
actually authenticates
+with once `forwardUser=true`.
+
+**4. Create the metalake and catalog.** Create the metalake `my_metalake`
first (via the
+Gravitino REST API, SDK, or CLI — see
+[Manage metalakes](../manage-metalake-using-gravitino.md#create-a-metalake)),
then create a
+REST-backed Iceberg catalog under it from the Trino CLI using the
+`gravitino.system.create_catalog` procedure. To also forward the end user's
token to the Iceberg
+REST catalog (IRC) itself, set
`trino.bypass.iceberg.rest-catalog.security=OAUTH2` and
+`trino.bypass.iceberg.rest-catalog.session=USER` on the catalog, alongside its
bootstrap
+`trino.bypass.iceberg.rest-catalog.oauth2.*` credentials:
+
+```sql
+call gravitino.system.create_catalog(
+ 'my_catalog',
+ 'lakehouse-iceberg',
+ map(
+ array['uri', 'catalog-backend', 'warehouse',
+ 'trino.bypass.iceberg.rest-catalog.security',
'trino.bypass.iceberg.rest-catalog.session',
+ 'trino.bypass.iceberg.rest-catalog.oauth2.credential',
'trino.bypass.iceberg.rest-catalog.oauth2.scope',
+ 'trino.bypass.iceberg.rest-catalog.oauth2.server-uri'
+ ],
+ array['http://irc-host:9001/iceberg', 'rest', 'my_catalog',
+ 'OAUTH2', 'USER',
+ 'service-account-id:service-account-secret', 'email',
+ 'http://your-idp/realms/gravitino/protocol/openid-connect/token'
+ ]
+ )
+);
+```
+
+This call itself runs with the connector's own shared service identity, not
any forwarded user
+token — `forwardUser` only affects `SELECT`/`SHOW`-style queries against the
catalog afterward,
+not catalog registration itself. `create_catalog` both creates the catalog in
Gravitino and loads
+it into Trino as its own top-level catalog — not as a schema nested under a
single `gravitino`
+catalog. If the two `trino.bypass.iceberg.rest-catalog.*` properties above are
omitted, the REST
+catalog keeps its own default security setting, independent of
+`gravitino.client.session.forwardUser`, and the end user's token never reaches
the IRC.
+
+**5. Query as a specific user.** With a real OIDC login flow, Trino populates
the forwarded token
+automatically after the user signs in — this automatic population is the part
that requires
+Starburst's DELEGATED-OAUTH2 (step 1). For manual testing on any Trino
distribution, including
+open-source Trino, the same extra-credential can instead be set directly on
the CLI, independent
+of how the coordinator is configured:
+
+```shell
+trino --server http://localhost:8080 \
+ --user alice \
+ --extra-credential token=<alice-idp-access-token> \
+ --execute "SHOW SCHEMAS IN my_catalog"
+```
+
+Gravitino sees this request as `alice`, not the shared service identity —
`alice`'s own
+privileges apply, and a request with a missing or invalid token is rejected
before it reaches the
+catalog. Because `my_catalog` was created with
`trino.bypass.iceberg.rest-catalog.session=USER`
+in step 4, the same forwarded token also reaches the IRC directly, so per-user
authorization
+applies consistently whether Trino talks to Gravitino's native API or straight
to the IRC.
-### Notes
+## Notes
- The Gravitino server must be configured with the corresponding
authentication mechanism enabled.
- For OAuth2 authentication, ensure the OAuth2 server is accessible from the
Trino coordinator and workers.
- For Kerberos authentication, ensure the Kerberos configuration is properly
set up on all Trino nodes.
- Authentication configuration is passed through the `gravitino.client.*`
prefix to the underlying Gravitino Java client.
-### See Also
+## See Also
- [Gravitino Server Authentication
Configuration](../security/how-to-authenticate.md)
- [Local users and groups](../security/local-users-and-groups.md)
diff --git
a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/IcebergCatalogWrapperManager.java
b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/IcebergCatalogWrapperManager.java
index c524d074fa..d5a2d27206 100644
---
a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/IcebergCatalogWrapperManager.java
+++
b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/IcebergCatalogWrapperManager.java
@@ -103,6 +103,10 @@ public class IcebergCatalogWrapperManager implements
AutoCloseable {
}
public CatalogWrapperForREST getCatalogWrapper(String catalogName) {
+ if (LOG.isDebugEnabled()) {
+ boolean cacheHit = catalogWrapperCache.getIfPresent(catalogName) != null;
+ LOG.debug("getCatalogWrapper catalogName={} cacheHit={}", catalogName,
cacheHit);
+ }
CatalogWrapperForREST catalogWrapperForREST =
catalogWrapperCache.get(catalogName, k ->
createCatalogWrapper(catalogName));
// Reload conf to reset UserGroupInformation or icebergTableOps will
always use
diff --git
a/server-common/src/main/java/org/apache/gravitino/server/authentication/AuthenticationFilter.java
b/server-common/src/main/java/org/apache/gravitino/server/authentication/AuthenticationFilter.java
index 8fa379922b..6f75837212 100644
---
a/server-common/src/main/java/org/apache/gravitino/server/authentication/AuthenticationFilter.java
+++
b/server-common/src/main/java/org/apache/gravitino/server/authentication/AuthenticationFilter.java
@@ -43,9 +43,13 @@ import org.apache.gravitino.exceptions.UnauthorizedException;
import org.apache.gravitino.server.web.HealthCheckPathMatcher;
import org.apache.gravitino.server.web.ObjectMapperProvider;
import org.apache.gravitino.utils.PrincipalUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
public class AuthenticationFilter implements Filter {
+ private static final Logger LOG =
LoggerFactory.getLogger(AuthenticationFilter.class);
+
private final List<Authenticator> filterAuthenticators;
/**
@@ -100,6 +104,13 @@ public class AuthenticationFilter implements Filter {
}
}
}
+ if (LOG.isDebugEnabled()) {
+ LOG.debug(
+ "uri={} hasAuthHeader={} principal={}",
+ req.getRequestURI(),
+ authData != null,
+ principal == null ? "null" : principal.getName());
+ }
if (principal == null) {
throw new UnauthorizedException("The provided credentials did not
support");
}
diff --git
a/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java
b/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java
index 01d419d1fe..e038f9671e 100644
---
a/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java
+++
b/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java
@@ -18,13 +18,16 @@
*/
package org.apache.gravitino.trino.connector;
+import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
import static io.trino.spi.StandardErrorCode.PERMISSION_DENIED;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.RemovalNotification;
+import com.google.common.util.concurrent.UncheckedExecutionException;
import io.trino.spi.TrinoException;
import io.trino.spi.connector.Connector;
import io.trino.spi.connector.ConnectorAccessControl;
@@ -225,37 +228,82 @@ public class GravitinoConnector implements Connector {
catalogConnectorContext.close();
}
- private CatalogConnectorMetadata resolveSessionMetadata(ConnectorSession
session) {
- String credKey = "simple:" + session.getUser();
+ @VisibleForTesting
+ CatalogConnectorMetadata resolveSessionMetadata(ConnectorSession session) {
+ String authType =
+ catalogConnectorContext
+ .getConfig()
+ .getClientConfig()
+ .getOrDefault(GravitinoAuthProvider.AUTH_TYPE_KEY, "simple");
+ String credentialKey =
+ catalogConnectorContext
+ .getConfig()
+ .getClientConfig()
+ .getOrDefault(
+ GravitinoAuthProvider.USER_TOKEN_CREDENTIAL_KEY,
+ GravitinoAuthProvider.DEFAULT_USER_TOKEN_CREDENTIAL_KEY);
+ String token =
session.getIdentity().getExtraCredentials().get(credentialKey);
+ String credKey = sessionCacheKey(authType, session.getUser(), token);
try {
return perUserSessionCache.get(
credKey,
() -> {
- GravitinoAdminClient userClient =
- GravitinoAuthProvider.buildForSession(
- catalogConnectorContext.getConfig(), session);
+ GravitinoAdminClient userClient = buildAuthClient(session);
GravitinoMetalake userMetalake =
userClient.loadMetalake(catalogConnectorContext.getMetalake().name());
return new UserSession(
userClient, new CatalogConnectorMetadata(userMetalake,
catalogIdentifier));
})
.metadata;
- } catch (ExecutionException e) {
- Throwable cause = e.getCause();
+ } catch (ExecutionException | UncheckedExecutionException e) {
+ Throwable cause = e.getCause() == null ? e : e.getCause();
LOG.warn(
- "Failed to create per-user Gravitino client for user '{}': {}",
- session.getUser(),
- cause.getMessage());
+ "Failed to create per-user Gravitino client for user '{}'",
session.getUser(), cause);
+ if (cause instanceof TrinoException) {
+ // Already carries a specific Trino error code (e.g. from
buildForSession); re-wrapping
+ // would swallow it.
+ throw (TrinoException) cause;
+ }
+ if (cause instanceof IllegalArgumentException
+ || cause instanceof UnsupportedOperationException) {
+ throw new TrinoException(
+ PERMISSION_DENIED,
+ "Failed to authenticate user '"
+ + session.getUser()
+ + "' with Gravitino: "
+ + cause.getMessage(),
+ cause);
+ }
throw new TrinoException(
- PERMISSION_DENIED,
- "Failed to authenticate user '"
+ GENERIC_INTERNAL_ERROR,
+ "Unexpected error while creating per-user Gravitino client for user
'"
+ session.getUser()
- + "' with Gravitino: "
+ + "': "
+ cause.getMessage(),
cause);
}
}
+ /**
+ * Builds the per-user Gravitino admin client for a forwarded session.
Extracted as an overridable
+ * seam so tests can substitute the client-building behavior without needing
to mock the static
+ * {@link GravitinoAuthProvider#buildForSession}.
+ *
+ * @param session the current Trino connector session
+ * @return the per-user Gravitino admin client
+ */
+ @VisibleForTesting
+ GravitinoAdminClient buildAuthClient(ConnectorSession session) {
+ return
GravitinoAuthProvider.buildForSession(catalogConnectorContext.getConfig(),
session);
+ }
+
+ @VisibleForTesting
+ static String sessionCacheKey(String authType, String user, String token) {
+ String tokenPart =
+ StringUtils.isBlank(token) ? "" : ":" +
Integer.toHexString(token.hashCode());
+ return authType + ":" + user + tokenPart;
+ }
+
private Cache<String, UserSession> buildSessionCache(GravitinoConfig config)
{
Map<String, String> clientConfig = config.getClientConfig();
String authTypeStr = clientConfig.get(GravitinoAuthProvider.AUTH_TYPE_KEY);
@@ -265,10 +313,11 @@ public class GravitinoConnector implements Connector {
"gravitino.client.session.forwardUser=true requires
gravitino.client.authType to be set");
}
GravitinoAuthProvider.AuthType authType =
GravitinoAuthProvider.parseAuthType(authTypeStr);
- if (authType != GravitinoAuthProvider.AuthType.SIMPLE) {
+ if (authType != GravitinoAuthProvider.AuthType.SIMPLE
+ && authType != GravitinoAuthProvider.AuthType.OAUTH2) {
throw new TrinoException(
GravitinoErrorCode.GRAVITINO_ILLEGAL_ARGUMENT,
- "gravitino.client.session.forwardUser=true only supports
authType=simple, got: "
+ "gravitino.client.session.forwardUser=true only supports
authType=simple or oauth2, got: "
+ authTypeStr);
}
diff --git
a/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/iceberg/IcebergCatalogPropertyConverter.java
b/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/iceberg/IcebergCatalogPropertyConverter.java
index 1802c6c64c..d6459b5988 100644
---
a/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/iceberg/IcebergCatalogPropertyConverter.java
+++
b/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/iceberg/IcebergCatalogPropertyConverter.java
@@ -161,9 +161,13 @@ public class IcebergCatalogPropertyConverter extends
CatalogPropertyConverter {
"Missing required property for Rest backend: " + missingProperty);
}
- Map<String, String> jdbcProperties = new HashMap<>();
- jdbcProperties.put("iceberg.catalog.type", "rest");
- jdbcProperties.put("iceberg.rest-catalog.uri",
properties.get(IcebergConstants.URI));
- return jdbcProperties;
+ Map<String, String> restProperties = new HashMap<>();
+ restProperties.put("iceberg.catalog.type", "rest");
+ restProperties.put("iceberg.rest-catalog.uri",
properties.get(IcebergConstants.URI));
+ if (properties.containsKey(IcebergConstants.WAREHOUSE)) {
+ restProperties.put(
+ "iceberg.rest-catalog.warehouse",
properties.get(IcebergConstants.WAREHOUSE));
+ }
+ return restProperties;
}
}
diff --git
a/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/GravitinoAuthProvider.java
b/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/GravitinoAuthProvider.java
index 21e2e25604..2f68a9a7b4 100644
---
a/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/GravitinoAuthProvider.java
+++
b/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/GravitinoAuthProvider.java
@@ -19,6 +19,7 @@
package org.apache.gravitino.trino.connector.security;
import com.google.common.base.Preconditions;
+import io.trino.spi.TrinoException;
import io.trino.spi.connector.ConnectorSession;
import java.io.File;
import java.util.Locale;
@@ -29,6 +30,7 @@ import org.apache.gravitino.client.GravitinoAdminClient;
import org.apache.gravitino.client.GravitinoClientConfiguration;
import org.apache.gravitino.client.KerberosTokenProvider;
import org.apache.gravitino.trino.connector.GravitinoConfig;
+import org.apache.gravitino.trino.connector.GravitinoErrorCode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -55,6 +57,18 @@ public class GravitinoAuthProvider {
public static final String FORWARD_SESSION_USER_KEY =
GravitinoClientConfiguration.GRAVITINO_CLIENT_CONFIG_PREFIX +
"session.forwardUser";
+ /**
+ * Configuration key naming the Trino session extra-credential that carries
the end user's
+ * forwarded IdP access token. Defaults to {@value
#DEFAULT_USER_TOKEN_CREDENTIAL_KEY}. Query
+ * engines differ in how they label the forwarded token, so this is
configurable rather than
+ * fixed.
+ */
+ public static final String USER_TOKEN_CREDENTIAL_KEY =
+ "gravitino.client.session.userTokenCredentialKey";
+
+ /** Default extra-credential name used when {@link
#USER_TOKEN_CREDENTIAL_KEY} is not set. */
+ public static final String DEFAULT_USER_TOKEN_CREDENTIAL_KEY = "token";
+
/** Built-in IdP username configuration key for Basic authentication. */
public static final String BASIC_USERNAME_KEY =
GravitinoClientConfiguration.GRAVITINO_CLIENT_CONFIG_PREFIX +
"basic.username";
@@ -165,8 +179,10 @@ public class GravitinoAuthProvider {
* connector session. This is the entry point for the per-user client cache
when {@code
* forwardUser=true}.
*
- * <p>Currently only {@code authType=simple} is supported: the Trino session
username is used as
- * the Gravitino simple-auth identity.
+ * <p>For {@code authType=simple} the Trino session username is used as the
Gravitino simple-auth
+ * identity. For {@code authType=oauth2} the end-user IdP token forwarded
into the connector
+ * session (extra-credential key {@code token}) is presented to Gravitino as
the bearer token, so
+ * the server authorizes against the end user rather than a shared service
identity.
*
* @param config the Gravitino connector configuration
* @param session the current Trino connector session
@@ -196,13 +212,38 @@ public class GravitinoAuthProvider {
GravitinoAdminClient.AdminClientBuilder builder =
GravitinoAdminClient.builder(uri);
- if (authType != AuthType.SIMPLE) {
- throw new UnsupportedOperationException(
- "Auth type "
- + authType
- + " does not support session forwarding. Only simple is
supported.");
+ switch (authType) {
+ case SIMPLE:
+ builder.withSimpleAuth(session.getUser());
+ break;
+ case OAUTH2:
+ {
+ String credentialKey =
+ clientConfig.getOrDefault(
+ USER_TOKEN_CREDENTIAL_KEY,
DEFAULT_USER_TOKEN_CREDENTIAL_KEY);
+ String userToken =
session.getIdentity().getExtraCredentials().get(credentialKey);
+ if (StringUtils.isBlank(userToken)) {
+ throw new TrinoException(
+ GravitinoErrorCode.GRAVITINO_ILLEGAL_ARGUMENT,
+ "No forwarded user token found in session extra-credentials
under key '"
+ + credentialKey
+ + "'. The Trino coordinator must populate this
extra-credential with the "
+ + "caller's OAuth2 access token, e.g. Starburst
Enterprise's "
+ + "http-server.authentication.type=DELEGATED-OAUTH2;
open-source Trino does not "
+ + "support this yet. If the coordinator labels the
forwarded token differently, "
+ + "set "
+ + USER_TOKEN_CREDENTIAL_KEY
+ + " to match.");
+ }
+ builder.withOAuth(new StaticUserTokenProvider(userToken));
+ }
+ break;
+ default:
+ throw new UnsupportedOperationException(
+ "Auth type "
+ + authType
+ + " does not support session forwarding. Only simple and
oauth2 are supported.");
}
- builder.withSimpleAuth(session.getUser());
removeAuthSpecificKeys(clientConfig);
builder.withClientConfig(clientConfig);
@@ -233,6 +274,7 @@ public class GravitinoAuthProvider {
clientConfig.remove(KERBEROS_PRINCIPAL_KEY);
clientConfig.remove(KERBEROS_KEYTAB_FILE_PATH_KEY);
clientConfig.remove(FORWARD_SESSION_USER_KEY);
+ clientConfig.remove(USER_TOKEN_CREDENTIAL_KEY);
clientConfig.remove(SESSION_CACHE_MAX_SIZE_KEY);
clientConfig.remove(SESSION_CACHE_EXPIRE_AFTER_ACCESS_SECONDS_KEY);
}
diff --git
a/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/StaticUserTokenProvider.java
b/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/StaticUserTokenProvider.java
new file mode 100644
index 0000000000..f477a633b3
--- /dev/null
+++
b/trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/StaticUserTokenProvider.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.gravitino.trino.connector.security;
+
+import com.google.common.base.Preconditions;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.gravitino.client.OAuth2TokenProvider;
+
+/**
+ * An {@link OAuth2TokenProvider} that returns a pre-fetched, already-valid
access token rather than
+ * minting one via client credentials. Used for per-user session forwarding:
the end user's IdP
+ * access token, forwarded by Trino into the connector session, is presented
directly to Gravitino
+ * so the server authorizes against the end user's identity instead of a
shared service identity.
+ *
+ * <p>The raw token is returned from {@link #getAccessToken()}. The {@code
Bearer } prefix is added
+ * by {@link OAuth2TokenProvider#getTokenData()}, so the token held here must
not include it.
+ */
+public final class StaticUserTokenProvider extends OAuth2TokenProvider {
+
+ private final String accessToken;
+
+ /**
+ * Constructs a provider that always returns the given access token.
+ *
+ * @param accessToken the raw bearer token (without the {@code Bearer }
prefix) to present to
+ * Gravitino
+ */
+ public StaticUserTokenProvider(String accessToken) {
+ Preconditions.checkArgument(
+ StringUtils.isNotBlank(accessToken), "accessToken must not be blank");
+ this.accessToken = accessToken;
+ }
+
+ @Override
+ protected String getAccessToken() {
+ return accessToken;
+ }
+}
diff --git
a/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/TestGravitinoConnectorForwardUser.java
b/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/TestGravitinoConnectorForwardUser.java
index 4380942be9..29b91f8ca3 100644
---
a/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/TestGravitinoConnectorForwardUser.java
+++
b/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/TestGravitinoConnectorForwardUser.java
@@ -18,8 +18,13 @@
*/
package org.apache.gravitino.trino.connector;
+import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
+import static io.trino.spi.StandardErrorCode.PERMISSION_DENIED;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
@@ -27,12 +32,18 @@ import static org.mockito.Mockito.when;
import com.google.common.collect.ImmutableMap;
import io.trino.spi.TrinoException;
+import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.security.ConnectorIdentity;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
import org.apache.gravitino.Catalog;
import org.apache.gravitino.NameIdentifier;
import org.apache.gravitino.SupportsSchemas;
+import org.apache.gravitino.client.GravitinoAdminClient;
import org.apache.gravitino.client.GravitinoMetalake;
import org.apache.gravitino.rel.TableCatalog;
import org.apache.gravitino.trino.connector.catalog.CatalogConnectorContext;
+import org.apache.gravitino.trino.connector.catalog.CatalogConnectorMetadata;
import org.apache.gravitino.trino.connector.metadata.GravitinoCatalog;
import org.apache.gravitino.trino.connector.security.GravitinoAuthProvider;
import org.junit.jupiter.api.Test;
@@ -63,15 +74,14 @@ class TestGravitinoConnectorForwardUser {
}
@Test
- void testForwardUserWithOAuth2AuthTypeThrowsAtConstruction() {
+ void testForwardUserWithOAuth2AuthTypeSucceeds() {
CatalogConnectorContext ctx =
mockContextWithConfig(
ImmutableMap.of(
GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true",
GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2"));
- TrinoException ex = assertThrows(TrinoException.class, () -> new
GravitinoConnector(ctx));
- assertEquals(GravitinoErrorCode.GRAVITINO_ILLEGAL_ARGUMENT.toErrorCode(),
ex.getErrorCode());
+ assertDoesNotThrow(() -> new GravitinoConnector(ctx));
}
@Test
@@ -91,6 +101,161 @@ class TestGravitinoConnectorForwardUser {
assertDoesNotThrow(() -> new GravitinoConnector(ctx));
}
+ @Test
+ void testSessionCacheKeyIsolatesDifferentUsers() {
+ assertNotEquals(
+ GravitinoConnector.sessionCacheKey("oauth2", "alice", "tok"),
+ GravitinoConnector.sessionCacheKey("oauth2", "bob", "tok"));
+ }
+
+ @Test
+ void testSessionCacheKeyIsolatesDifferentAuthTypes() {
+ assertNotEquals(
+ GravitinoConnector.sessionCacheKey("oauth2", "alice", "tok"),
+ GravitinoConnector.sessionCacheKey("simple", "alice", "tok"));
+ }
+
+ @Test
+ void testSessionCacheKeyIsolatesDifferentTokens() {
+ assertNotEquals(
+ GravitinoConnector.sessionCacheKey("oauth2", "alice", "token-a"),
+ GravitinoConnector.sessionCacheKey("oauth2", "alice", "token-b"));
+ }
+
+ @Test
+ void testSessionCacheKeyIsStableForSameUserAuthTypeAndToken() {
+ assertEquals(
+ GravitinoConnector.sessionCacheKey("oauth2", "alice", "tok"),
+ GravitinoConnector.sessionCacheKey("oauth2", "alice", "tok"));
+ }
+
+ @Test
+ void testSessionCacheKeyIgnoresBlankTokenForSimpleAuth() {
+ assertEquals(
+ GravitinoConnector.sessionCacheKey("simple", "alice", null),
+ GravitinoConnector.sessionCacheKey("simple", "alice", ""));
+ }
+
+ @Test
+ void testResolveSessionMetadataBuildsNewClientOnTokenRotation() {
+ CatalogConnectorContext ctx =
+ mockContextWithConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true",
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2"));
+ AtomicInteger buildCount = new AtomicInteger();
+ GravitinoConnector connector =
+ newConnectorWithAuthClient(ctx, session ->
mockAdminClient(ctx.getMetalake(), buildCount));
+
+ CatalogConnectorMetadata first =
+ connector.resolveSessionMetadata(mockSession("alice", "token-a"));
+ CatalogConnectorMetadata second =
+ connector.resolveSessionMetadata(mockSession("alice", "token-a"));
+ CatalogConnectorMetadata third =
+ connector.resolveSessionMetadata(mockSession("alice", "token-b"));
+
+ assertSame(first, second, "same user/token should reuse the cached
client");
+ assertNotSame(first, third, "a rotated token must not reuse the stale
cached client");
+ assertEquals(2, buildCount.get(), "a rotated token must trigger a fresh
client build");
+ }
+
+ @Test
+ void testResolveSessionMetadataMapsAuthSpecificFailureToPermissionDenied() {
+ CatalogConnectorContext ctx =
+ mockContextWithConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true",
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2"));
+ GravitinoConnector connector =
+ newConnectorWithAuthClient(
+ ctx,
+ session -> {
+ throw new IllegalArgumentException("No forwarded user token
found");
+ });
+
+ TrinoException ex =
+ assertThrows(
+ TrinoException.class,
+ () -> connector.resolveSessionMetadata(mockSession("alice",
"token-a")));
+ assertEquals(PERMISSION_DENIED.toErrorCode(), ex.getErrorCode());
+ }
+
+ @Test
+ void testResolveSessionMetadataPreservesTrinoExceptionErrorCode() {
+ CatalogConnectorContext ctx =
+ mockContextWithConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true",
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2"));
+ GravitinoConnector connector =
+ newConnectorWithAuthClient(
+ ctx,
+ session -> {
+ throw new TrinoException(
+ GravitinoErrorCode.GRAVITINO_ILLEGAL_ARGUMENT, "No forwarded
user token found");
+ });
+
+ TrinoException ex =
+ assertThrows(
+ TrinoException.class,
+ () -> connector.resolveSessionMetadata(mockSession("alice",
"token-a")));
+ assertEquals(GravitinoErrorCode.GRAVITINO_ILLEGAL_ARGUMENT.toErrorCode(),
ex.getErrorCode());
+ }
+
+ @Test
+ void testResolveSessionMetadataMapsUnexpectedFailureToGenericInternalError()
{
+ CatalogConnectorContext ctx =
+ mockContextWithConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true",
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2"));
+ GravitinoConnector connector =
+ newConnectorWithAuthClient(
+ ctx,
+ session -> {
+ throw new RuntimeException("connection refused");
+ });
+
+ TrinoException ex =
+ assertThrows(
+ TrinoException.class,
+ () -> connector.resolveSessionMetadata(mockSession("alice",
"token-a")));
+ assertEquals(GENERIC_INTERNAL_ERROR.toErrorCode(), ex.getErrorCode());
+ }
+
+ /**
+ * Creates a {@link GravitinoConnector} whose {@code buildAuthClient} is
overridden with the given
+ * function, so tests can control what building the per-user client does
without mocking the
+ * static {@link GravitinoAuthProvider#buildForSession}.
+ */
+ private static GravitinoConnector newConnectorWithAuthClient(
+ CatalogConnectorContext ctx,
+ Function<ConnectorSession, GravitinoAdminClient> authClientBuilder) {
+ return new GravitinoConnector(ctx) {
+ @Override
+ GravitinoAdminClient buildAuthClient(ConnectorSession session) {
+ return authClientBuilder.apply(session);
+ }
+ };
+ }
+
+ private static GravitinoAdminClient mockAdminClient(
+ GravitinoMetalake metalake, AtomicInteger buildCount) {
+ buildCount.incrementAndGet();
+ GravitinoAdminClient client = mock(GravitinoAdminClient.class);
+ when(client.loadMetalake(any())).thenReturn(metalake);
+ return client;
+ }
+
+ private static ConnectorSession mockSession(String user, String token) {
+ ConnectorIdentity identity = mock(ConnectorIdentity.class);
+ when(identity.getExtraCredentials()).thenReturn(ImmutableMap.of("token",
token));
+ ConnectorSession session = mock(ConnectorSession.class);
+ when(session.getUser()).thenReturn(user);
+ when(session.getIdentity()).thenReturn(identity);
+ return session;
+ }
+
private static CatalogConnectorContext mockContextWithConfig(
ImmutableMap<String, String> extraConfig) {
GravitinoCatalog mockCatalog = mock(GravitinoCatalog.class);
diff --git
a/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/catalog/iceberg/TestIcebergCatalogPropertyConverter.java
b/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/catalog/iceberg/TestIcebergCatalogPropertyConverter.java
index b9fda2c4b3..9ef2d72d40 100644
---
a/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/catalog/iceberg/TestIcebergCatalogPropertyConverter.java
+++
b/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/catalog/iceberg/TestIcebergCatalogPropertyConverter.java
@@ -57,6 +57,33 @@ public class TestIcebergCatalogPropertyConverter {
"Missing required property for Hive backend: [uri]");
}
+ @Test
+ public void testRestBackendProperty() {
+ PropertyConverter propertyConverter = new
IcebergCatalogPropertyConverter();
+ Map<String, String> gravitinoIcebergConfig =
+ ImmutableMap.<String, String>builder()
+ .put("uri", "http://localhost:9001/iceberg";)
+ .put("catalog-backend", "rest")
+ .put("warehouse", "gt_iceberg_rest")
+ .build();
+ Map<String, String> restBackendConfig =
+ propertyConverter.gravitinoToEngineProperties(gravitinoIcebergConfig);
+
+ Assertions.assertEquals(restBackendConfig.get("iceberg.catalog.type"),
"rest");
+ Assertions.assertEquals(
+ restBackendConfig.get("iceberg.rest-catalog.uri"),
"http://localhost:9001/iceberg";);
+ Assertions.assertEquals(
+ restBackendConfig.get("iceberg.rest-catalog.warehouse"),
"gt_iceberg_rest");
+
+ Map<String, String> wrongMap = Maps.newHashMap(gravitinoIcebergConfig);
+ wrongMap.remove("uri");
+
+ Assertions.assertThrows(
+ TrinoException.class,
+ () -> propertyConverter.gravitinoToEngineProperties(wrongMap),
+ "Missing required property for Rest backend: [uri]");
+ }
+
@Test
public void testJDBCBackendProperty() {
PropertyConverter propertyConverter = new
IcebergCatalogPropertyConverter();
diff --git
a/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/security/TestGravitinoAuthProvider.java
b/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/security/TestGravitinoAuthProvider.java
index 131eee6af2..03ede70a6a 100644
---
a/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/security/TestGravitinoAuthProvider.java
+++
b/trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/security/TestGravitinoAuthProvider.java
@@ -25,7 +25,9 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import com.google.common.collect.ImmutableMap;
+import io.trino.spi.TrinoException;
import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.security.ConnectorIdentity;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
@@ -259,6 +261,63 @@ public class TestGravitinoAuthProvider {
() -> GravitinoAuthProvider.buildForSession(config, session));
}
+ @Test
+ public void testBuildForSessionOAuth2() {
+ GravitinoConfig config =
+ buildConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2",
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true"));
+
+ ConnectorSession session = mock(ConnectorSession.class);
+ when(session.getUser()).thenReturn("alice");
+ when(session.getIdentity())
+ .thenReturn(
+ ConnectorIdentity.forUser("alice")
+ .withExtraCredentials(ImmutableMap.of("token",
"forwarded-user-jwt"))
+ .build());
+
+ GravitinoAdminClient client =
GravitinoAuthProvider.buildForSession(config, session);
+ assertNotNull(client);
+ }
+
+ @Test
+ public void testBuildForSessionOAuth2WithCustomCredentialKey() {
+ GravitinoConfig config =
+ buildConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2",
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true",
+ GravitinoAuthProvider.USER_TOKEN_CREDENTIAL_KEY,
"access_token"));
+
+ ConnectorSession session = mock(ConnectorSession.class);
+ when(session.getUser()).thenReturn("alice");
+ when(session.getIdentity())
+ .thenReturn(
+ ConnectorIdentity.forUser("alice")
+ .withExtraCredentials(ImmutableMap.of("access_token",
"alice-jwt"))
+ .build());
+
+ GravitinoAdminClient client =
GravitinoAuthProvider.buildForSession(config, session);
+ assertNotNull(client);
+ }
+
+ @Test
+ public void testBuildForSessionOAuth2ThrowsWhenTokenMissing() {
+ GravitinoConfig config =
+ buildConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2",
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true"));
+
+ ConnectorSession session = mock(ConnectorSession.class);
+ when(session.getUser()).thenReturn("alice");
+
when(session.getIdentity()).thenReturn(ConnectorIdentity.forUser("alice").build());
+
+ assertThrows(
+ TrinoException.class, () ->
GravitinoAuthProvider.buildForSession(config, session));
+ }
+
private GravitinoConfig buildConfig(ImmutableMap<String, String> authConfig)
{
ImmutableMap.Builder<String, String> builder =
ImmutableMap.<String, String>builder()
