Copilot commented on code in PR #10730:
URL: https://github.com/apache/gravitino/pull/10730#discussion_r3083643922
##########
trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java:
##########
@@ -64,6 +83,39 @@ public GravitinoConnector(CatalogConnectorContext
catalogConnectorContext) {
this.catalogConnectorContext = catalogConnectorContext;
this.connectorMetadata =
new CatalogConnectorMetadata(catalogConnectorContext.getMetalake(),
this.catalogIdentifier);
+
+ GravitinoConfig config = catalogConnectorContext.getConfig();
+ Map<String, String> clientConfig = config.getClientConfig();
+ this.forwardUser =
+ Boolean.parseBoolean(
+
clientConfig.getOrDefault(GravitinoAuthProvider.FORWARD_SESSION_USER_KEY,
"false"));
+ String authTypeStr =
clientConfig.getOrDefault(GravitinoAuthProvider.AUTH_TYPE_KEY, "none");
+ this.authType = GravitinoAuthProvider.parseAuthType(authTypeStr);
+ this.oauth2CredentialKey =
clientConfig.get(GravitinoAuthProvider.OAUTH2_TOKEN_CREDENTIAL_KEY);
+
+ if (forwardUser && authType == GravitinoAuthProvider.AuthType.OAUTH2) {
+ Preconditions.checkArgument(
+ StringUtils.isNotBlank(oauth2CredentialKey),
+ "oauth2 with forwardUser=true requires '%s' to be set",
+ GravitinoAuthProvider.OAUTH2_TOKEN_CREDENTIAL_KEY);
+ }
+
+ if (forwardUser) {
+ this.perUserSessionCache =
+ CacheBuilder.newBuilder()
+ .maximumSize(500)
+ .expireAfterAccess(1, TimeUnit.HOURS)
+ .removalListener(
+ (RemovalNotification<String, UserSession> notification) -> {
+ UserSession session = notification.getValue();
+ if (session != null) {
+ session.close();
+ }
+ })
+ .build();
Review Comment:
Evicting/expiring entries from `perUserSessionCache` closes the associated
`GravitinoAdminClient`. Trino can hold on to a `ConnectorMetadata` instance for
the lifetime of a query; if the underlying cached session is evicted while a
query is still using its metadata, this could close the HTTP client mid-query
and cause intermittent failures. Consider either (a) not closing sessions on
cache eviction (only on connector shutdown), or (b) making the per-query
metadata/client lifecycle explicit (e.g., tying it to transaction/query scope)
so clients are only closed when no longer in use.
##########
trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java:
##########
@@ -64,6 +83,39 @@ public GravitinoConnector(CatalogConnectorContext
catalogConnectorContext) {
this.catalogConnectorContext = catalogConnectorContext;
this.connectorMetadata =
new CatalogConnectorMetadata(catalogConnectorContext.getMetalake(),
this.catalogIdentifier);
+
+ GravitinoConfig config = catalogConnectorContext.getConfig();
+ Map<String, String> clientConfig = config.getClientConfig();
+ this.forwardUser =
+ Boolean.parseBoolean(
+
clientConfig.getOrDefault(GravitinoAuthProvider.FORWARD_SESSION_USER_KEY,
"false"));
+ String authTypeStr =
clientConfig.getOrDefault(GravitinoAuthProvider.AUTH_TYPE_KEY, "none");
+ this.authType = GravitinoAuthProvider.parseAuthType(authTypeStr);
+ this.oauth2CredentialKey =
clientConfig.get(GravitinoAuthProvider.OAUTH2_TOKEN_CREDENTIAL_KEY);
+
+ if (forwardUser && authType == GravitinoAuthProvider.AuthType.OAUTH2) {
+ Preconditions.checkArgument(
+ StringUtils.isNotBlank(oauth2CredentialKey),
+ "oauth2 with forwardUser=true requires '%s' to be set",
+ GravitinoAuthProvider.OAUTH2_TOKEN_CREDENTIAL_KEY);
Review Comment:
When `forwardUser` is enabled, the connector currently only validates the
OAuth2 credential-key case. If `authType` is `none` or `kerberos`, the
connector will initialize successfully but later fail at query time when
`buildForSession` is invoked. Consider validating at construction time that
`forwardUser=true` is only allowed with `authType=simple` or `authType=oauth2`
(and keep the existing OAuth2 credential-key validation), so misconfiguration
fails fast with a clear error.
```suggestion
if (forwardUser) {
Preconditions.checkArgument(
authType == GravitinoAuthProvider.AuthType.SIMPLE
|| authType == GravitinoAuthProvider.AuthType.OAUTH2,
"forwardUser=true requires authType to be 'simple' or 'oauth2',
but found '%s'",
authTypeStr);
if (authType == GravitinoAuthProvider.AuthType.OAUTH2) {
Preconditions.checkArgument(
StringUtils.isNotBlank(oauth2CredentialKey),
"oauth2 with forwardUser=true requires '%s' to be set",
GravitinoAuthProvider.OAUTH2_TOKEN_CREDENTIAL_KEY);
}
```
##########
trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java:
##########
@@ -90,6 +142,45 @@ public ConnectorMetadata getMetadata(
ConnectorMetadata internalMetadata =
internalConnector.getMetadata(session,
gravitinoTransactionHandle.getInternalHandle());
Preconditions.checkArgument(internalMetadata != null, "Internal metadata
must not be null");
+
+ if (forwardUser) {
+ String credKey =
+ authType == GravitinoAuthProvider.AuthType.OAUTH2
+ ? "oauth2:" +
session.getIdentity().getExtraCredentials().get(oauth2CredentialKey)
+ : "simple:" + session.getUser();
+ UserSession userSession;
+ try {
+ userSession =
+ perUserSessionCache.get(
+ credKey,
+ () -> {
+ GravitinoAdminClient userClient =
+ GravitinoAuthProvider.buildForSession(
+ catalogConnectorContext.getConfig(), session);
+ GravitinoMetalake userMetalake =
+
userClient.loadMetalake(catalogConnectorContext.getMetalake().name());
+ CatalogConnectorMetadata userMetadata =
+ new CatalogConnectorMetadata(userMetalake,
catalogIdentifier);
+ return new UserSession(userClient, userMetadata);
+ });
+ } catch (ExecutionException e) {
+ Throwable cause = e.getCause();
+ LOG.warn(
+ "Failed to create per-user Gravitino client for user '{}': {}",
+ session.getUser(),
+ cause.getMessage());
+ throw new TrinoException(
+ PERMISSION_DENIED,
+ "Failed to authenticate user '"
+ + session.getUser()
+ + "' with Gravitino: "
+ + cause.getMessage(),
+ cause);
+ }
+ return createGravitinoMetadata(
+ userSession.metadata, catalogConnectorContext.getMetadataAdapter(),
internalMetadata);
+ }
Review Comment:
The new per-user session forwarding path in `getMetadata()` (per-user cache
keying + building per-session `GravitinoAdminClient` and `GravitinoMetalake`)
is a significant behavioral change but isn’t covered by unit tests in the
connector module. Adding tests that exercise `forwardUser=true` (simple and
oauth2), cache reuse, and failure cases (e.g., missing token / unsupported
authType) would help prevent regressions.
##########
trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/CatalogConnectorManager.java:
##########
@@ -414,7 +416,8 @@ public CatalogConnectorContext
createCatalogConnectorContext(
catalogConnectorFactory.createCatalogConnectorContextBuilder(catalog);
builder
.withMetalake(metalakes.computeIfAbsent(catalog.getMetalake(),
this::retrieveMetalake))
- .withContext(context);
+ .withContext(context)
+ .withConfig(this.config);
Review Comment:
`createCatalogConnectorContext(...)` receives a per-catalog `config`
parameter, but the builder is currently populated with
`withConfig(this.config)` (the manager’s global config captured from the first
initialization). This means per-catalog settings (including
`gravitino.client.*` auth/session forwarding options) may be ignored for
dynamically loaded catalogs, and multiple catalogs with different connector
properties would behave incorrectly. Consider passing the method parameter
`config` into the context instead (or rename/remove the parameter if
per-catalog config is intentionally unsupported).
```suggestion
.withConfig(config);
```
##########
trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java:
##########
@@ -90,6 +142,45 @@ public ConnectorMetadata getMetadata(
ConnectorMetadata internalMetadata =
internalConnector.getMetadata(session,
gravitinoTransactionHandle.getInternalHandle());
Preconditions.checkArgument(internalMetadata != null, "Internal metadata
must not be null");
+
+ if (forwardUser) {
+ String credKey =
+ authType == GravitinoAuthProvider.AuthType.OAUTH2
+ ? "oauth2:" +
session.getIdentity().getExtraCredentials().get(oauth2CredentialKey)
+ : "simple:" + session.getUser();
Review Comment:
The cache key for OAuth2 session forwarding is built from the raw Bearer
token value. This keeps access tokens in cleartext in heap keys (and
potentially in diagnostics like heap dumps) and can also amplify memory usage
for long tokens. Consider using a stable, non-sensitive key (e.g., a hash of
the token, or a token ID if available) and avoid retaining the raw token beyond
what’s needed to build the per-user client.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
